From ecfc49934244992cd318bc42c1457d9d0333d359 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 3 Sep 2026 11:14:05 +0200 Subject: [PATCH] Implement LinearAlgebra's storage-specific `mul!` methods Julia 1.13-rc4 bumped LinearAlgebra to a version that includes JuliaLang/LinearAlgebra.jl#1671, which renamed the storage-level `generic_matvecmul!`/`generic_matmatmul!` entry points into 6- and 7-argument `mul!` methods, with the old names kept only as a fallback behind the generic `mul!` method. LinearAlgebra's strided BlasFloat `mul!` methods are more specific than that fallback, so on 1.13 our `generic_matvecmul!` overload for strided GPU arrays was never reached and products ended up in CPU BLAS (which only works for GPU array types that happen to expose a host pointer). The `generic_matmatmul_wrapper!` hooks likewise forwarded to the legacy name, which on 1.13 bypasses the `mul!` methods that back-ends provide. Implement the generic fallback as the new `mul!` methods, route the wrapper hooks through them, and keep the legacy names only as forwarding shims for Julia before 1.13, so that they can be dropped together once 1.13 becomes the floor. Back-ends should overload the same `mul!` signatures for their own array types. --- Project.toml | 2 +- src/host/linalg.jl | 56 +++++++++++++++++++++++++++++++--------------- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/Project.toml b/Project.toml index 71e184f4..dc16f5f3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "GPUArrays" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "11.5.13" +version = "11.5.14" [workspace] projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"] diff --git a/src/host/linalg.jl b/src/host/linalg.jl index f5cbf7ac..2d19fb46 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -597,28 +597,39 @@ end end end -@static if !isdefined(LinearAlgebra, Symbol("@stable_muladdmul")) # @stable_muladdmul was added in 1.12 -function LinearAlgebra.generic_matvecmul!(C::AnyStridedGPUVector, tA::AbstractChar, A::AnyStridedGPUMatrix, B::AnyStridedGPUVector, _add::MulAddMul = MulAddMul()) +# Storage-level products below LinearAlgebra's 5-argument `mul!`, used for GPU arrays whose +# back-end does not provide a native implementation. Back-ends overload these same `mul!` +# signatures for their own array types to route products to vendor libraries. +function LinearAlgebra.mul!(C::AnyStridedGPUVector, tA::AbstractChar, A::AnyStridedGPUMatrix, B::AnyStridedGPUVector, a::Number, b::Number) Cm = reshape(C, length(C), 1) Bm = reshape(B, length(B), 1) - generic_matmatmul!(Cm, tA, 'N', A, Bm, _add) + LinearAlgebra.mul!(Cm, tA, 'N', A, Bm, a, b) return C end -function LinearAlgebra.generic_matmatmul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, _add::MulAddMul=MulAddMul()) - generic_matmatmul!(C, tA, tB, A, B, _add) +@static if isdefined(LinearAlgebra, Symbol("@stable_muladdmul")) # @stable_muladdmul was added in 1.12 +function LinearAlgebra.mul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, a::Number, b::Number) + LinearAlgebra.@stable_muladdmul generic_matmatmul!(C, tA, tB, A, B, MulAddMul(a, b)) end else -function LinearAlgebra.generic_matvecmul!(C::AnyStridedGPUVector, tA::AbstractChar, A::AnyStridedGPUMatrix, B::AnyStridedGPUVector, a::Number, b::Number) - Cm = reshape(C, length(C), 1) - Bm = reshape(B, length(B), 1) - LinearAlgebra.@stable_muladdmul generic_matmatmul!(Cm, tA, 'N', A, Bm, MulAddMul(a, b)) - return C +function LinearAlgebra.mul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, a::Number, b::Number) + generic_matmatmul!(C, tA, tB, A, B, MulAddMul(a, b)) end - -function LinearAlgebra.generic_matmatmul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, a::Number, b::Number) - LinearAlgebra.@stable_muladdmul generic_matmatmul!(C, tA, tB, A, B, MulAddMul(a, b)) 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!(C::AnyStridedGPUVector, tA::AbstractChar, A::AnyStridedGPUMatrix, B::AnyStridedGPUVector, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, A, B, a, b) + LinearAlgebra.generic_matvecmul!(C::AnyStridedGPUVector, tA::AbstractChar, A::AnyStridedGPUMatrix, B::AnyStridedGPUVector, _add::MulAddMul = MulAddMul()) = + LinearAlgebra.mul!(C, tA, A, B, _add.alpha, _add.beta) + LinearAlgebra.generic_matmatmul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, a::Number, b::Number) = + LinearAlgebra.mul!(C, tA, tB, A, B, a, b) + LinearAlgebra.generic_matmatmul!(C::AnyStridedGPUVecOrMat, tA, tB, A::AnyStridedGPUVecOrMat, B::AnyStridedGPUVecOrMat, _add::MulAddMul = MulAddMul()) = + LinearAlgebra.mul!(C, tA, tB, A, B, _add.alpha, _add.beta) end # triangular × triangular matmul: C = α·(A·B) + β·C, with both A and B triangular. @@ -660,23 +671,32 @@ end @static if VERSION ≥ v"1.12.0-rc" # we need to use the generic wrapper to avoid dispatch to the 2x2or3x3 method using LinearAlgebra: generic_matmatmul_wrapper!, BlasFlag + # The wrappers hand strided GPU arrays back to the storage-level product. Before Julia 1.13 + # that is still the legacy name, which back-ends released for those versions overload. + @static if VERSION < v"1.13.0-rc4" + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) = + LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + else + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) = + LinearAlgebra.mul!(C, tA, tB, A, B, alpha, beta) + end function LinearAlgebra.generic_matmatmul_wrapper!(C::AnyStridedGPUMatrix{T}, tA::AbstractChar, tB::AbstractChar, A::AnyStridedGPUVecOrMat{T}, B::AnyStridedGPUVecOrMat{T}, alpha::Number, beta::Number, val::LinearAlgebra.BlasFlag.SyrkHerkGemm) where {T} - LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) end # Symmetric/Hermitian inputs with BLAS eltypes would otherwise dispatch to BLAS.symm!/ # hemm!: GPU arrays are DenseArrays, so they match the StridedMatrix{<:BlasFloat} methods function LinearAlgebra.generic_matmatmul_wrapper!(C::AnyStridedGPUMatrix{T}, tA::AbstractChar, tB::AbstractChar, A::AnyStridedGPUVecOrMat{T}, B::AnyStridedGPUVecOrMat{T}, alpha::Number, beta::Number, val::LinearAlgebra.BlasFlag.SymmHemmGeneric) where {T} - LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) end # need to support mixed complex/real types too #function LinearAlgebra.generic_matmatmul_wrapper!(C::AbstractGPUMatrix{Complex{T}}, tA::AbstractChar, tB::AbstractChar, A::AbstractGPUVecOrMat{Complex{T}}, B::AbstractGPUVecOrMat{T}, alpha::Number, beta::Number, val::V) where {T<:BlasReal, V<:LinearAlgebra.BlasFlag.SyrkHerkGemm} - # LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + # storage_matmatmul!(C, tA, tB, A, B, alpha, beta) #end function LinearAlgebra.generic_matmatmul_wrapper!(C::AnyStridedGPUMatrix{Complex{T}}, tA::AbstractChar, tB::AbstractChar, A::AnyStridedGPUVecOrMat{Complex{T}}, B::AnyStridedGPUVecOrMat{T}, alpha::Number, beta::Number, val::Val{LinearAlgebra.BlasFlag.GEMM}) where T<:Union{Float32, Float64} - LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) end function LinearAlgebra.generic_matmatmul_wrapper!(C::AnyStridedGPUMatrix{Complex{T}}, tA::AbstractChar, tB::AbstractChar, A::AnyStridedGPUVecOrMat{T}, B::AnyStridedGPUVecOrMat{Complex{T}}, alpha::Number, beta::Number, val::Val{LinearAlgebra.BlasFlag.GEMM}) where T<:Union{Float32, Float64} - LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, alpha, beta) + storage_matmatmul!(C, tA, tB, A, B, alpha, beta) end # Julia 1.12 introduced generic_mul! for scalar * array operations function LinearAlgebra.generic_mul!(C::AbstractGPUVecOrMat, X::AbstractGPUVecOrMat, s::Number, alpha::Number, beta::Number)