Merge pull request #1409 from JuliaGPU/tb/sparse_bsr

maleadt · web-flow · commit 285c7a5b994c · 2022-02-25T11:54:39.000+01:00
CUSPARSE BSR improvements
diff --git a/lib/cusparse/array.jl b/lib/cusparse/array.jl
@@ -102,7 +102,7 @@ mutable struct CuSparseMatrixBSR{Tv, Ti} <: AbstractCuSparseMatrix{Tv, Ti}
     dims::NTuple{2,Int}
     blockDim::Ti
     dir::SparseChar
-    nnz::Ti
+    nnzb::Ti
 
     function CuSparseMatrixBSR{Tv, Ti}(rowPtr::CuVector{<:Integer}, colVal::CuVector{<:Integer},
                                    nzVal::CuVector, dims::NTuple{2,<:Integer},
@@ -268,6 +268,8 @@ LinearAlgebra.istril(M::LowerTriangular{T,S}) where {T<:BlasFloat, S<:AbstractCu
 
 Hermitian{T}(Mat::CuSparseMatrix{T}) where T = Hermitian{T,typeof(Mat)}(Mat,'U')
 
+SparseArrays.nnz(g::CuSparseMatrixBSR) = g.nnzb * g.blockDim * g.blockDim
+
 
 ## indexing
 
@@ -297,6 +299,13 @@ end
 Base.getindex(A::CuSparseMatrixCSC, i::Integer, ::Colon) = CuSparseVector(sparse(A[i, 1:end]))  # TODO: optimize
 Base.getindex(A::CuSparseMatrixCSR, ::Colon, j::Integer) = CuSparseVector(sparse(A[1:end, j]))  # TODO: optimize
 
+function Base.getindex(A::CuSparseVector{Tv, Ti}, i::Integer) where {Tv, Ti}
+    @boundscheck checkbounds(A, i)
+    ii = searchsortedfirst(A.iPtr, convert(Ti, i))
+    (ii > nnz(A) || A.iPtr[ii] != i) && return zero(Tv)
+    A.nzVal[ii]
+end
+
 function Base.getindex(A::CuSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T
     @boundscheck checkbounds(A, i0, i1)
     r1 = Int(A.colPtr[i1])
@@ -327,10 +336,17 @@ function Base.getindex(A::CuSparseMatrixCOO{T}, i0::Integer, i1::Integer) where
     nonzeros(A)[c1]
 end
 
-function SparseArrays._spgetindex(m::Integer, nzind::CuVector{Ti}, nzval::CuVector{Tv},
-                                  i::Integer) where {Tv,Ti}
-    ii = searchsortedfirst(nzind, convert(Ti, i))
-    (ii <= m && nzind[ii] == i) ? nzval[ii] : zero(Tv)
+function Base.getindex(A::CuSparseMatrixBSR{T}, i0::Integer, i1::Integer) where T
+    @boundscheck checkbounds(A, i0, i1)
+    i0_block, i0_idx = fldmod1(i0, A.blockDim)
+    i1_block, i1_idx = fldmod1(i1, A.blockDim)
+    block_idx = (i0_idx - 1) * A.blockDim + i1_idx - 1
+    c1 = Int(A.rowPtr[i0_block])
+    c2 = Int(A.rowPtr[i0_block+1]-1)
+    (c1 > c2) && return zero(T)
+    c1 = searchsortedfirst(A.colVal, i1_block, c1, c2, Base.Order.Forward)
+    (c1 > c2 || A.colVal[c1] != i1_block) && return zero(T)
+    nonzeros(A)[c1+block_idx]
 end
 
 
@@ -407,7 +423,7 @@ function Base.copyto!(dst::CuSparseVector, src::CuSparseVector)
     end
     copyto!(nonzeroinds(dst), nonzeroinds(src))
     copyto!(nonzeros(dst), nonzeros(src))
-    dst.nnz = nnz(src)
+    dst.nnz = src.nnz
     dst
 end
 
@@ -418,7 +434,7 @@ function Base.copyto!(dst::CuSparseMatrixCSC, src::CuSparseMatrixCSC)
     copyto!(dst.colPtr, src.colPtr)
     copyto!(rowvals(dst), rowvals(src))
     copyto!(nonzeros(dst), nonzeros(src))
-    dst.nnz = nnz(src)
+    dst.nnz = src.nnz
     dst
 end
 
@@ -429,7 +445,7 @@ function Base.copyto!(dst::CuSparseMatrixCSR, src::CuSparseMatrixCSR)
     copyto!(dst.rowPtr, src.rowPtr)
     copyto!(dst.colVal, src.colVal)
     copyto!(nonzeros(dst), nonzeros(src))
-    dst.nnz = nnz(src)
+    dst.nnz = src.nnz
     dst
 end
 
@@ -441,7 +457,7 @@ function Base.copyto!(dst::CuSparseMatrixBSR, src::CuSparseMatrixBSR)
     copyto!(dst.colVal, src.colVal)
     copyto!(nonzeros(dst), nonzeros(src))
     dst.dir = src.dir
-    dst.nnz = nnz(src)
+    dst.nnzb = src.nnzb
     dst
 end
 
@@ -452,7 +468,7 @@ function Base.copyto!(dst::CuSparseMatrixCOO, src::CuSparseMatrixCOO)
     copyto!(dst.rowInd, src.rowInd)
     copyto!(dst.colInd, src.colInd)
     copyto!(nonzeros(dst), nonzeros(src))
-    dst.nnz = nnz(src)
+    dst.nnz = src.nnz
     dst
 end
 
@@ -537,7 +553,7 @@ function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseMatrixBSR)
         adapt(to, x.colVal),
         adapt(to, x.nzVal),
         size(x), x.blockDim,
-        x.dir, x.nnz
+        x.dir, x.nnzb
     )
 end
 
diff --git a/lib/cusparse/conversions.jl b/lib/cusparse/conversions.jl
@@ -216,8 +216,8 @@ for (fname,elty) in ((:cusparseScsr2bsr, :Float32),
                                           indc::SparseChar='O')
             m,n = size(csr)
             nnz_ref = Ref{Cint}(1)
-            mb = div((m + blockDim - 1),blockDim)
-            nb = div((n + blockDim - 1),blockDim)
+            mb = cld(m, blockDim)
+            nb = cld(n, blockDim)
             bsrRowPtr = CUDA.zeros(Cint,mb + 1)
             cudesca = CuMatrixDescriptor('G', 'L', 'N', inda)
             cudescc = CuMatrixDescriptor('G', 'L', 'N', indc)
@@ -242,19 +242,20 @@ for (fname,elty) in ((:cusparseSbsr2csr, :Float32),
         function CuSparseMatrixCSR{$elty}(bsr::CuSparseMatrixBSR{$elty};
                                           inda::SparseChar='O', indc::SparseChar='O')
             m,n = size(bsr)
-            mb = div(m,bsr.blockDim)
-            nb = div(n,bsr.blockDim)
-            nnzVal = nnz(bsr) * bsr.blockDim * bsr.blockDim
+            mb = cld(m, bsr.blockDim)
+            nb = cld(n, bsr.blockDim)
             cudesca = CuMatrixDescriptor('G', 'L', 'N', inda)
             cudescc = CuMatrixDescriptor('G', 'L', 'N', indc)
             csrRowPtr = CUDA.zeros(Cint, m + 1)
-            csrColInd = CUDA.zeros(Cint, nnzVal)
-            csrNzVal  = CUDA.zeros($elty, nnzVal)
+            csrColInd = CUDA.zeros(Cint, nnz(bsr))
+            csrNzVal  = CUDA.zeros($elty, nnz(bsr))
             $fname(handle(), bsr.dir, mb, nb,
                    cudesca, nonzeros(bsr), bsr.rowPtr, bsr.colVal,
                    bsr.blockDim, cudescc, csrNzVal, csrRowPtr,
                    csrColInd)
-            CuSparseMatrixCSR(csrRowPtr, csrColInd, csrNzVal, size(bsr))
+            # XXX: the size here may not match the expected size, when the matrix dimension
+            #      is not a multiple of the block dimension!
+            CuSparseMatrixCSR(csrRowPtr, csrColInd, csrNzVal, (mb*bsr.blockDim, nb*bsr.blockDim))
         end
     end
 end
diff --git a/lib/cusparse/level2.jl b/lib/cusparse/level2.jl
@@ -70,7 +70,7 @@ for (bname,aname,sname,elty) in ((:cusparseSbsrsv2_bufferSize, :cusparseSbsrsv2_
             if m != n
                 throw(DimensionMismatch("A must be square, but has dimensions ($m,$n)!"))
             end
-            mb = div(m,A.blockDim)
+            mb = cld(m, A.blockDim)
             mX = length(X)
             if mX != m
                 throw(DimensionMismatch("X must have length $m, but has length $mX"))
@@ -80,21 +80,21 @@ for (bname,aname,sname,elty) in ((:cusparseSbsrsv2_bufferSize, :cusparseSbsrsv2_
 
             function bufferSize()
                 out = Ref{Cint}(1)
-                $bname(handle(), A.dir, transa, mb, nnz(A),
+                $bname(handle(), A.dir, transa, mb, A.nnzb,
                        desc, nonzeros(A), A.rowPtr, A.colVal, A.blockDim,
                        info[1], out)
                 return out[]
             end
             with_workspace(bufferSize) do buffer
-                $aname(handle(), A.dir, transa, mb, nnz(A),
+                $aname(handle(), A.dir, transa, mb, A.nnzb,
                         desc, nonzeros(A), A.rowPtr, A.colVal, A.blockDim,
                         info[1], CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
                 posit = Ref{Cint}(1)
                 cusparseXbsrsv2_zeroPivot(handle(), info[1], posit)
                 if posit[] >= 0
                     error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
                 end
-                $sname(handle(), A.dir, transa, mb, nnz(A),
+                $sname(handle(), A.dir, transa, mb, A.nnzb,
                         alpha, desc, nonzeros(A), A.rowPtr, A.colVal,
                         A.blockDim, info[1], X, X,
                         CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
diff --git a/lib/cusparse/level3.jl b/lib/cusparse/level3.jl
@@ -28,8 +28,8 @@ for (fname,elty) in ((:cusparseSbsrmm, :Float32),
                      index::SparseChar)
             desc = CuMatrixDescriptor('G', 'L', 'N', index)
             m,k = size(A)
-            mb = div(m,A.blockDim)
-            kb = div(k,A.blockDim)
+            mb = cld(m, A.blockDim)
+            kb = cld(k, A.blockDim)
             n = size(C)[2]
             if transa == 'N' && transb == 'N'
                 chkmmdims(B,C,k,n,m,n)
@@ -43,7 +43,7 @@ for (fname,elty) in ((:cusparseSbsrmm, :Float32),
             ldb = max(1,stride(B,2))
             ldc = max(1,stride(C,2))
             $fname(handle(), A.dir,
-                   transa, transb, mb, n, kb, nnz(A),
+                   transa, transb, mb, n, kb, A.nnzb,
                    alpha, desc, nonzeros(A),A.rowPtr, A.colVal,
                    A.blockDim, B, ldb, beta, C, ldc)
             C
@@ -156,7 +156,7 @@ for (bname,aname,sname,elty) in ((:cusparseSbsrsm2_bufferSize, :cusparseSbsrsm2_
             if m != n
                  throw(DimensionMismatch("A must be square, but has dimensions ($m,$n)!"))
             end
-            mb = div(m,A.blockDim)
+            mb = cld(m, A.blockDim)
             mX,nX = size(X)
             if transxy == 'N' && (mX != m)
                 throw(DimensionMismatch(""))
@@ -171,14 +171,14 @@ for (bname,aname,sname,elty) in ((:cusparseSbsrsm2_bufferSize, :cusparseSbsrsm2_
             function bufferSize()
                 out = Ref{Cint}(1)
                 $bname(handle(), A.dir, transa, transxy,
-                       mb, nX, nnz(A), desc, nonzeros(A), A.rowPtr,
+                       mb, nX, A.nnzb, desc, nonzeros(A), A.rowPtr,
                        A.colVal, A.blockDim, info[],
                        out)
                 return out[]
             end
             with_workspace(bufferSize) do buffer
                 $aname(handle(), A.dir, transa, transxy,
-                        mb, nX, nnz(A), desc, nonzeros(A), A.rowPtr,
+                        mb, nX, A.nnzb, desc, nonzeros(A), A.rowPtr,
                         A.colVal, A.blockDim, info[],
                         CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
                 posit = Ref{Cint}(1)
@@ -187,7 +187,7 @@ for (bname,aname,sname,elty) in ((:cusparseSbsrsm2_bufferSize, :cusparseSbsrsm2_
                     error("Structural/numerical zero in A at ($(posit[]),$(posit[])))")
                 end
                 $sname(handle(), A.dir, transa, transxy, mb,
-                        nX, nnz(A), alpha, desc, nonzeros(A), A.rowPtr,
+                        nX, A.nnzb, alpha, desc, nonzeros(A), A.rowPtr,
                         A.colVal, A.blockDim, info[], X, ldx, X, ldx,
                         CUSPARSE_SOLVE_POLICY_USE_LEVEL, buffer)
             end
diff --git a/test/cusparse/device.jl b/test/cusparse/device.jl
@@ -24,7 +24,6 @@ using CUDA.CUSPARSE: CuSparseDeviceVector, CuSparseDeviceMatrixCSC, CuSparseDevi
     cuA = CuSparseMatrixCOO(A)
     @test cudaconvert(cuA) isa CuSparseDeviceMatrixCOO{Float64, Cint, 1}
 
-    # Roger-Luo: I'm not sure how to create a BSR matrix
-    # cuA = CuSparseMatrixBSR(A)
-    # @test cudaconvert(cuA) isa CuSparseDeviceMatrixBSR
+    cuA = CuSparseMatrixBSR(A, 2)
+    @test cudaconvert(cuA) isa CuSparseDeviceMatrixBSR
 end
