Fix isblasmatrix for GPU StridedViews

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"

ext/StridedAMDGPUExt.jl

@@ -1,27 +1,17 @@
 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, getblasmatrix
 
 const ROCStridedView{T, N, A <: ROCArray{T}} = StridedViews.StridedView{T, N, A}
 
 function Strided.blas_mul!(C::ROCStridedView{T, 2}, A::ROCStridedView{T, 2}, B::ROCStridedView{T, 2}, α::Number, β::Number) where {T <: LinearAlgebra.BlasFloat}
-    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)
+    isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β)
+    A2, CA = getblasmatrix(A)
+    B2, CB = getblasmatrix(B)
+    C2, _ = getblasmatrix(C)
+    AMDGPU.rocBLAS.gemm!(CA, CB, convert(T, α), A2, B2, convert(T, β), C2)
     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

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,17 +161,25 @@ function Strided._mapreduce_block!(
     return nothing
 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)
-    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)
-    else # should never happen
-        return false
-    end
+function Strided.getblasmatrix(A::GPUStridedView{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
+
+_applyblasop(M, c::Char) = c == 'N' ? M : c == 'T' ? transpose(M) : adjoint(M)
+
+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
+    A2, CA = Strided.getblasmatrix(A)
+    B2, CB = Strided.getblasmatrix(B)
+    C2, _ = Strided.getblasmatrix(C) # `_mul!` guarantees `C.op == identity && stride(C, 1) < stride(C, 2)`
+    LinearAlgebra.mul!(C2, _applyblasop(A2, CA), _applyblasop(B2, CB), convert(T, α), convert(T, β))
+    return C
 end
 
 end

ext/StridedcuBLASExt.jl

@@ -1,18 +1,15 @@
 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, getblasmatrix
 
 const CuStridedView{T, N, A <: CuArray{T}} = StridedViews.StridedView{T, N, A}
 
 function Strided.blas_mul!(C::CuStridedView{T, 2}, A::CuStridedView{T, 2}, B::CuStridedView{T, 2}, α::Number, β::Number) where {T <: LinearAlgebra.BlasFloat}
-    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)
+    A2, CA = getblasmatrix(A)
+    B2, CB = getblasmatrix(B)
+    C2, _ = getblasmatrix(C)
+    cuBLAS.gemm!(CA, CB, convert(T, α), A2, B2, convert(T, β), C2)
     return C
 end
 

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
@@ -82,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'
@@ -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, α, β)

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)