More tests for sparse matrix dimension checks (#2796)

Author: kshyatt

test/libraries/cusparse/bmm.jl

@@ -15,6 +15,25 @@ if CUSPARSE.version() ≥ v"11.7.2"
     α = rand(elty) 
     β = rand(elty) 
 
+    @testset "Dimension checks" begin
+        A1 = CuSparseMatrixCSR{elty}(sprand(elty, m, k, p))
+        A2 = copy(A1)
+        A2.nzVal = CUDA.rand(elty, size(A2.nzVal)...)
+        A = cat(A1, A2; dims=3)
+
+        B = CUDA.rand(elty, k, n, 2)
+        C = CUDA.rand(elty, m, n, 3)
+
+        @test_throws ArgumentError("C must have same batch-dimension as max(size(A,3)=$(size(A,3)), size(B,3)=$(size(B,3))), got $(size(C,3)).") CUSPARSE.bmm!('N', 'N', α, A, B, β, C, 'O') 
+        
+        C = CUDA.rand(elty, m, 1, 2)
+        @test_throws ArgumentError("bmm! does not work for n==1 and b>1 due to CUDA error.") CUSPARSE.bmm!('N', 'N', α, A, B, β, C, 'O') 
+
+        C = CUDA.rand(elty, m, n, 2)
+        B = CUDA.rand(elty, k+1, n, 2)
+        @test_throws DimensionMismatch("B has dimensions $(size(B)) but needs ($k,$n)") CUSPARSE.bmm!('N', 'N', α, A, B, β, C, 'O') 
+    end
+
     @testset "C = αAB + βC" begin
         A1 = CuSparseMatrixCSR{elty}(sprand(elty, m, k, p))
         A2 = copy(A1)

test/libraries/cusparse/generic.jl

@@ -6,19 +6,28 @@ using LinearAlgebra
 
 if CUSPARSE.version() >= v"11.4.1"
     @testset "generic mv!" for T in [Float32, Float64]
-        A = sprand(T, 10, 10, 0.1)
-        x = rand(Complex{T}, 10)
+        m = 10
+        A = sprand(T, m, m, 0.1)
+        x = rand(Complex{T}, m)
         y = similar(x)
         dx = adapt(CuArray, x)
         dy = adapt(CuArray, y)
 
         dA = adapt(CuArray, A)
-        mv!('N', T(1.0), dA, dx, T(0.0), dy, 'O')
+        mv!('N', one(T), dA, dx, zero(T), dy, 'O')
         @test Array(dy) ≈ A * x
 
         dA = CuSparseMatrixCSR(dA)
-        mv!('N', T(1.0), dA, dx, T(0.0), dy, 'O')
+        mv!('N', one(T), dA, dx, zero(T), dy, 'O')
         @test Array(dy) ≈ A * x
+        
+        A_bad = sprand(T, m+1, m, 0.1)
+        dA_bad = adapt(CuArray, A_bad)
+        @test_throws DimensionMismatch("Y must have length $(m+1), but has length $m") mv!('N', one(T), dA_bad, dx, zero(T), dy, 'O')
+
+        A_bad = sprand(T, m, m+1, 0.1)
+        dA_bad = adapt(CuArray, A_bad)
+        @test_throws DimensionMismatch("X must have length $(m+1), but has length $m") mv!('N', one(T), dA_bad, dx, zero(T), dy, 'O')
     end
 end