From 77b9aafa213a3f650e24f4487053a94391f08f0d Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 10:34:32 -0400 Subject: [PATCH 1/5] Fix isblasmatrix for GPU StridedViews --- Project.toml | 4 ++-- ext/StridedAMDGPUExt.jl | 27 ++++++++++----------- ext/StridedGPUArraysExt.jl | 29 ++++++++++++++++++---- ext/StridedcuBLASExt.jl | 19 +++++++++------ test/gpu.jl | 49 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 100 insertions(+), 28 deletions(-) diff --git a/Project.toml b/Project.toml index 5cb347a..35ea02b 100644 --- a/Project.toml +++ b/Project.toml @@ -15,9 +15,9 @@ cuBLAS = "182d3088-87b7-4494-8cad-fc6afaa545bc" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" [extensions] -StridedcuBLASExt = "cuBLAS" +StridedcuBLASExt = ["GPUArrays", "cuBLAS"] StridedGPUArraysExt = "GPUArrays" -StridedAMDGPUExt = "AMDGPU" +StridedAMDGPUExt = ["AMDGPU", "GPUArrays"] [compat] AMDGPU = "2" diff --git a/ext/StridedAMDGPUExt.jl b/ext/StridedAMDGPUExt.jl index e14d804..464e99e 100644 --- a/ext/StridedAMDGPUExt.jl +++ b/ext/StridedAMDGPUExt.jl @@ -1,27 +1,24 @@ module StridedAMDGPUExt -using Strided, StridedViews, AMDGPU, AMDGPU.rocBLAS, LinearAlgebra -import Strided: blas_mul!, _get_op +using Strided, StridedViews, GPUArrays, AMDGPU, AMDGPU.rocBLAS, LinearAlgebra +import Strided: blas_mul!, isblasmatrix, _get_op const ROCStridedView{T, N, A <: ROCArray{T}} = StridedViews.StridedView{T, N, A} +# A view of the `(lda, n)` matrix keeps the column stride as the leading dimension. The last +# column might not span `lda` elements in the parent, but `gemm!` only accesses `1:m`. +function _stridedstorage(A::ROCStridedView{T, 2}) where {T} + m, n = size(A) + lda = max(stride(A, 2), m, 1) + return view(GPUArrays.derive(T, parent(A), (lda, n), A.offset), 1:m, :) +end + function Strided.blas_mul!(C::ROCStridedView{T, 2}, A::ROCStridedView{T, 2}, B::ROCStridedView{T, 2}, α::Number, β::Number) where {T <: LinearAlgebra.BlasFloat} + isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) A2, CA = Strided.getblasmatrix(A) B2, CB = Strided.getblasmatrix(B) - C2, CC = Strided.getblasmatrix(C) - A2a = Base.unsafe_wrap(ROCMatrix{T}, pointer(A2), size(A2)) - B2a = Base.unsafe_wrap(ROCMatrix{T}, pointer(B2), size(B2)) - C2a = Base.unsafe_wrap(ROCMatrix{T}, pointer(C2), size(C2)) - AMDGPU.rocBLAS.gemm!(CA, CB, convert(T, α), A2a, B2a, convert(T, β), C2a) + AMDGPU.rocBLAS.gemm!(CA, CB, convert(T, α), _stridedstorage(A2), _stridedstorage(B2), convert(T, β), _stridedstorage(C)) return C end -_conj(x) = real(x) - imag(x) * im -@static if VERSION < v"1.11.0-rc" - # work around compiler issue on AMD on 1.10 - _get_op(A::ROCStridedView) = A.op == conj ? _conj : A.op -else - _get_op(A::ROCStridedView) = A.op -end - end diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index d34465d..17ddf80 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/ext/StridedGPUArraysExt.jl @@ -163,15 +163,36 @@ end function Strided.isblasmatrix(A::GPUStridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} if A.op == identity - # unsafe wrap approach doesn't work if second condition not met - return stride(A, 1) == 1 && size(A, 1) == size(parent(A), 1) + if stride(A, 1) == 1 + return stride(A, 2) ≥ size(A, 1) || size(A, 2) <= 1 + elseif stride(A, 2) == 1 + return stride(A, 1) ≥ size(A, 2) || size(A, 1) <= 1 + else + return false + end elseif A.op == conj # this is converted to adjoint - # unsafe wrap approach doesn't work if second condition not met - return stride(A, 2) == 1 && size(A, 2) == size(parent(A), 2) + return stride(A, 2) == 1 && (stride(A, 1) ≥ size(A, 2) || size(A, 1) <= 1) else # should never happen return false end end +function _densematrix(A::GPUStridedView{T, 2}) where {T} + A2, CA = Strided.getblasmatrix(A) + m, n = size(A2) + lda = max(stride(A2, 2), m, 1) + M = view(GPUArrays.derive(T, parent(A2), (lda, n), A2.offset), 1:m, 1:n) + return CA == 'N' ? M : CA == 'T' ? transpose(M) : adjoint(M) +end + +function Strided.blas_mul!( + C::GPUStridedView{T, 2}, A::GPUStridedView{T, 2}, B::GPUStridedView{T, 2}, + α::Number, β::Number + ) where {T <: LinearAlgebra.BlasFloat} + isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) + LinearAlgebra.mul!(_densematrix(C), _densematrix(A), _densematrix(B), convert(T, α), convert(T, β)) + return C +end + end diff --git a/ext/StridedcuBLASExt.jl b/ext/StridedcuBLASExt.jl index 3c41a2c..38c0709 100644 --- a/ext/StridedcuBLASExt.jl +++ b/ext/StridedcuBLASExt.jl @@ -1,18 +1,23 @@ module StridedcuBLASExt -using Strided, StridedViews, cuBLAS, cuBLAS.CUDACore, LinearAlgebra -import Strided: blas_mul! +using Strided, StridedViews, GPUArrays, cuBLAS, cuBLAS.CUDACore, LinearAlgebra +import Strided: blas_mul!, isblasmatrix const CuStridedView{T, N, A <: CuArray{T}} = StridedViews.StridedView{T, N, A} +# A view of the `(lda, n)` matrix keeps the column stride as the leading dimension. The last +# column might not span `lda` elements in the parent, but `gemm!` only accesses `1:m`. +function _stridedstorage(A::CuStridedView{T, 2}) where {T} + m, n = size(A) + lda = max(stride(A, 2), m, 1) + return view(GPUArrays.derive(T, parent(A), (lda, n), A.offset), 1:m, :) +end + function Strided.blas_mul!(C::CuStridedView{T, 2}, A::CuStridedView{T, 2}, B::CuStridedView{T, 2}, α::Number, β::Number) where {T <: LinearAlgebra.BlasFloat} + isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) A2, CA = Strided.getblasmatrix(A) B2, CB = Strided.getblasmatrix(B) - C2, CC = Strided.getblasmatrix(C) - A2a = Base.unsafe_wrap(CuMatrix{T}, pointer(A2), size(A2)) - B2a = Base.unsafe_wrap(CuMatrix{T}, pointer(B2), size(B2)) - C2a = Base.unsafe_wrap(CuMatrix{T}, pointer(C2), size(C2)) - cuBLAS.gemm!(CA, CB, convert(T, α), A2a, B2a, convert(T, β), C2a) + cuBLAS.gemm!(CA, CB, convert(T, α), _stridedstorage(A2), _stridedstorage(B2), convert(T, β), _stridedstorage(C)) return C end diff --git a/test/gpu.jl b/test/gpu.jl index 9a19d3d..ab38519 100644 --- a/test/gpu.jl +++ b/test/gpu.jl @@ -26,6 +26,22 @@ Metal.functional() && push!(ATs, MtlArray) @test !Strided.isblasmatrix(A2) A3 = view(conj!(A1), 1:4:20, 1:20) # stride(A3, 2) is not 1 @test !Strided.isblasmatrix(A3) + A4 = view(A1, 1:20, 1:2:20) # unit stride, columns are not packed but do not overlap + @test Strided.isblasmatrix(A4) + A5 = sreshape(StridedView(AT(randn(T, 4, 5, 20))), (20, 20)) # parent has 3 dims + @test Strided.isblasmatrix(A5) + A6 = permutedims(A1, (2, 1)) # packed rows, used as 'T' + @test Strided.isblasmatrix(A6) + A7 = permutedims(view(A1, 1:10, 1:20), (2, 1)) # unit stride along rows, not packed + @test Strided.isblasmatrix(A7) + A8 = adjoint(view(A1, 1:10, 1:20)) # conj op with unit stride along rows, not packed + @test Strided.isblasmatrix(A8) + if T <: Complex + A9 = conj(view(A1, 1:10, 1:20)) # conj op without a transpose has no BLAS equivalent + @test !Strided.isblasmatrix(A9) + end + @test Strided.isblasmatrix(StridedView(AT(zeros(T, 0, 5)))) # empty matrices + @test Strided.isblasmatrix(StridedView(AT(zeros(T, 0, 0)))) end end @@ -67,6 +83,39 @@ end @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, vA1, f1(vA2), f2(vA3)) end end + # matrices reshaped out of higher-dimensional parents + rA1, rA2, rA3 = (StridedView(rand(T, (4, 8, 32))) for _ in 1:3) + @test compare(AT, rA1, rA2, rA3) do C, A, B + return mul!(sreshape(C, (32, 32)), sreshape(A, (32, 32)), sreshape(B, (32, 32)), α, β) + end + # packed A and B, but C has unit stride without packed columns + vA1 = view(StridedView(rand(T, (64, 32))), 1:32, 1:32) + vA2, vA3 = (StridedView(rand(T, (32, 32))) for _ in 1:2) + @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, vA1, vA2, vA3) + # unit stride and non-packed columns and all transposition combinations + pA1, pA2, pA3 = (StridedView(rand(T, (64, 64))) for _ in 1:3) + for (r1, r2) in ((1:32, 1:32), (33:64, 33:64)) + vA1 = view(pA1, 1:32, 1:32) + vA2 = view(pA2, r1, 1:32) + vA3 = view(pA3, r2, 33:64) + @testset for f1 in (identity, conj, adjoint, transpose), f2 in (identity, conj, adjoint, transpose) + @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, vA1, f1(vA2), f2(vA3)) + end + end + # conjugated views: non-packed adjoints as inputs, and conjugated non-packed destination + vA1, vA2, vA3 = (view(p, 1:32, 1:32) for p in (pA1, pA2, pA3)) + @test compare((C, A, B) -> mul!(C, A', B', α, β), AT, vA1, vA2, vA3) + @test compare((C, A, B) -> mul!(C, conj(A), B, α, β), AT, vA1, vA2, vA3) + @test compare((C, A, B) -> mul!(conj(C), A', B, α, β), AT, vA1, vA2, vA3) + @test compare((C, A, B) -> mul!(C', A, B, α, β), AT, vA1, vA2, vA3) + # non-packed destination + vC = view(StridedView(rand(T, (64, 64))), 33:64, 1:32) + @test compare((C, A, B) -> mul!(C, A, B, α, β), AT, vC, view(pA2, 1:32, 1:32), view(pA3, 1:32, 1:32)) + # empty matrices + @test compare( + (C, A, B) -> mul!(C, A, B, α, β), AT, + StridedView(zeros(T, 0, 5)), StridedView(zeros(T, 0, 3)), StridedView(zeros(T, 3, 5)) + ) # non-BLAS fallback path vA1 = view(StridedView(rand(T, (32, 32))), 1:32, 1:32) vA2 = view(StridedView(rand(T, (32, 64))), 1:32, 1:2:64) From bffd5ad9fe2b5eb89c691c3a4b6998050078de14 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 16:44:55 +0100 Subject: [PATCH 2/5] Exit early for empty matrices --- ext/StridedGPUArraysExt.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index 17ddf80..0a0b9a1 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/ext/StridedGPUArraysExt.jl @@ -190,6 +190,7 @@ function Strided.blas_mul!( C::GPUStridedView{T, 2}, A::GPUStridedView{T, 2}, B::GPUStridedView{T, 2}, α::Number, β::Number ) where {T <: LinearAlgebra.BlasFloat} + isempty(C) && return C isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) LinearAlgebra.mul!(_densematrix(C), _densematrix(A), _densematrix(B), convert(T, α), convert(T, β)) return C From cde4c197fb08666ef69804c19d8fee8b391e03e1 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 16:56:18 +0100 Subject: [PATCH 3/5] Use isblasmatrix check for C --- ext/StridedGPUArraysExt.jl | 1 - src/linalg.jl | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index 0a0b9a1..2724a1e 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/ext/StridedGPUArraysExt.jl @@ -191,7 +191,6 @@ function Strided.blas_mul!( α::Number, β::Number ) where {T <: LinearAlgebra.BlasFloat} isempty(C) && return C - isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) LinearAlgebra.mul!(_densematrix(C), _densematrix(A), _densematrix(B), convert(T, α), convert(T, β)) return C end diff --git a/src/linalg.jl b/src/linalg.jl index 4a1bbbe..1cfdd63 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -99,7 +99,7 @@ function _mul!( C::StridedView{T, 2}, A::StridedView{T, 2}, B::StridedView{T, 2}, α::Number, β::Number ) where {T <: LinearAlgebra.BlasFloat} - if stride(C, 1) == 1 && isblasmatrix(A) && isblasmatrix(B) + if isblasmatrix(A) && isblasmatrix(B) && isblasmatrix(C) return blas_mul!(C, A, B, α, β) else return __mul!(C, A, B, α, β) From 0f9428c6a599a7e82c99abe0955fb306779b2022 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 14:24:52 -0400 Subject: [PATCH 4/5] Consolidate --- ext/StridedGPUArraysExt.jl | 19 +------------------ src/linalg.jl | 1 - 2 files changed, 1 insertion(+), 19 deletions(-) diff --git a/ext/StridedGPUArraysExt.jl b/ext/StridedGPUArraysExt.jl index 2724a1e..3b1db6f 100644 --- a/ext/StridedGPUArraysExt.jl +++ b/ext/StridedGPUArraysExt.jl @@ -5,7 +5,7 @@ using GPUArrays: Adapt, KernelAbstractions using GPUArrays.KernelAbstractions: @kernel, @index using StridedViews: ParentIndex -import Strided: isblasmatrix, _get_op +import Strided: _get_op ALL_FS = Union{typeof(adjoint), typeof(conj), typeof(identity), typeof(transpose)} @@ -161,23 +161,6 @@ function Strided._mapreduce_block!( return nothing end -function Strided.isblasmatrix(A::GPUStridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} - if A.op == identity - if stride(A, 1) == 1 - return stride(A, 2) ≥ size(A, 1) || size(A, 2) <= 1 - elseif stride(A, 2) == 1 - return stride(A, 1) ≥ size(A, 2) || size(A, 1) <= 1 - else - return false - end - elseif A.op == conj - # this is converted to adjoint - return stride(A, 2) == 1 && (stride(A, 1) ≥ size(A, 2) || size(A, 1) <= 1) - else # should never happen - return false - end -end - function _densematrix(A::GPUStridedView{T, 2}) where {T} A2, CA = Strided.getblasmatrix(A) m, n = size(A2) diff --git a/src/linalg.jl b/src/linalg.jl index 1cfdd63..d93d357 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -73,7 +73,6 @@ function LinearAlgebra.mul!( end function isblasmatrix(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} - pointer(A) isa Ptr || return false if A.op == identity return stride(A, 1) == 1 || stride(A, 2) == 1 elseif A.op == conj From a0216a4416a7e7a26dd57fac1391803f498c19c6 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sat, 31 Oct 2026 02:48:16 -0400 Subject: [PATCH 5/5] getblasmatrix for all --- ext/StridedcuBLASExt.jl | 18 ++++++++++-------- src/linalg.jl | 3 ++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/ext/StridedcuBLASExt.jl b/ext/StridedcuBLASExt.jl index 38c0709..8ae8fd8 100644 --- a/ext/StridedcuBLASExt.jl +++ b/ext/StridedcuBLASExt.jl @@ -1,23 +1,25 @@ module StridedcuBLASExt using Strided, StridedViews, GPUArrays, cuBLAS, cuBLAS.CUDACore, LinearAlgebra -import Strided: blas_mul!, isblasmatrix +import Strided: blas_mul!, isblasmatrix, getblasmatrix const CuStridedView{T, N, A <: CuArray{T}} = StridedViews.StridedView{T, N, A} # A view of the `(lda, n)` matrix keeps the column stride as the leading dimension. The last # column might not span `lda` elements in the parent, but `gemm!` only accesses `1:m`. -function _stridedstorage(A::CuStridedView{T, 2}) where {T} - m, n = size(A) - lda = max(stride(A, 2), m, 1) - return view(GPUArrays.derive(T, parent(A), (lda, n), A.offset), 1:m, :) +function Strided.getblasmatrix(A::CuStridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} + A2, CA = Strided._blasview(A) + m, n = size(A2) + lda = max(stride(A2, 2), m, 1) + return view(GPUArrays.derive(T, parent(A2), (lda, n), A2.offset), 1:m, :), CA end function Strided.blas_mul!(C::CuStridedView{T, 2}, A::CuStridedView{T, 2}, B::CuStridedView{T, 2}, α::Number, β::Number) where {T <: LinearAlgebra.BlasFloat} isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β) - A2, CA = Strided.getblasmatrix(A) - B2, CB = Strided.getblasmatrix(B) - cuBLAS.gemm!(CA, CB, convert(T, α), _stridedstorage(A2), _stridedstorage(B2), convert(T, β), _stridedstorage(C)) + A2, CA = getblasmatrix(A) + B2, CB = getblasmatrix(B) + C2, _ = getblasmatrix(C) # `_mul!` guarantees `C.op == identity && stride(C, 1) < stride(C, 2)` + cuBLAS.gemm!(CA, CB, convert(T, α), A2, B2, convert(T, β), C2) return C end diff --git a/src/linalg.jl b/src/linalg.jl index d93d357..02aa920 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -81,7 +81,8 @@ function isblasmatrix(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} return false end end -function getblasmatrix(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} +getblasmatrix(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} = _blasview(A) +function _blasview(A::StridedView{T, 2}) where {T <: LinearAlgebra.BlasFloat} if A.op == identity if stride(A, 1) == 1 return A, 'N'