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..2724a1e 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} + isempty(C) && return C + 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/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, α, β) 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)