Specialize _matmul! on ForwardDiff.Duals (#117)

* _matmul! with matrix multipliying array of ForwardDiff.Dual numbers

* _matmul! with matrix of ForwardDiff.Dual numbers multiplying array

* fall back to standard method if possible

* add compat bound for Requires

* increase test coverage

* set version to v0.3.5

Author: ranocha

Project.toml

@@ -1,14 +1,15 @@
 name = "Octavian"
 uuid = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 authors = ["Mason Protter", "Chris Elrod", "Dilum Aluthge", "contributors"]
-version = "0.3.4"
+version = "0.3.5"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 ManualMemory = "d125e4d3-2237-4719-b19c-fa641b8a4667"
 PolyesterWeave = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 ThreadingUtilities = "8290d209-cae3-49c0-8002-c8c24d57dab5"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
@@ -19,6 +20,7 @@ IfElse = "0.1"
 LoopVectorization = "0.12.86"
 ManualMemory = "0.1.1"
 PolyesterWeave = "0.1.1"
+Requires = "1"
 Static = "0.2, 0.3"
 ThreadingUtilities = "0.4.6"
 VectorizationBase = "0.21.15"
@@ -27,6 +29,7 @@ julia = "1.6"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
@@ -35,4 +38,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "InteractiveUtils", "LinearAlgebra", "LoopVectorization", "Random", "VectorizationBase", "Test"]
+test = ["Aqua", "BenchmarkTools", "ForwardDiff", "InteractiveUtils", "LinearAlgebra", "LoopVectorization", "Random", "VectorizationBase", "Test"]

src/Octavian.jl

@@ -1,10 +1,12 @@
 module Octavian
 
+using Requires: @require
+
 using VectorizationBase, ArrayInterface, LoopVectorization
 
 using VectorizationBase: align, AbstractStridedPointer, zstridedpointer, vsub_nsw, assume,
     static_sizeof, StridedPointer, gesp, pause, pick_vector_width, has_feature,
-    cache_size, num_cores, num_cores, cache_inclusive, cache_linesize
+    cache_size, num_cores, cache_inclusive, cache_linesize
 using LoopVectorization: preserve_buffer, CloseOpen, UpperBoundedInteger
 using ArrayInterface: size, strides, offsets, indices, axes, StrideIndex
 using IfElse: ifelse

src/forward_diff.jl

@@ -0,0 +1,45 @@
+
+real_rep(a::AbstractArray{DualT}) where {TAG, T, DualT<:ForwardDiff.Dual{TAG, T}} = reinterpret(reshape, T, a)
+
+# multiplication of dual vector/matrix by standard matrix from the left
+@inline function _matmul!(_C::AbstractVecOrMat{DualT}, A::AbstractMatrix, _B::AbstractVecOrMat{DualT},
+                          α=One(), β=Zero(), nthread::Nothing=nothing, MKN=nothing, contig_axis=nothing) where {DualT <: ForwardDiff.Dual}
+    B = real_rep(_B)
+    C = real_rep(_C)
+
+    @tturbo for n ∈ indices((C, B), 3), m ∈ indices((C, A), (2, 1)), l in indices((C, B), 1)
+        Cₗₘₙ = zero(eltype(C))
+        for k ∈ indices((A, B), 2)
+            Cₗₘₙ += A[m, k] * B[l, k, n]
+        end
+        C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
+    end
+
+    _C
+end
+
+# multiplication of dual matrix by standard vector/matrix from the right
+@inline function _matmul!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, B::AbstractVecOrMat,
+                          α=One(), β=Zero(), nthread::Nothing=nothing, MKN=nothing) where {TAG, T, DualT <: ForwardDiff.Dual{TAG, T}}
+    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
+        A = reinterpret(T, _A)
+        C = reinterpret(T, _C)
+        _matmul!(C, A, B, α, β, nthread, MKN)
+    else
+        # we cannot use the standard method directly
+        A = real_rep(_A)
+        C = real_rep(_C)
+
+        @tturbo for n ∈ indices((C, B), (3, 2)), m ∈ indices((C, A), 2), l in indices((C, A), 1)
+            Cₗₘₙ = zero(eltype(C))
+            for k ∈ indices((A, B), (3, 1))
+                Cₗₘₙ += A[l, m, k] * B[k, n]
+            end
+            C[l, m, n] = α * Cₗₘₙ + β * C[l, m, n]
+        end
+    end
+
+    _C
+end

src/init.jl

@@ -1,4 +1,6 @@
 function __init__()
+  @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include("forward_diff.jl")
+
   init_acache()
   init_bcache()
   nt = init_num_tasks()

src/macrokernels.jl

@@ -109,9 +109,9 @@ end
     for k ∈ CloseOpen(K)
       Aₘₖ = A[m,k]
       Cₘₙ += Aₘₖ * B[k,n]
-      Ãₚ[m,k] = Aₘₖ 
+      Ãₚ[m,k] = Aₘₖ
     end
-    C[m,n] = α * Cₘₙ + β * C[m,n] 
+    C[m,n] = α * Cₘₙ + β * C[m,n]
   end
 end
 @inline function alloc_a_pack(K, ::Val{T}) where {T}

test/forward_diff.jl

@@ -0,0 +1,63 @@
+
+@time @testset "ForwardDiff.jl" begin
+    m = 5
+    n = 6
+    k = 7
+
+    A1 = rand(Float64, m, k)
+    B1 = rand(Float64, k, n)
+    C1 = rand(Float64, m, n)
+
+    A2 = deepcopy(A1)
+    B2 = deepcopy(B1)
+    C2 = deepcopy(C1)
+
+    α = Float64(2.0)
+    β = Float64(2.0)
+
+    Octavian.matmul!(C1, A1, B1, α, β)
+    LinearAlgebra.mul!(C2, A2, B2, α, β)
+    @test C1 ≈ C2
+
+    @testset "real array from the left" begin
+        config = ForwardDiff.JacobianConfig(nothing, C1, B1)
+        I = LinearAlgebra.I(size(B1, 2))
+
+        J1 = ForwardDiff.jacobian((C, B) -> Octavian.matmul!(C, A1, B), C1, B1, config)
+        @test J1 ≈ kron(I, A1)
+
+        J2 = ForwardDiff.jacobian((C, B) -> LinearAlgebra.mul!(C, A2, B), C2, B2, config)
+        @test J1 ≈ kron(I, A2)
+        @test J1 ≈ J2
+    end
+
+    @testset "real array from the right" begin
+        # dense and column-major arrays
+        config = ForwardDiff.JacobianConfig(nothing, C1, A1)
+
+        J1 = ForwardDiff.jacobian((C, A) -> Octavian.matmul!(C, A, B1), C1, A1, config)
+        J2 = ForwardDiff.jacobian((C, A) -> LinearAlgebra.mul!(C, A, B2), C2, A2, config)
+        @test J1 ≈ J2
+
+        # transposed arrays
+        A1new = Matrix(A1')'
+        A2new = Matrix(A2')'
+        config = ForwardDiff.JacobianConfig(nothing, C1, A1new)
+
+        J1 = ForwardDiff.jacobian((C, A) -> Octavian.matmul!(C, A, B1), C1, A1new, config)
+        J2 = ForwardDiff.jacobian((C, A) -> LinearAlgebra.mul!(C, A, B2), C2, A2new, config)
+        @test J1 ≈ J2
+
+        # direct version using dual numbers
+        A1dual = zeros(eltype(config), reverse(size(A1))...)
+        A1dual .= A1'
+        C1dual = zeros(eltype(config), size(C1)...)
+
+        A2dual = deepcopy(A1dual)
+        C2dual = deepcopy(C1dual)
+
+        Octavian.matmul!(C1dual, A1dual', B1)
+        Octavian.matmul!(C2dual, A2dual', B2)
+        @test C1dual ≈ C2dual
+    end
+end

test/runtests.jl

@@ -2,6 +2,7 @@ import Octavian
 
 import Aqua
 import BenchmarkTools
+import ForwardDiff
 import InteractiveUtils
 import LinearAlgebra
 import LoopVectorization
@@ -24,6 +25,7 @@ include("integer_division.jl")
 include("macrokernels.jl")
 include("matmul_coverage.jl")
 include("utils.jl")
+include("forward_diff.jl")
 
 if !coverage
     include("matmul_main.jl")