Set default buffer size in CUSPARSE mm! functions (#2298)

lpawela · amontoison · web-flow · commit ad5ddf820214 · 2024-03-21T14:08:35.000+01:00
This works around a bug in CUSPARSE.

Co-authored-by: Alexis Montoison &lt;35051714+amontoison@users.noreply.github.com&gt;
diff --git a/lib/cusparse/generic.jl b/lib/cusparse/generic.jl
@@ -235,8 +235,10 @@ function mm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::Union{CuS
     #     cusparseCsrSetStridedBatch(obj, batchsize, 0, nnz(A))
     # end
 
+    # Set default buffer for small matrices (10000 chosen arbitrarly)
+    # Otherwise tries to allocate 120TB of memory (see #2296)
     function bufferSize()
-        out = Ref{Csize_t}()
+        out = Ref{Csize_t}(10000) 
         cusparseSpMM_bufferSize(
             handle(), transa, transb, Ref{T}(alpha), descA, descB, Ref{T}(beta),
             descC, T, algo, out)
@@ -311,8 +313,10 @@ function bmm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::CuSparse
     strideC = stride(C, 3)
     cusparseDnMatSetStridedBatch(descC, b, strideC)
 
+    # Set default buffer for small matrices (10000 chosen arbitrarly)
+    # Otherwise tries to allocate 120TB of memory (see #2296)
     function bufferSize()
-        out = Ref{Csize_t}()
+        out = Ref{Csize_t}(10000)
         cusparseSpMM_bufferSize(
             handle(), transa, transb, Ref{T}(alpha), descA, descB, Ref{T}(beta),
             descC, T, algo, out)
@@ -337,7 +341,6 @@ function mm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::DenseCuMa
              beta::Number, C::DenseCuMatrix{T}, index::SparseChar, algo::cusparseSpMMAlg_t=CUSPARSE_SPMM_ALG_DEFAULT) where {T}
 
     CUSPARSE.version() < v"11.7.4" && throw(ErrorException("This operation is not supported by the current CUDA version."))
-
     # Support transa = 'C' and `transb = 'C' for real matrices
     transa = T <: Real && transa == 'C' ? 'T' : transa
     transb = T <: Real && transb == 'C' ? 'T' : transb
@@ -370,8 +373,10 @@ function mm!(transa::SparseChar, transb::SparseChar, alpha::Number, A::DenseCuMa
     descB = CuSparseMatrixDescriptor(B, index, transposed=true)
     descC = CuDenseMatrixDescriptor(C, transposed=true)
 
+    # Set default buffer for small matrices (10000 chosen arbitrarly)
+    # Otherwise tries to allocate 120TB of memory (see #2296)
     function bufferSize()
-        out = Ref{Csize_t}()
+        out = Ref{Csize_t}(10000)
         cusparseSpMM_bufferSize(
             handle(), transb, transa, Ref{T}(alpha), descB, descA, Ref{T}(beta),
             descC, T, algo, out)
diff --git a/test/libraries/cusparse/bmm.jl b/test/libraries/cusparse/bmm.jl
@@ -71,6 +71,120 @@ if CUSPARSE.version() ≥ v"11.7.2"
     end
 
 
+    @testset "C = αAᵀBᵀ + βC" begin
+        A1 = CuSparseMatrixCSR{elty}(sprand(elty, k, m, p))
+        A2 = copy(A1)
+        A2.nzVal = CUDA.rand(elty, size(A2.nzVal)...)
+        A = cat(A1, A2; dims=3)
+
+        B = CUDA.rand(elty, n, k, 2)
+        C = CUDA.rand(elty, m, n, 2)
+        D = copy(C)
+
+        CUSPARSE.bmm!('C', 'C', α, A, B, β, C, 'O') 
+
+        D[:,:,1] = α * A1' * B[:,:,1]' + β * D[:,:,1]
+        D[:,:,2] = α * A2' * B[:,:,2]' + β * D[:,:,2]
+
+        @test D ≈ C
+    end
+
+    @testset "extended batch-dims" begin
+        A1 = CuSparseMatrixCSR{elty}(sprand(elty, m, k, p))
+        A2 = copy(A1)
+        A2.nzVal = CUDA.rand(elty, size(A2.nzVal)...)
+        A3 = cat(A1, A2; dims=3)
+
+        A4 = copy(A3)
+        A4.nzVal = CUDA.rand(elty, size(A3.nzVal)...)
+
+        A5 = copy(A3)
+        A5.nzVal = CUDA.rand(elty, size(A3.nzVal)...)
+
+        A = cat(A3, A4, A5; dims=4)
+
+        B = CUDA.rand(elty, k, n, 2, 3)
+        C = CUDA.rand(elty, m, n, 2, 3)
+        D = copy(C)
+
+        CUSPARSE.bmm!('N', 'N', α, A, B, β, C, 'O') 
+
+        for c in CartesianIndices((2,3))
+            CUDA.@allowscalar D[:,:,c] = α * A[:,:,c.I...] * B[:,:,c] + β*D[:,:,c]
+        end
+
+        @test D ≈ C
+    end
+end
+
+m = 1
+n = 2
+# error when n == 1 and batchsize > 1 as cusparseSpMM fallsback to cusparseSpMV, which doesn't do batched computations.
+# see https://docs.nvidia.com/cuda/cusparse/#cusparsespmm
+k = 1
+p = 1.
+
+@testset "Sparse-Dense $elty bmm! for small matrices" for elty in (Float64, Float32, ComplexF64, ComplexF32)
+    # check if #2296 returns
+    α = rand(elty) 
+    β = rand(elty) 
+
+    @testset "C = αAB + βC" 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, 2)
+        D = copy(C)
+
+        CUSPARSE.bmm!('N', 'N', α, A, B, β, C, 'O') 
+
+        D[:,:,1] = α * A1 * B[:,:,1] + β * D[:,:,1]
+        D[:,:,2] = α * A2 * B[:,:,2] + β * D[:,:,2]
+
+        @test D ≈ C
+    end
+
+    @testset "C = αAᵀB + βC" begin
+        A1 = CuSparseMatrixCSR{elty}(sprand(elty, k, m, 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, 2)
+        D = copy(C)
+
+        CUSPARSE.bmm!('C', 'N', α, A, B, β, C, 'O') 
+
+        D[:,:,1] = α * A1' * B[:,:,1] + β * D[:,:,1]
+        D[:,:,2] = α * A2' * B[:,:,2] + β * D[:,:,2]
+
+        @test D ≈ C
+    end
+
+
+    @testset "C = αABᵀ + βC" 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, n, k, 2)
+        C = CUDA.rand(elty, m, n, 2)
+        D = copy(C)
+
+        CUSPARSE.bmm!('N', 'C', α, A, B, β, C, 'O') 
+
+        D[:,:,1] = α * A1 * B[:,:,1]' + β * D[:,:,1]
+        D[:,:,2] = α * A2 * B[:,:,2]' + β * D[:,:,2]
+
+        @test D ≈ C
+    end
+
+
     @testset "C = αAᵀBᵀ + βC" begin
         A1 = CuSparseMatrixCSR{elty}(sprand(elty, k, m, p))
         A2 = copy(A1)
diff --git a/test/libraries/cusparse/generic.jl b/test/libraries/cusparse/generic.jl
@@ -134,6 +134,22 @@ if CUSPARSE.version() >= v"11.7.4"
                         beta = rand(T)
                         mm!(transa, transb, alpha, dA, dB, beta, dC, 'O', algo)
                         @test alpha * opa(A) * opb(B) + beta * C ≈ collect(dC)
+
+                        # add tests for very small matrices (see #2296)
+                        # skip conjugate transpose - causes errors with 1x1 matrices
+                        # CUSPARSE_SPMM_CSR_ALG3 also fails
+                        (algo == CUSPARSE.CUSPARSE_SPMM_CSR_ALG3 || transa == 'C') && continue
+                        A = rand(T, 1, 1)
+                        B = sprand(T, 1, 1, 1.)
+                        C = rand(T, 1, 1)
+                        dA = CuArray(A)
+                        dB = SparseMatrixType(B)
+                        dC = CuArray(C)
+
+                        alpha = rand(T)
+                        beta = rand(T)
+                        mm!(transa, transb, alpha, dA, dB, beta, dC, 'O', algo)
+                        @test alpha * opa(A) * opb(B) + beta * C ≈ collect(dC)
                     end
                 end
             end
