From 97ef6e87bef05d07d4b0f81a851c5e6010047a26 Mon Sep 17 00:00:00 2001 From: Abdulmalik Alquwayfili Date: Thu, 15 Oct 2026 21:00:29 +0200 Subject: [PATCH] Forward strided-view matrix products to the storage-level mul! MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since #760, any product with a strided GPU view goes to the generic matmul kernel, which bypasses the vendor libraries back-ends hook up through the storage-level `mul!(C, tA, tB, A, B, α, β)`. For the view(A, 1:4000, :) * B case in #799 that is about 10x slower than cuBLAS. Forward plain matrix-matrix products (wrapper chars N, T or C) to that method again when every matrix is a valid BLAS operand: unit row stride and a leading dimension of at least the number of rows. The gemm wrappers in the back-ends take stride(X, 2) as the leading dimension without validating it, so row-stepping views would give wrong results and reversed columns invalid arguments. Dimension mismatches also stay on the generic kernel, which reports them, since Metal.jl's native method does not. Symmetric and Hermitian operands and matrix-vector products keep the generic kernels, as AMDGPU.jl's symm!, hemm! and gemv! do not accept views. Fixes #799. Co-authored-by: Tim Besard --- src/host/linalg.jl | 27 ++++++++++++++++++++++++--- test/testsuite/linalg.jl | 13 +++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/src/host/linalg.jl b/src/host/linalg.jl index a873f48f..c9bbadbd 100644 --- a/src/host/linalg.jl +++ b/src/host/linalg.jl @@ -525,7 +525,8 @@ end ## matrix multiplication # GPU-backed members of Base's StridedArray union match LinearAlgebra's BLAS methods, -# but cannot be converted to host pointers. Route them to the generic GPU kernels. +# but cannot be converted to host pointers. Route them to the generic GPU kernels, or to the +# storage-level `mul!` methods where the back-end can take them (see `strided_view_matmatmul!`). const StridedGPUSubArray{T,N} = Base.StridedSubArray{T,N,<:AbstractGPUArray} const AnyStridedGPUArray{T,N} = Union{AbstractGPUArray{T,N},StridedGPUSubArray{T,N}} const AnyStridedGPUVector{T} = AnyStridedGPUArray{T,1} @@ -546,6 +547,26 @@ const AnyStridedGPUVecOrMatOperand{T} = Union{ has_strided_gpu_view(As...) = any(A -> LinearAlgebra._unwrap(A) isa StridedGPUSubArray, As) +# Back-ends overload the storage-level `mul!` to call their vendor library. Their gemm wrappers +# accept views, but some pass `stride(X, 2)` as the leading dimension without validating the +# layout, so only forward views that are valid BLAS operands: unit row stride and a leading +# dimension of at least the number of rows (which rules out reversed columns). Not every +# back-end checks dimensions either, so leave mismatches to the generic kernel to report. +# Symmetric and Hermitian operands and products with a vector keep the generic kernels, as not +# every back-end accepts views for those. +blas_compatible(X) = + X isa AbstractMatrix && stride(X, 1) == 1 && stride(X, 2) >= max(1, size(X, 1)) + +function strided_view_matmatmul!(C, A, B, a, b) + tA, tB = LinearAlgebra.wrapper_char(A), LinearAlgebra.wrapper_char(B) + pA, pB = LinearAlgebra._unwrap(A), LinearAlgebra._unwrap(B) + if tA in ('N', 'T', 'C') && tB in ('N', 'T', 'C') && all(blas_compatible, (C, pA, pB)) && + size(A, 2) == size(B, 1) && size(C) == (size(A, 1), size(B, 2)) + return LinearAlgebra.mul!(C, tA, tB, pA, pB, a, b) + end + return generic_matmatmul!(C, A, B, a, b) +end + # Intercept before LinearAlgebra unwraps operands and dispatches to backend BLAS methods. The # signatures also cover view-free products to avoid overlapping methods for each possible view # position; those calls are sent back through LinearAlgebra's original implementation. @@ -564,7 +585,7 @@ end @inline function LinearAlgebra.mul!(C::AnyStridedGPUMatrix, A::AnyStridedGPUVecOrMatOperand, B::AnyStridedGPUVecOrMatOperand, a::Number, b::Number) if has_strided_gpu_view(C, A, B) - return generic_matmatmul!(C, A, B, a, b) + return strided_view_matmatmul!(C, A, B, a, b) end invoke(LinearAlgebra.mul!, Tuple{AbstractMatrix,LinearAlgebra.AbstractVecOrMat, @@ -587,7 +608,7 @@ end @inline function LinearAlgebra._mul!(C::AnyStridedGPUMatrix, A::AnyStridedGPUVecOrMatOperand, B::AnyStridedGPUVecOrMatOperand, a::Number, b::Number) if has_strided_gpu_view(C, A, B) - return generic_matmatmul!(C, A, B, a, b) + return strided_view_matmatmul!(C, A, B, a, b) end invoke(LinearAlgebra._mul!, Tuple{AbstractMatrix,LinearAlgebra.AbstractVecOrMat, diff --git a/test/testsuite/linalg.jl b/test/testsuite/linalg.jl index d1a0e1c2..23251cdf 100644 --- a/test/testsuite/linalg.jl +++ b/test/testsuite/linalg.jl @@ -763,12 +763,25 @@ end C end + # reversed columns (negative stride along the second dimension) + @test compare((A, B) -> view(A, :, 6:-1:1) * view(B, 6:-1:1, :), AT, As, Bs) + @test compare((A, B) -> A * view(B, :, 10:-1:1), AT, As, Bs) + @test compare(A -> view(A, :, 6:-1:1)' * A, AT, As) + Cr = rand(T, 8, 10) + @test compare(AT, Cr, As, Bs) do C, A, B + mul!(view(C, :, 10:-1:1), A, B, T(2), T(1)) + C + end + # strided destination with contiguous operands B3 = rand(T, 3, 5) @test compare(AT, C, A2, B3) do C, A, B mul!(view(C, 1:2, :), A, B) C end + + # mismatched destination + @test_throws DimensionMismatch mul!(AT(rand(T, 1, 5)), view(AT(A), 1:2, :), AT(B3)) end if Float32 in eltypes && ComplexF32 in eltypes