Despecialize blas_mul! for improved compilation time

Project.toml

@@ -5,6 +5,7 @@ authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd19
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
@@ -29,6 +30,7 @@ GPUArrays = "11.4.1"
 JLArrays = "0.3.1"
 LinearAlgebra = "1.6"
 Metal = "1.9"
+PrecompileTools = "1"
 Random = "1.6"
 StridedViews = "0.5.1"
 Test = "1.6"

src/Strided.jl

@@ -57,4 +57,6 @@ include("broadcast.jl")
 include("macros.jl")
 include("convert.jl")
 
+include("precompile.jl")
+
 end

src/linalg.jl

@@ -111,50 +111,47 @@ function blas_mul!(
         C::StridedView{T, 2}, A::StridedView{T, 2}, B::StridedView{T, 2},
         α::Number, β::Number
     ) where {T <: LinearAlgebra.BlasFloat}
+    m, n = size(C)
+    m == size(A, 1) && n == size(B, 2) || throw(DimensionMismatch())
     nthreads = use_threaded_mul() ? get_num_threads() : 1
-    return _threaded_blas_mul!(C, A, B, α, β, nthreads)
+    A2, CA = getblasmatrix(A)
+    B2, CB = getblasmatrix(B)
+    return _threaded_blas_mul!(C, A2, CA, B2, CB, convert(T, α), convert(T, β), nthreads)
 end
 
+_split_row(A, m) = (A[1:m, :], A[(m + 1):end, :])
+_split_col(A, m) = (A[:, 1:m], A[:, (m + 1):end])
+
 function _threaded_blas_mul!(
-        C::StridedView{T, 2}, A::StridedView{T, 2}, B::StridedView{T, 2},
+        C::StridedView{T, 2},
+        A::StridedView{T, 2}, CA,
+        B::StridedView{T, 2}, CB,
         α::Number, β::Number,
         nthreads
     ) where {T <: LinearAlgebra.BlasFloat}
     m, n = size(C)
-    m == size(A, 1) && n == size(B, 2) || throw(DimensionMismatch())
-    return if nthreads == 1 || m * n < 1024
-        A2, CA = getblasmatrix(A)
-        B2, CB = getblasmatrix(B)
-        LinearAlgebra.BLAS.gemm!(CA, CB, convert(T, α), A2, B2, convert(T, β), C)
+    if nthreads == 1 || m * n < 1024
+        LinearAlgebra.BLAS.gemm!(CA, CB, α, A, B, β, C)
+        return C
+    end
+
+    if m > n
+        m2 = round(Int, m / 16) * 8
+        Chead, Ctail = _split_row(C, m2)
+        Ahead, Atail = CA == 'N' ? _split_row(A, m2) : _split_col(A, m2)
+        Bhead = Btail = B
     else
-        if m > n
-            m2 = round(Int, m / 16) * 8
-            nthreads2 = nthreads >> 1
-            t = Threads.@spawn _threaded_blas_mul!(
-                C[1:($m2), :], A[1:($m2), :], B, α, β,
-                $nthreads2
-            )
-            _threaded_blas_mul!(
-                C[(m2 + 1):m, :], A[(m2 + 1):m, :], B, α, β,
-                nthreads - nthreads2
-            )
-            wait(t)
-            return C
-        else
-            n2 = round(Int, n / 16) * 8
-            nthreads2 = nthreads >> 1
-            t = Threads.@spawn _threaded_blas_mul!(
-                C[:, 1:($n2)], A, B[:, 1:($n2)], α, β,
-                $nthreads2
-            )
-            _threaded_blas_mul!(
-                C[:, (n2 + 1):n], A, B[:, (n2 + 1):n], α, β,
-                nthreads - nthreads2
-            )
-            wait(t)
-            return C
-        end
+        n2 = round(Int, n / 16) * 8
+        Chead, Ctail = _split_col(C, n2)
+        Ahead = Atail = A
+        Bhead, Btail = CB == 'N' ? _split_col(B, n2) : _split_row(B, n2)
     end
+    nthreads1 = nthreads >> 1
+    nthreads2 = nthreads - nthreads1
+    t = Threads.@spawn _threaded_blas_mul!(Chead, Ahead, CA, Bhead, CB, α, β, nthreads1)
+    _threaded_blas_mul!(Ctail, Atail, CA, Btail, CB, α, β, nthreads2)
+    wait(t)
+    return C
 end
 
 # This implementation is faster than LinearAlgebra.generic_matmatmul

src/mapreduce.jl

@@ -253,17 +253,6 @@ function _mapreduce_threaded!(
     return nothing
 end
 
-@generated function _mapreduce_kernel!(
-        (f), (op),
-        (initop), dims::NTuple{N, Int},
-        blocks::NTuple{N, Int},
-        arrays::NTuple{M, StridedView},
-        strides::NTuple{M, NTuple{N, Int}},
-        offsets::NTuple{M, Int}
-    ) where {N, M}
-    return _mapreduce_kernel_expr(f, op, initop, N, M)
-end
-
 function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
 
     # many variables
@@ -503,6 +492,17 @@ function _mapreduce_kernel_expr(f, op, initop, N::Int, M::Int)
     return ex
 end
 
+@generated function _mapreduce_kernel!(
+        (f), (op),
+        (initop), dims::NTuple{N, Int},
+        blocks::NTuple{N, Int},
+        arrays::NTuple{M, StridedView},
+        strides::NTuple{M, NTuple{N, Int}},
+        offsets::NTuple{M, Int}
+    ) where {N, M}
+    return _mapreduce_kernel_expr(f, op, initop, N, M)
+end
+
 function indexorder(strides::NTuple{N, Int}) where {N}
     # returns order such that strides[i] is the order[i]th smallest element of strides, not
     # counting zero strides zero strides have order 1

src/precompile.jl

@@ -0,0 +1,23 @@
+using PrecompileTools: @setup_workload, @compile_workload
+
+@setup_workload begin
+    Ts = [Float32, Float64, ComplexF32, ComplexF64]
+    blasops = [identity, conj, transpose, adjoint]
+    @compile_workload begin
+        for T in Ts
+            A = StridedView(zeros(T, 2, 2))
+            B = StridedView(zeros(T, 2, 2))
+            C = StridedView(zeros(T, 2, 2))
+            # op variety on the operands, both scalar-argument forms
+            for opA in blasops, opB in blasops
+                mul!(C, opA(A), opB(B))
+                mul!(C, opA(A), opB(B), one(T), zero(T))
+            end
+            # op variety on the destination (mul! normalization branches)
+            for opC in (transpose, adjoint)
+                mul!(opC(C), A, B)
+                mul!(opC(C), A, B, one(T), zero(T))
+            end
+        end
+    end
+end