don't multithread x86 (#146)

* don't multithread x86

* matmul for x86

* prefer chained && over all

* forwardiff updates

* don't test forward_diff on 32 bit Windows

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.13"
+version = "0.3.14"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -16,7 +16,7 @@ ThreadingUtilities = "8290d209-cae3-49c0-8002-c8c24d57dab5"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
-ArrayInterface = "3.1.14, 5.0.1"
+ArrayInterface = "3.1.14, 5.0.1, 6"
 CPUSummary = "0.1.1 - 0.1.8, 0.1.14"
 IfElse = "0.1"
 LoopVectorization = "0.12.86"

src/forward_diff.jl

@@ -1,58 +1,59 @@
 
 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,:,:])
+
 
 # multiplication of dual vector/matrix by standard matrix from the left
 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)
+                  α, β=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]
+  @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
+  _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, 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
+  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
+  _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}}
+                          α=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 all((ArrayInterface.is_dense(_C), ArrayInterface.is_column_major(_C),
-          ArrayInterface.is_dense(_A), ArrayInterface.is_column_major(_A)))
+  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)
@@ -77,3 +78,80 @@ _view1(B::AbstractArray{<:Any,3}) = @view(B[1,:,:])
   end
   _C
 end
+
+
+# multiplication of dual vector/matrix by standard matrix from the left
+function _matmul_serial!(_C::AbstractVecOrMat{DualT}, A::AbstractMatrix, _B::AbstractVecOrMat{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
+@inline function _matmul_serial!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, B::AbstractVecOrMat,
+                                 α, β, 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
+
+@inline function _matmul_serial!(_C::AbstractVecOrMat{DualT}, _A::AbstractMatrix{DualT}, _B::AbstractVecOrMat{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
+

src/matmul.jl

@@ -115,6 +115,16 @@ end
   matmul_serial!(C, A, B, One(), Zero(), (M,K,N), ArrayInterface.contiguous_axis(C))
   return C
 end
+@inline function matmul_serial(A::AbstractMatrix, B::AbstractVecOrMat, α)
+  C, (M,K,N) = alloc_matmul_product(A, B)
+  matmul_serial!(C, A, B, α, Zero(), (M,K,N), ArrayInterface.contiguous_axis(C))
+  return C
+end
+@inline function matmul_serial(A::AbstractMatrix, B::AbstractVecOrMat, α, β)
+  C, (M,K,N) = alloc_matmul_product(A, B)
+  matmul_serial!(C, A, B, α, β, (M,K,N), ArrayInterface.contiguous_axis(C))
+  return C
+end
 
 
 # These methods must be compile time constant
@@ -165,7 +175,7 @@ Otherwise, based on the array's size, whether they are transposed, and whether t
 """
 @inline function _matmul_serial!(
   C::AbstractMatrix{T}, A::AbstractMatrix, B::AbstractMatrix, α, β, MKN
-) where {T}
+) where {T <: Base.HWReal}
   ((β ≢ Zero()) && iszero(β)) && return _matmul_serial!(C, A, B, α, Zero(), MKN)
   (β isa Bool) && return _matmul_serial!(C, A, B, α, One(), MKN)
   M, K, N = MKN === nothing ? matmul_sizes(C, A, B) : MKN
@@ -216,7 +226,7 @@ function matmul_st_pack_dispatcher!(pC::AbstractStridedPointer{T}, pA, pB, α, 
     nothing
 end
 
-
+if sizeof(Int) >= 8
 """
     matmul(A, B)
 
@@ -365,9 +375,7 @@ function __matmul!(
   else
     clamp(div_fast(M * N, StaticInt{256}() * W), 0, _nthread-1)
   end
-  # nkern = cld_fast(M * N,  MᵣW * Nᵣ)
   threads, torelease = PolyesterWeave.__request_threads(_nrequest % UInt32, PolyesterWeave.worker_pointer(), nothing)
-  # _threads, _torelease = PolyesterWeave.request_threads(Threads.threadid()%UInt32, _nrequest)
 
   nrequest = threads.i
   iszero(nrequest) && @goto SINGLETHREAD
@@ -401,9 +409,6 @@ end
 
 # If tasks is [0,1,2,3] (e.g., `CloseOpen(0,4)`), it will wait on `MULTASKS[i]` for `i = [1,2,3]`.
 function waitonmultasks(threads, nthread)
-  # for (_,tid) ∈ threads
-  #   wait(tid)
-  # end
   (tnum, tuu) = PolyesterWeave.initial_state(threads)
   for _ ∈ CloseOpen(One(), nthread)
     (tnum, tuu) = PolyesterWeave.iter(tnum, tuu)
@@ -524,7 +529,11 @@ function sync_mul!(
   end
   nothing
 end
-
+else
+  @inline matmul(args::Vararg{Any,K}) where {K} = matmul_serial(args...)
+  @inline matmul!(args::Vararg{Any,K}) where {K} = matmul_serial!(args...)
+end
+  
 function _matmul!(y::AbstractVector{T}, A::AbstractMatrix, x::AbstractVector, α, β, _, __) where {T}
   @tturbo for m ∈ indices((A,y),1)
     yₘ = zero(T)

test/_matmul.jl

@@ -78,7 +78,7 @@ function matmul_pack_ab!(C, A, B)
   zc, za, zb = Octavian.zstridedpointer.((C,A,B))
   nspawn = min(Threads.nthreads(), Octavian.num_cores())
   GC.@preserve C A B begin
-    if nspawn > 1
+    if nspawn > 1 && sizeof(Int) >= 8
       threads, torelease = Octavian.PolyesterWeave.__request_threads((nspawn-1)%UInt32, Octavian.PolyesterWeave.worker_pointer(), nothing)
       @assert threads.i < Threads.nthreads()
       Octavian.matmul_pack_A_and_B!(

test/forward_diff.jl

@@ -30,6 +30,11 @@ randdual(x, ::Val{N}=Val(3)) where {N} = ForwardDiff.Dual(x, ntuple(_ -> randn()
         J2 = ForwardDiff.jacobian((C, B) -> LinearAlgebra.mul!(C, A2, B), C2, B2, config)
         @test J1 ≈ kron(I, A2)
         @test J1 ≈ J2
+
+        J3 = ForwardDiff.jacobian((C, B) -> Octavian.matmul_serial!(C, A1, B), C1, B1, config)
+        @test J3 ≈ kron(I, A1)
+        @test J3 ≈ J2
+
     end
 
     @testset "real array from the right" begin
@@ -38,7 +43,8 @@ randdual(x, ::Val{N}=Val(3)) where {N} = ForwardDiff.Dual(x, ntuple(_ -> randn()
 
         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
+        J3 = ForwardDiff.jacobian((C, A) -> Octavian.matmul_serial!(C, A, B1), C1, A1, config)
+        @test J1 ≈ J2 ≈ J3
 
         # transposed arrays
         A1new = Matrix(A1')'
@@ -47,7 +53,8 @@ randdual(x, ::Val{N}=Val(3)) where {N} = ForwardDiff.Dual(x, ntuple(_ -> randn()
 
         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
+        J3 = ForwardDiff.jacobian((C, A) -> Octavian.matmul_serial!(C, A, B1), C1, A1new, config)
+        @test J1 ≈ J2 ≈ J3
 
         # direct version using dual numbers
         A1dual = zeros(eltype(config), reverse(size(A1))...)
@@ -56,16 +63,18 @@ randdual(x, ::Val{N}=Val(3)) where {N} = ForwardDiff.Dual(x, ntuple(_ -> randn()
 
         A2dual = deepcopy(A1dual)
         C2dual = deepcopy(C1dual)
-
+        C3dual = similar(C1dual); C4dual = similar(C2dual)
         Octavian.matmul!(C1dual, A1dual', B1)
         Octavian.matmul!(C2dual, A2dual', B2)
-        @test C1dual ≈ C2dual
+        Octavian.matmul_serial!(C3dual, A1dual', B1)
+        Octavian.matmul_serial!(C4dual, A2dual', B2)
+        @test C1dual ≈ C2dual ≈ C3dual ≈ C4dual
     end
 
   @testset "two dual arrays" begin
     A1d = randdual.(A1)
     B1d = randdual.(B1)
-    @test reinterpret(Float64, Octavian.matmul(A1d, B1d, 1.3)) ≈ reinterpret(Float64, (A1d * B1d) .* 1.3)
-    @test reinterpret(Float64, Octavian.matmul(@view(A1d[begin:end-1,:]), B1d)) ≈ reinterpret(Float64, @view(A1d[begin:end-1,:]) * B1d)
+    @test reinterpret(Float64, Octavian.matmul(A1d, B1d, 1.3)) ≈ reinterpret(Float64, Octavian.matmul_serial(A1d, B1d, 1.3)) ≈ reinterpret(Float64, (A1d * B1d) .* 1.3)
+    @test reinterpret(Float64, Octavian.matmul(@view(A1d[begin:end-1,:]), B1d)) ≈ reinterpret(Float64, Octavian.matmul_serial(@view(A1d[begin:end-1,:]), B1d)) ≈ reinterpret(Float64, @view(A1d[begin:end-1,:]) * B1d)
   end
 end

test/runtests.jl

@@ -34,6 +34,8 @@ include("_matmul.jl")
 coverage || include("matmul_main.jl")
 include("matmul_coverage.jl")
 include("utils.jl")
-include("forward_diff.jl")
+if sizeof(Int) >= 8 || !Sys.iswindows()
+  include("forward_diff.jl")
+end
 
 include("aqua.jl") # run the Aqua.jl tests last