Add support for both arrays containing dual numbers. (#118)

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "Octavian"
 uuid = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 authors = ["Mason Protter", "Chris Elrod", "Dilum Aluthge", "contributors"]
-version = "0.3.5"
+version = "0.3.6"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/forward_diff.jl

@@ -26,7 +26,7 @@ end
         # we can avoid the reshape and call the standard method
         A = reinterpret(T, _A)
         C = reinterpret(T, _C)
-        _matmul!(C, A, B, α, β, nthread, MKN)
+        _matmul!(C, A, B, α, β, nthread, nothing)
     else
         # we cannot use the standard method directly
         A = real_rep(_A)
@@ -43,3 +43,33 @@ end
 
     _C
 end
+
+_view1(B::AbstractMatrix) = @view(B[1,:])
+_view1(B::AbstractArray{<:Any,3}) = @view(B[1,:,:])
+@inline function _matmul!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, _B::AbstractVecOrMat{DualT},
+                          α=One(), β=Zero(), nthread::Nothing=nothing, MKN=nothing) where {TAG, T, P, DualT <: ForwardDiff.Dual{TAG, T, P}}
+  A = real_rep(_A)
+  C = real_rep(_C)
+  B = real_rep(_B)
+  if all((ArrayInterface.is_dense(_C), ArrayInterface.is_column_major(_C),
+          ArrayInterface.is_dense(_A), ArrayInterface.is_column_major(_A)))
+    # we can avoid the reshape and call the standard method
+    Ar = reinterpret(T, _A)
+    Cr = reinterpret(T, _C)
+    _matmul!(Cr, Ar, _view1(B), α, β, nthread, nothing)
+  else
+    # we cannot use the standard method directly
+    @tturbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), 2), l in indices((C, A), 1)
+      Cₗₘₙ = zero(eltype(C))
+      for k ∈ indices((A, B), (3, 2))
+        Cₗₘₙ += A[l, m, k] * B[1, k, n]
+      end
+      C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
+    end
+  end
+  Pstatic = static(P)
+  @tturbo for n ∈ indices((B,C),3), m ∈ indices((A,C),2), p ∈ 1:Pstatic, k ∈ indices((A,B),(3,2))
+    C[p+1,m,n] += A[1,m,k] * B[p+1,k,n]
+  end
+  _C
+end

test/forward_diff.jl

@@ -1,8 +1,9 @@
 
+randdual(x, ::Val{N}=Val(3)) where {N} = ForwardDiff.Dual(x, ntuple(_ -> randn(), Val(N))...)
 @time @testset "ForwardDiff.jl" begin
-    m = 5
-    n = 6
-    k = 7
+    m = 53
+    n = 63
+    k = 73
 
     A1 = rand(Float64, m, k)
     B1 = rand(Float64, k, n)
@@ -60,4 +61,11 @@
         Octavian.matmul!(C2dual, A2dual', B2)
         @test C1dual ≈ C2dual
     end
+
+  @testset "two dual arrays" begin
+    A1d = randdual.(A1)
+    B1d = randdual.(B1)
+    @test reinterpret(Float64, Octavian.matmul(A1d, B1d)) ≈ reinterpret(Float64, A1d * B1d)
+    @test reinterpret(Float64, Octavian.matmul(@view(A1d[begin:end-1,:]), B1d)) ≈ reinterpret(Float64, @view(A1d[begin:end-1,:]) * B1d)
+  end
 end