Skip to content

Storage-level mul! gives wrong results or MethodErrors for SubArray operands #1161

Open

Description

@maleadt

The storage-level mul! methods in src/blas/highlevel.jl pass SubArrays of ROCArrays to rocBLAS without checking their layout, and some of the wrappers they call don't accept views at all:

using AMDGPU, LinearAlgebra
A = ROCArray(rand(Float32, 8, 4)); B = ROCArray(rand(Float32, 4, 5)); x = ROCArray(rand(Float32, 4))
S = ROCArray(rand(Float32, 8, 8))

# row-stepping view: silently wrong result (lda = stride(A, 2) and stride(A, 1) is ignored)
C = mul!(ROCArray(zeros(Float32, 4, 5)), 'N', 'N', view(A, 1:2:8, :), B, true, false)
Array(C) ≈ Array(A)[1:2:8, :] * Array(B)   # false

# reversed columns: ROCBLASError(rocblas_status_invalid_size)
mul!(ROCArray(zeros(Float32, 8, 5)), 'N', 'N', view(A, :, 4:-1:1), B, true, false)

# any view, even a contiguous one, on the gemv/symv/hemv/symm/hemm paths: MethodError
mul!(ROCArray(zeros(Float32, 4)), 'N', view(A, 1:4, :), x, true, false)       # no method matching gemv!
mul!(ROCArray(zeros(Float32, 4, 5)), 'S', 'N', view(S, 1:4, 1:4), B, true, false)   # no method matching symm!

The mul! methods take StridedROCMatrix/StridedROCVector. gemm! and gemm_ex! take strided arrays too, but compute lda = max(1, stride(A, 2)) without checking stride(A, 1). gemv!, symv!, hemv!, symm! and hemm! (src/blas/wrappers.jl) only accept dense ROCMatrix/ROCVector. The 'S'/'H' branch also has no check that eltype(A) == eltype(B) == eltype(C), so mixed element types end up in a MethodError too.

The public mul!(C, A, B) isn't affected, because GPUArrays intercepts products with views before they reach these methods. JuliaGPU/GPUArrays.jl#800 starts forwarding matrix-matrix products with views here again, but only N/T/C products where every matrix has a unit row stride and stride(X, 2) >= size(X, 1). It keeps Symmetric/Hermitian and matrix-vector products away because of the MethodErrors above. Calling the storage-level mul! directly is public API since Julia 1.13 (JuliaLang/LinearAlgebra.jl#1671), though. CUDA.jl added the equivalent guards in JuliaGPU/CUDA.jl#3316:

  • only call rocBLAS when every matrix has stride(X, 1) == 1 and stride(X, 2) >= max(1, size(X, 1)), and every vector has a positive stride;
  • check that element types match on the symm/hemm branch;
  • widen the BLAS 2 and symm/hemm wrappers to strided arrays, or fall back for views;
  • send everything else to GPUArrays.generic_matmatmul!.

With those in place, GPUArrays could forward Symmetric, Hermitian and matrix-vector products with views as well.

Reproduced with AMDGPU.jl 2.9.1 and GPUArrays 12.1 on a gfx1036 iGPU, Julia 1.12.7.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions