diff --git a/Project.toml b/Project.toml index da9069b..b76a0b4 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ authors = ["Lukas Devos ", "Maarten Van Damme 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 diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 5de8e6c..7a48290 100644 --- a/src/mapreduce.jl +++ b/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 diff --git a/src/precompile.jl b/src/precompile.jl new file mode 100644 index 0000000..37d56d3 --- /dev/null +++ b/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