update Requires.jl code to package extensions (#175)

* update Requires.jl code to package extensions

* drop Requires completely and make ForwardDiff a dependency

* fix older Julia versions

* ignore ForwardDiff staledependency in Aqua tests

* resolve method ambiguities on julia nightly

* fix more ambiguities

Project.toml

@@ -1,15 +1,15 @@
 name = "Octavian"
 uuid = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 authors = ["Chris Elrod", "Dilum Aluthge", "Mason Protter", "contributors"]
-version = "0.3.21"
+version = "0.3.22"
 
 [deps]
 CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 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"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
@@ -18,18 +18,21 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
 CPUSummary = "0.1.26, 0.2.1"
+ForwardDiff = "0.10"
 IfElse = "0.1"
 LoopVectorization = "0.12.86"
 ManualMemory = "0.1.1"
 PolyesterWeave = "0.1.1, 0.2"
-Requires = "1"
 SnoopPrecompile = "1"
 Static = "0.8.4"
 StaticArrayInterface = "1"
 ThreadingUtilities = "0.5"
 VectorizationBase = "0.21.15"
 julia = "1.6"
 
+[extensions]
+ForwardDiffExt = "ForwardDiff"
+
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -43,3 +46,6 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [targets]
 test = ["Aqua", "BenchmarkTools", "ForwardDiff", "InteractiveUtils", "LinearAlgebra", "LoopVectorization", "Random", "VectorizationBase", "Test"]
+
+[weakdeps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

ext/ForwardDiffExt.jl

@@ -0,0 +1,235 @@
+module ForwardDiffExt
+
+using ForwardDiff: ForwardDiff
+
+using Octavian: ArrayInterface,
+                 @turbo, @tturbo,
+                 One, Zero,
+                 indices, static
+import Octavian: real_rep, _matmul!, _matmul_serial!
+
+real_rep(a::AbstractArray{DualT}) where {TAG,T,DualT<:ForwardDiff.Dual{TAG,T}} =
+  reinterpret(reshape, T, a)
+_view1(B::AbstractMatrix) = @view(B[1, :])
+_view1(B::AbstractArray{<:Any,3}) = @view(B[1, :, :])
+
+for AbstractVectorOrMatrix in (:AbstractVector, :AbstractMatrix)
+  # multiplication of dual vector/matrix by standard matrix from the left
+  @eval function _matmul!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    A::AbstractMatrix,
+    _B::$(AbstractVectorOrMatrix){DualT},
+    α,
+    β = 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
+  @eval @inline function _matmul!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    _A::AbstractMatrix{DualT},
+    B::$(AbstractVectorOrMatrix),
+    α = One(),
+    β = Zero(),
+    nthread::Nothing = nothing,
+    MKN = nothing
+  ) where {TAG,T,DualT<:ForwardDiff.Dual{TAG,T}}
+    if Bool(ArrayInterface.is_dense(_C)) &&
+      Bool(ArrayInterface.is_column_major(_C)) &&
+      Bool(ArrayInterface.is_dense(_A)) &&
+      Bool(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, nothing)
+    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
+
+  @eval @inline function _matmul!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    _A::AbstractMatrix{DualT},
+    _B::$(AbstractVectorOrMatrix){DualT},
+    α = One(),
+    β = Zero(),
+    nthread::Nothing = nothing,
+    MKN = nothing,
+    contig = nothing
+  ) where {TAG,T,P,DualT<:ForwardDiff.Dual{TAG,T,P}}
+    A = real_rep(_A)
+    C = real_rep(_C)
+    B = real_rep(_B)
+    if Bool(ArrayInterface.is_dense(_C)) &&
+      Bool(ArrayInterface.is_column_major(_C)) &&
+      Bool(ArrayInterface.is_dense(_A)) &&
+      Bool(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
+      Cₚₘₙ = zero(eltype(C))
+      for k ∈ indices((A, B), (3, 2))
+        Cₚₘₙ += A[1, m, k] * B[p+1, k, n]
+      end
+      C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ
+    end
+    _C
+  end
+
+  # multiplication of dual vector/matrix by standard matrix from the left
+  @eval function _matmul_serial!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    A::AbstractMatrix,
+    _B::$(AbstractVectorOrMatrix){DualT},
+    α,
+    β,
+    MKN
+  ) where {DualT<:ForwardDiff.Dual}
+    B = real_rep(_B)
+    C = real_rep(_C)
+
+    @turbo 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
+  @eval @inline function _matmul_serial!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    _A::AbstractMatrix{DualT},
+    B::$(AbstractVectorOrMatrix),
+    α,
+    β,
+    MKN
+  ) where {TAG,T,DualT<:ForwardDiff.Dual{TAG,T}}
+    if Bool(ArrayInterface.is_dense(_C)) &&
+      Bool(ArrayInterface.is_column_major(_C)) &&
+      Bool(ArrayInterface.is_dense(_A)) &&
+      Bool(ArrayInterface.is_column_major(_A))
+      # we can avoid the reshape and call the standard method
+      A = reinterpret(T, _A)
+      C = reinterpret(T, _C)
+      _matmul_serial!(C, A, B, α, β, nothing)
+    else
+      # we cannot use the standard method directly
+      A = real_rep(_A)
+      C = real_rep(_C)
+
+      @turbo 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
+
+  @eval @inline function _matmul_serial!(
+    _C::$(AbstractVectorOrMatrix){DualT},
+    _A::AbstractMatrix{DualT},
+    _B::$(AbstractVectorOrMatrix){DualT},
+    α,
+    β,
+    MKN
+  ) where {TAG,T,P,DualT<:ForwardDiff.Dual{TAG,T,P}}
+    A = real_rep(_A)
+    C = real_rep(_C)
+    B = real_rep(_B)
+    if Bool(ArrayInterface.is_dense(_C)) &&
+      Bool(ArrayInterface.is_column_major(_C)) &&
+      Bool(ArrayInterface.is_dense(_A)) &&
+      Bool(ArrayInterface.is_column_major(_A))
+      # we can avoid the reshape and call the standard method
+      Ar = reinterpret(T, _A)
+      Cr = reinterpret(T, _C)
+      _matmul_serial!(Cr, Ar, _view1(B), α, β, nothing)
+    else
+      # we cannot use the standard method directly
+      @turbo 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)
+    @turbo for n ∈ indices((B, C), 3), m ∈ indices((A, C), 2), p ∈ 1:Pstatic
+      Cₚₘₙ = zero(eltype(C))
+      for k ∈ indices((A, B), (3, 2))
+        Cₚₘₙ += A[1, m, k] * B[p+1, k, n]
+      end
+      C[p+1, m, n] = C[p+1, m, n] + α * Cₚₘₙ
+    end
+    _C
+  end
+end # for
+
+end # module

src/Octavian.jl

@@ -1,7 +1,5 @@
 module Octavian
 
-using Requires: @require
-
 using VectorizationBase, StaticArrayInterface, LoopVectorization
 
 using VectorizationBase:
@@ -67,6 +65,14 @@ include("complex_matmul.jl")
 
 include("init.jl") # `Octavian.__init__()` is defined in this file
 
+# TODO: This loads ForwardDiff unconditionally on Julia v1.6 - v1.8.
+#       It could be reconsidered when these older versions are not supported
+#       anymore. In this case, ForwardDiff should be removed from the
+#       dependencies and treated as weak dependency.
+if !isdefined(Base, :get_extension)
+  include("../ext/ForwardDiffExt.jl")
+end
+
 @static if VERSION >= v"1.8.0-beta1"
   @precompile_setup begin
     # Putting some things in `setup` can reduce the size of the