Fixes and tests for COO indexing, exclude more kernels from coverage (#2668)

Author: kshyatt

lib/cusparse/array.jl

@@ -385,6 +385,34 @@ function Base.getindex(x::CuSparseMatrixCSR, i::Integer, ::Colon)
     CuSparseVector(x.colVal[c1:c2], nonzeros(x)[c1:c2], size(x, 2))
 end
 
+function Base.getindex(x::CuSparseMatrixCOO{T}, i::Integer, ::Colon) where {T}
+    checkbounds(x, i, :)
+    if issorted(x.rowInd)
+        row_start = searchsortedfirst(x.rowInd, i)
+        row_end   = min(searchsortedlast(x.rowInd, i), length(x.rowInd))
+        row_start == length(x.rowInd) + 1 && return CuSparseVector(similar(x.rowInd, 0), CUDA.zeros(T, 0), size(x, 2))
+        CuSparseVector(x.colInd[row_start:row_end], x.nzVal[row_start:row_end], size(x, 2))
+    else
+        row_inds = findall(ix->ix == i, x.rowInd)
+        isnothing(row_inds) && return CuSparseVector(similar(x.rowInd, 0), CUDA.zeros(T, 0), size(x, 2))
+        CuSparseVector(x.colInd[row_inds], x.nzVal[row_inds], size(x, 2))
+    end
+end
+
+function Base.getindex(x::CuSparseMatrixCOO{T}, ::Colon, j::Integer) where {T}
+    checkbounds(x, :, j)
+    if issorted(x.colInd)
+        col_start = searchsortedfirst(x.colInd, j)
+        col_end   = min(searchsortedlast(x.colInd, j), length(x.colInd))
+        col_start == length(x.colInd) + 1 && return CuSparseVector(similar(x.colInd, 0), CUDA.zeros(T, 0), size(x, 2))
+        CuSparseVector(x.rowInd[col_start:col_end], x.nzVal[col_start:col_end], size(x, 1))
+    else
+        col_inds = findall(ix->ix == j, x.colInd)
+        isnothing(col_inds) && return CuSparseVector(similar(x.colInd, 0), CUDA.zeros(T, 0), size(x, 1))
+        CuSparseVector(x.rowInd[col_inds], x.nzVal[col_inds], size(x, 1))
+    end
+end
+
 # row slices
 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
@@ -420,9 +448,9 @@ function Base.getindex(A::CuSparseMatrixCOO{T}, i0::Integer, i1::Integer) where
     @boundscheck checkbounds(A, i0, i1)
     r1 = searchsortedfirst(A.rowInd, i0, Base.Order.Forward)
     (r1 > length(A.rowInd) || A.rowInd[r1] > i0) && return zero(T)
-    r2 = searchsortedfirst(A.rowInd, i0+1, Base.Order.Forward)
+    r2 = min(searchsortedfirst(A.rowInd, i0+1, Base.Order.Forward), length(A.rowInd))
     c1 = searchsortedfirst(A.colInd, i1, r1, r2, Base.Order.Forward)
-    (c1 > r2 || A.colInd[c1] > i1) && return zero(T)
+    (c1 > r2 || c1 == length(A.colInd) + 1 || A.colInd[c1] > i1) && return zero(T)
     nonzeros(A)[c1]
 end
 

lib/cusparse/broadcast.jl

@@ -101,6 +101,7 @@ end
 @inline _zeros_eltypes(A, Bs...) = (zero(eltype(A)), _zeros_eltypes(Bs...)...)
 
 
+## COV_EXCL_START
 ## iteration helpers
 
 """
@@ -300,7 +301,6 @@ _getindex(arg, I, ptr) = Broadcast._broadcast_getindex(arg, I)
 ## sparse broadcast implementation
 
 # TODO: unify CSC/CSR kernels
-## COV_EXCL_START
 # kernel to count the number of non-zeros in a row, to determine the row offsets
 function compute_offsets_kernel(::Type{<:CuSparseMatrixCSR}, offsets::AbstractVector{Ti},
                                 args...) where Ti

lib/cusparse/reduce.jl

@@ -50,6 +50,7 @@ function Base.mapreduce(f, op, A::Union{CuSparseMatrixCSR,CuSparseMatrixCSC};
     end
 end
 
+## COV_EXCL_START
 function csr_reduce_kernel(f::F, op::OP, neutral, output::CuDeviceArray, args...) where {F, OP}
     # every thread processes an entire row
     row = threadIdx().x + (blockIdx().x - 1i32) * blockDim().x
@@ -95,3 +96,4 @@ function csc_reduce_kernel(f::F, op::OP, neutral, output::CuDeviceArray, args...
     @inbounds output[col] = val
     return
 end
+## COV_EXCL_STOP

test/libraries/cusparse.jl

@@ -75,6 +75,38 @@ blockdim = 5
     y = sprand(k,n,0.2)
     d_y = CuSparseMatrixCSC(y)
     @test_throws ArgumentError copyto!(d_y,d_x)
+    x = sprand(m,n,0.2)
+    d_x = CuSparseMatrixCOO(x)
+    @test CuSparseMatrixCOO(d_x) === d_x
+    @test length(d_x) == m*n
+    @test size(d_x)   == (m,n)
+    @test size(d_x,1) == m
+    @test size(d_x,2) == n
+    @test size(d_x,3) == 1
+    @test ndims(d_x)  == 2
+    d_x2 = copy(d_x)
+    @test d_x2 isa CuSparseMatrixCOO
+    @test size(d_x2) == size(d_x)
+    @test length(d_x) == length(x)
+    CUDA.@allowscalar begin
+        @test Array(d_x[:])        == x[:]
+        @test d_x[firstindex(d_x)] == x[firstindex(x)]
+        @test d_x[div(end, 2)]     == x[div(end, 2)]
+        @test d_x[end]             == x[end]
+        @test d_x[firstindex(d_x), firstindex(d_x)] == x[firstindex(x), firstindex(x)]
+        @test d_x[div(end, 2), div(end, 2)]         == x[div(end, 2), div(end, 2)]
+        @test d_x[end, end]        == x[end, end]
+        @test Array(d_x[firstindex(d_x):end]) == x[:]
+        for i in 1:size(x, 2)
+            @test Array(d_x[:, i]) == x[:, i]
+        end
+        for i in 1:size(x, 1)
+            @test Array(d_x[i, :]) == x[i, :]
+        end
+    end
+    y = sprand(k,n,0.2)
+    d_y = CuSparseMatrixCOO(y)
+    @test_throws ArgumentError copyto!(d_y,d_x)
     d_y = CuSparseMatrixCSR(d_y)
     d_x = CuSparseMatrixCSR(d_x)
     @test CuSparseMatrixCSR(d_x) === d_x
@@ -151,6 +183,7 @@ end
             @test size(similar(d_x, (3, 4))) == (3, 4)
             @test similar(d_x, Float32) isa CuSparseMatrixCSC{Float32}
             @test similar(d_x, Float32, n, m) isa CuSparseMatrixCSC{Float32}
+            @test similar(d_x, Float32, (n, m)) isa CuSparseMatrixCSC{Float32}
         end
 
         @testset "CSR" begin
@@ -162,6 +195,7 @@ end
             @test size(similar(d_x, (3, 4))) == (3, 4)
             @test similar(d_x, Float32) isa CuSparseMatrixCSR{Float32}
             @test similar(d_x, Float32, n, m) isa CuSparseMatrixCSR{Float32}
+            @test similar(d_x, Float32, (n, m)) isa CuSparseMatrixCSR{Float32}
         end
 
         @testset "COO" begin
@@ -171,7 +205,9 @@ end
             @test similar(d_x) isa CuSparseMatrixCOO{elty}
             @test similar(d_x, (3, 4)) isa CuSparseMatrixCOO{elty}
             @test size(similar(d_x, (3, 4))) == (3, 4)
+            @test size(similar(d_x, Float64, (3, 4))) == (3, 4)
             @test similar(d_x, Float32) isa CuSparseMatrixCOO{Float32}
+            @test CuSparseMatrixCOO(d_x) === d_x
         end
 
         @testset "BSR" begin