Implement LinearAlgebra's storage-specific mul! methods (#630)

Author: maleadt

Project.toml

@@ -38,7 +38,7 @@ AcceleratedKernels = "0.3.1, 0.4"
 Adapt = "4"
 CEnum = "0.4, 0.5"
 ExprTools = "0.1"
-GPUArrays = "11.2.1"
+GPUArrays = "11.5.14"
 GPUCompiler = "2"
 GPUToolbox = "0.1, 0.2, 0.3, 1, 3"
 KernelAbstractions = "0.9.39"

lib/mkl/interfaces.jl

@@ -2,40 +2,49 @@
 
 using LinearAlgebra: BlasComplex, BlasFloat, BlasReal, MulAddMul
 
-# legacy methods with final MulAddMul argument
-LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSR{T}, B::oneVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
-    LinearAlgebra.generic_matvecmul!(C, tA, A, B, _add.alpha, _add.beta)
-LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSC{T}, B::oneVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
-    LinearAlgebra.generic_matvecmul!(C, tA, A, B, _add.alpha, _add.beta)
-LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSR{T}, B::oneMatrix{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
-    LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
-LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSC{T}, B::oneMatrix{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
-    LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
-
-function LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSR{T}, B::oneVector{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
+function LinearAlgebra.mul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSR{T}, B::oneVector{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     return sparse_gemv!(tA, alpha, A, B, beta, C)
 end
 
-function LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSC{T}, B::oneVector{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
+function LinearAlgebra.mul!(C::oneVector{T}, tA::AbstractChar, A::oneSparseMatrixCSC{T}, B::oneVector{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
     # sparse_gemv! already maps op(A) onto the transposed CSR handle, so tA is passed through
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     return sparse_gemv!(tA, alpha, A, B, beta, C)
 end
 
-function LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSR{T}, B::oneMatrix{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
+function LinearAlgebra.mul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSR{T}, B::oneMatrix{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     tB = tB in ('S', 's', 'H', 'h') ? 'N' : tB
     return sparse_gemm!(tA, tB, alpha, A, B, beta, C)
 end
 
-function LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSC{T}, B::oneMatrix{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
+function LinearAlgebra.mul!(C::oneMatrix{T}, tA, tB, A::oneSparseMatrixCSC{T}, B::oneMatrix{T}, alpha::Number, beta::Number) where {T <: BlasFloat}
     # sparse_gemm! already maps op(A) onto the transposed CSR handle, so tA is passed through
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     tB = tB in ('S', 's', 'H', 'h') ? 'N' : tB
     return sparse_gemm!(tA, tB, alpha, A, B, beta, C)
 end
 
+# Julia < 1.13 dispatches on the non-public `generic_matvecmul!` and `generic_matmatmul!`,
+# which JuliaLang/LinearAlgebra.jl#1671 superseded by the `mul!` methods above. Forward from
+# the old names, both the alpha/beta variants (1.12) and the ones taking a final MulAddMul
+# (1.10 and 1.11).
+@static if VERSION < v"1.13.0-rc4"
+    for SparseMatrixType in (:oneSparseMatrixCSR, :oneSparseMatrixCSC)
+        @eval begin
+            LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::$SparseMatrixType{T}, B::oneVector{T}, alpha::Number, beta::Number) where {T <: BlasFloat} =
+                LinearAlgebra.mul!(C, tA, A, B, alpha, beta)
+            LinearAlgebra.generic_matvecmul!(C::oneVector{T}, tA::AbstractChar, A::$SparseMatrixType{T}, B::oneVector{T}, _add::MulAddMul) where {T <: BlasFloat} =
+                LinearAlgebra.mul!(C, tA, A, B, _add.alpha, _add.beta)
+            LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::$SparseMatrixType{T}, B::oneMatrix{T}, alpha::Number, beta::Number) where {T <: BlasFloat} =
+                LinearAlgebra.mul!(C, tA, tB, A, B, alpha, beta)
+            LinearAlgebra.generic_matmatmul!(C::oneMatrix{T}, tA, tB, A::$SparseMatrixType{T}, B::oneMatrix{T}, _add::MulAddMul) where {T <: BlasFloat} =
+                LinearAlgebra.mul!(C, tA, tB, A, B, _add.alpha, _add.beta)
+        end
+    end
+end
+
 function LinearAlgebra.generic_trimatdiv!(C::oneVector{T}, uploc, isunitc, tfun::Function, A::oneSparseMatrixCSR{T}, B::oneVector{T}) where {T <: BlasFloat}
     return sparse_trsv!(uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), A, B, C)
 end

lib/mkl/linalg.jl

@@ -70,9 +70,7 @@ end
 #
 # BLAS 2
 
-LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::oneStridedMatrix, B::oneStridedVector, _add::MulAddMul) =
-    LinearAlgebra.generic_matvecmul!(Y, tA, A, B, _add.alpha, _add.beta)
-function LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::oneStridedMatrix, B::oneStridedVector, a::Number, b::Number)
+function LinearAlgebra.mul!(Y::oneVector, tA::AbstractChar, A::oneStridedMatrix, B::oneStridedVector, a::Number, b::Number)
     mA, nA = tA == 'N' ? size(A) : reverse(size(A))
 
     if nA != length(B)
@@ -96,8 +94,8 @@ function LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::one
             if tA in ('N', 'T', 'C')
                 return gemv!(tA, alpha, A, B, beta, Y)
             elseif tA in ('S', 's') && T <: Real
-                # complex symv! is not wrapped; fall through to generic_matmatmul!,
-                # which can use symm! instead
+                # complex symv! is not wrapped; fall through to the matrix-matrix
+                # `mul!`, which can use symm! instead
                 return symv!(tA == 'S' ? 'U' : 'L', alpha, A, B, beta, Y)
             elseif tA in ('H', 'h')
                 # hemv! only supports complex eltypes, but a real Hermitian matrix
@@ -107,7 +105,7 @@ function LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::one
             end
         end
     end
-    return LinearAlgebra.generic_matmatmul!(Y, tA, 'N', A, B, alpha, beta)
+    return LinearAlgebra.mul!(Y, tA, 'N', A, B, alpha, beta)
 end
 
 # triangular
@@ -123,11 +121,7 @@ LinearAlgebra.generic_trimatdiv!(C::oneStridedVector{T}, uploc, isunitc, tfun::F
 # BLAS 3
 #
 
-LinearAlgebra.generic_matmatmul!(
-    C::oneStridedVecOrMat, tA, tB, A::oneStridedVecOrMat,
-    B::oneStridedVecOrMat, _add::MulAddMul,
-) = LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
-function LinearAlgebra.generic_matmatmul!(
+function LinearAlgebra.mul!(
         C::oneStridedVecOrMat, tA, tB, A::oneStridedVecOrMat,
         B::oneStridedVecOrMat, alpha::Number, beta::Number,
     )
@@ -185,6 +179,21 @@ function LinearAlgebra.generic_matmatmul!(
     GPUArrays.generic_matmatmul!(C, wrap(A, tA), wrap(B, tB), alpha, beta)
 end
 
+# Julia < 1.13 dispatches on the non-public `generic_matvecmul!` and `generic_matmatmul!`,
+# which JuliaLang/LinearAlgebra.jl#1671 superseded by the `mul!` methods above. Forward from
+# the old names, both the alpha/beta variants (1.12) and the ones taking a final MulAddMul
+# (1.10 and 1.11).
+@static if VERSION < v"1.13.0-rc4"
+    LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::oneStridedMatrix, B::oneStridedVector, alpha::Number, beta::Number) =
+        LinearAlgebra.mul!(Y, tA, A, B, alpha, beta)
+    LinearAlgebra.generic_matvecmul!(Y::oneVector, tA::AbstractChar, A::oneStridedMatrix, B::oneStridedVector, _add::MulAddMul) =
+        LinearAlgebra.mul!(Y, tA, A, B, _add.alpha, _add.beta)
+    LinearAlgebra.generic_matmatmul!(C::oneStridedVecOrMat, tA, tB, A::oneStridedVecOrMat, B::oneStridedVecOrMat, alpha::Number, beta::Number) =
+        LinearAlgebra.mul!(C, tA, tB, A, B, alpha, beta)
+    LinearAlgebra.generic_matmatmul!(C::oneStridedVecOrMat, tA, tB, A::oneStridedVecOrMat, B::oneStridedVecOrMat, _add::MulAddMul) =
+        LinearAlgebra.mul!(C, tA, tB, A, B, _add.alpha, _add.beta)
+end
+
 # triangular
 LinearAlgebra.generic_trimatmul!(C::oneStridedMatrix{T}, uploc, isunitc, tfun::Function, A::oneStridedMatrix{T}, B::oneStridedMatrix{T}) where {T<:onemklFloat} =
     trmm!('L', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), A, C === B ? C : copyto!(C, B))