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