Add SparseArrays functionality for CuSparseDeviceColumnView (#2904)

Author: Zinoex

lib/cusparse/device.jl

@@ -23,7 +23,7 @@ Base.length(g::CuSparseDeviceVector) = g.len
 Base.size(g::CuSparseDeviceVector) = (g.len,)
 SparseArrays.nnz(g::CuSparseDeviceVector) = g.nnz
 
-struct CuSparseDeviceMatrixCSC{Tv,Ti,A} <: AbstractSparseMatrix{Tv,Ti}
+struct CuSparseDeviceMatrixCSC{Tv,Ti,A} <: SparseArrays.AbstractSparseMatrixCSC{Tv,Ti}
     colPtr::CuDeviceVector{Ti, A}
     rowVal::CuDeviceVector{Ti, A}
     nzVal::CuDeviceVector{Tv, A}
@@ -38,6 +38,29 @@ SparseArrays.rowvals(g::CuSparseDeviceMatrixCSC) = g.rowVal
 SparseArrays.getcolptr(g::CuSparseDeviceMatrixCSC) = g.colPtr
 SparseArrays.getnzval(g::CuSparseDeviceMatrixCSC) = g.nzVal
 SparseArrays.nzrange(g::CuSparseDeviceMatrixCSC, col::Integer) = SparseArrays.getcolptr(g)[col]:(SparseArrays.getcolptr(g)[col+1]-1)
+SparseArrays.nonzeros(g::CuSparseDeviceMatrixCSC) = g.nzVal
+
+const CuSparseDeviceColumnView{Tv, Ti} = SubArray{Tv, 1, <:CuSparseDeviceMatrixCSC{Tv, Ti}, Tuple{Base.Slice{Base.OneTo{Int}}, Int}}
+function SparseArrays.nonzeros(x::CuSparseDeviceColumnView)
+    rowidx, colidx = parentindices(x)
+    A = parent(x)
+    @inbounds y = view(SparseArrays.nonzeros(A), SparseArrays.nzrange(A, colidx))
+    return y
+end
+
+function SparseArrays.nonzeroinds(x::CuSparseDeviceColumnView)
+    rowidx, colidx = parentindices(x)
+    A = parent(x)
+    @inbounds y = view(SparseArrays.rowvals(A), SparseArrays.nzrange(A, colidx))
+    return y
+end
+SparseArrays.rowvals(x::CuSparseDeviceColumnView) = SparseArrays.nonzeroinds(x)
+
+function SparseArrays.nnz(x::CuSparseDeviceColumnView)
+    rowidx, colidx = parentindices(x)
+    A = parent(x)
+    return length(SparseArrays.nzrange(A, colidx))
+end
 
 struct CuSparseDeviceMatrixCSR{Tv,Ti,A} <: AbstractSparseMatrix{Tv,Ti}
     rowPtr::CuDeviceVector{Ti, A}

test/libraries/cusparse/device.jl

@@ -27,3 +27,82 @@ using CUDA.CUSPARSE: CuSparseDeviceVector, CuSparseDeviceMatrixCSC, CuSparseDevi
     cuA = CuSparseMatrixBSR(A, 2)
     @test cudaconvert(cuA) isa CuSparseDeviceMatrixBSR{Float64, Cint, AS.Global}
 end
+
+@testset "device SparseArrays api" begin
+    @testset "nnz per column" begin
+        function nnz_per_column(A::CuSparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
+            function nnz_per_column_kernel(out, A)
+                i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+                col = @view A[:, i]
+                out[i] = SparseArrays.nnz(col)
+                nothing
+            end
+
+            out = CuVector{Ti}(undef, size(A, 2))
+            @cuda threads=size(A, 2) nnz_per_column_kernel(out, A)
+            out
+        end
+
+        nnz_per_column(A::SparseMatrixCSC) = map(SparseArrays.nnz, eachcol(A))
+
+        A = sprand(10, 10, 0.5)
+        cuA = CuSparseMatrixCSC(A)
+
+        @test nnz_per_column(A) == Vector(nnz_per_column(cuA))
+    end
+
+    @testset "sum per column" begin
+        function sum_per_column(A::CuSparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
+            function sum_per_column_kernel(out, A)
+                j = blockIdx().x
+                col = @view A[:, j]
+                
+                v = zero(Tv)
+                i = threadIdx().x
+                while i <= SparseArrays.nnz(col)
+                    v += nonzeros(col)[i]
+                    i += blockDim().x
+                end
+                v = CUDA.reduce_warp(+, v)
+
+                if threadIdx().x == 1
+                    out[j] = v
+                end
+                nothing
+            end
+
+            out = CuVector{Tv}(undef, size(A, 2))
+            @cuda threads=32 blocks=size(A, 2) sum_per_column_kernel(out, A)
+            out
+        end
+
+        sum_per_column(A::SparseMatrixCSC) = vec(sum(A; dims=1))
+
+        A = sprand(10, 10, 0.5)
+        cuA = CuSparseMatrixCSC(A)
+
+        @test sum_per_column(A) ≈ Vector(sum_per_column(cuA))
+    end
+
+    @testset "last nonzero per column" begin
+        function last_nz_per_column(A::CuSparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
+            function last_nz_per_column_kernel(out, A)
+                i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+                col = @view A[:, i]
+                out[i] = last(SparseArrays.rowvals(col))
+                nothing
+            end
+
+            out = CuVector{Ti}(undef, size(A, 2))
+            @cuda threads=size(A, 2) last_nz_per_column_kernel(out, A)
+            out
+        end
+
+        last_nz_per_column(A::SparseMatrixCSC) = map(last ∘ SparseArrays.rowvals, eachcol(A))
+
+        A = sprand(10, 10, 0.5)
+        cuA = CuSparseMatrixCSC(A)
+
+        @test last_nz_per_column(A) == Vector(last_nz_per_column(cuA))
+    end
+end