Fix isblasmatrix for GPU StridedViews

Project.toml

@@ -10,14 +10,10 @@ StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
 [weakdeps]
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-cuBLAS = "182d3088-87b7-4494-8cad-fc6afaa545bc"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 
 [extensions]
-StridedcuBLASExt = "cuBLAS"
 StridedGPUArraysExt = "GPUArrays"
-StridedAMDGPUExt = "AMDGPU"
 
 [compat]
 AMDGPU = "2"

ext/StridedGPUArraysExt.jl

@@ -163,15 +163,35 @@ 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 = GPUArrays.derive(T, parent(A2), size(A2), A2.offset)
+    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}
+    # `_mul!` only ensures `stride(C, 1) == 1`; the dense wrapper also needs packed columns
+    isblasmatrix(C) || return Strided.__mul!(C, A, B, α, β)
+    LinearAlgebra.mul!(_densematrix(C), _densematrix(A), _densematrix(B), convert(T, α), convert(T, β))
+    return C
+end
+
 end

test/gpu.jl

@@ -26,6 +26,14 @@ 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, but the columns are not packed
+        @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, but not packed
+        @test !Strided.isblasmatrix(A7)
     end
 end
 
@@ -67,6 +75,15 @@ 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)
     # 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)