Strided SubArray views of CuArray no longer reach cuBLAS in mul! since v11.5.11

[jaemolihm]

Since GPUArrays v11.5.11 (PR #760), mul! with a strided (non-contiguous) view of a CuArray runs GPUArrays' generic gpu_coalesced_matmul_kernel instead of cuBLAS. cuBLAS handles these views natively: CUDA.jl's StridedCuMatrix methods pass the leading dimension, and calling cuBLAS.gemm! on the same view works. For Float32 the generic kernel is about 10× slower here.

I am not sure this should be fixed in the CUDA side or GPUArrays side. But I guess this affects all device backends so I opened an issue here.

MWE

using CUDA, LinearAlgebra

A  = CUDA.rand(Float32, 4096, 1024)
Av = view(A, 1:4000, :)          # strided (non-contiguous) view of a CuMatrix
Ac = copy(Av)                    # contiguous copy of the same data
B  = CUDA.rand(Float32, 1024, 1024)
C  = CUDA.zeros(Float32, 4000, 1024)

display(CUDA.@profile mul!(C, Av, B))                                   # view
display(CUDA.@profile mul!(C, Ac, B))                                   # contiguous copy
display(CUDA.@profile CUDA.cuBLAS.gemm!('N', 'N', 1f0, Av, B, 0f0, C))  # cuBLAS on the view itself

Device-side kernel (NVIDIA RTX A6000):

call	GPUArrays v11.5.8	GPUArrays v11.5.14
mul!(C, Av, B), view	ampere_sgemm_128x64_nn, 0.40 ms	gpu_coalesced_matmul_kernel, 3.98 ms
mul!(C, Ac, B), contiguous copy	ampere_sgemm_128x64_nn, 0.40 ms	ampere_sgemm_128x64_nn, 0.40 ms
cuBLAS.gemm! on Av	ampere_sgemm_128x64_nn, 0.40 ms	ampere_sgemm_128x64_nn, 0.40 ms

[aalquwayfili]

The view branch in _mul! (and mul! before 1.11) calls generic_matmatmul! directly, so it never reaches the storage-level mul!(C, tA, tB, A, B, α, β) hook that CUDA.jl overloads since JuliaGPU/CUDA.jl#3257. Forwarding to that hook instead (wrapper_char plus _unwrap on the operands) brings the MWE back to sgemm (1.6 ms instead of 15.7 ms on an RTX 8000), and backends without their own hook still land in GPUArrays' generic method, so JLArrays keeps working.

Doing only that isn't safe yet, though. CUDA's hook assumes contiguous columns and cuBLAS element types, and nothing sent views there before, so two bugs were never hit. Calling the hook directly on current main:

• a row-step view like view(A, 1:2:8, :) goes to gemm!/gemv! and gives wrong results
• Symmetric of a Float16 view goes to symm! and throws a MethodError (the case @christiangnrd mentioned in Fix multiplication with strided GPU array views #760)

With the forwarding alone, 13 of the linalg/mul!/strided-views tests fail on CuArray for exactly these. Guarding the hook so it only calls cuBLAS when every operand has stride(X, 1) == 1 and, for symm/hemm, a CublasFloat element type shared by all operands, fixes both. It falls through to the generic kernel it already ends with otherwise.

With both changes, every strided-view case in the testsuite passes on CuArray and JLArray, and views reach gemm, gemv, symm, hemm, and CUTLASS for Float16. Only Int32 and mixed-eltype products stay generic.

Would two PRs work, the CUDA.jl guard first and the GPUArrays forwarding after it's released? The GPUArrays change alone would give CUDA users wrong results for row-step views. If you'd rather make views opt-in per backend instead, I can do it that way.

[jaemolihm]

Thanks for the comments! I don't know much about the GPUArrays/CUDA packages. As a user, any fix that lets mul!(C, Av, B) reach the optimized method would work.