Fixes for COO indices. It works now

Author: alonsoC1s

src/device/indexing.jl

@@ -7,7 +7,10 @@ Base.IndexStyle(::Type{GPUSparseDeviceVector}) = Base.IndexLinear()
 # Scalar indexing
 ## Adapted from SparseArrays.AbstractSparseVector
 
-@propagate_inbounds function Base.getindex(v::GPUSparseDeviceVector{Tv,Ti}, i::Integer) where {Tv,Ti}
+@propagate_inbounds function Base.getindex(
+    v::GPUSparseDeviceVector{Tv,Ti},
+    i::Integer,
+) where {Tv,Ti}
     @boundscheck checkbounds(v, i)
     m = nnz(v)
     nzind = nonzeroinds(v)
@@ -20,10 +23,17 @@ end
 # TODO: Logical indexing
 
 # Indexing by colon not implemented. Non-scalar indexing would allocate in device code
-@propagate_inbounds Base.getindex(A::AbstractGPUSparseDeviceMatrix, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2])
+@propagate_inbounds Base.getindex(
+    A::AbstractGPUSparseDeviceMatrix,
+    I::Tuple{Integer,Integer},
+) = getindex(A, I[1], I[2])
 
 ## Adapted logic from SparseArrays.AbstractSparseMatrixCSC
-@propagate_inbounds function Base.getindex(A::GPUSparseDeviceMatrixCSC{Tv,Ti}, i::Integer, j::Integer) where {Tv,Ti}
+@propagate_inbounds function Base.getindex(
+    A::GPUSparseDeviceMatrixCSC{Tv,Ti},
+    i::Integer,
+    j::Integer,
+) where {Tv,Ti}
     @boundscheck checkbounds(A, i, j)
     colPtr, rowVal, nzVal = getcolptr(A), rowvals(A), nonzeros(A)
 
@@ -32,12 +42,15 @@ end
     rr = convert(Ti, @inbounds colPtr[j+1] - 1)
     (rl > rr) && return zero(Tv)
 
-    # possible_row = @view rowVal[rl:rr]
     ii = searchsortedfirst(rowVal, convert(Ti, i), rl, rr, Base.Order.Forward)
     (ii <= nnz(A) && rowVal[ii] == i) ? nzVal[ii] : zero(Tv)
 end
 
-@propagate_inbounds function Base.getindex(A::GPUSparseDeviceMatrixCSR{Tv,Ti}, i::Integer, j::Integer) where {Tv,Ti}
+@propagate_inbounds function Base.getindex(
+    A::GPUSparseDeviceMatrixCSR{Tv,Ti},
+    i::Integer,
+    j::Integer,
+) where {Tv,Ti}
     @boundscheck checkbounds(A, i, j)
     rowPtr, colVal, nzVal = A.rowPtr, A.colVal, A.nzVal
 
@@ -46,27 +59,28 @@ end
     rb = convert(Ti, @inbounds rowPtr[i+1] - 1)
     (rt > rb) && return zero(Tv)
 
-    # possible_col = @view colVal[rt:rb]
     jj = searchsortedfirst(colVal, convert(Ti, j), rt, rb, Base.Order.Forward)
     (jj <= nnz(A) && colVal[jj] == j) ? nzVal[jj] : zero(Tv)
 end
 
 ## Adapted from CUDA.jl/blob/lib/cusparse/src/array.jl#L490
-# FIXME: Currently not correct
-@propagate_inbounds function Base.getindex(A::GPUSparseDeviceMatrixCOO{Tv,Ti}, i::Integer, j::Integer) where {Tv,Ti}
+@propagate_inbounds function Base.getindex(
+    A::GPUSparseDeviceMatrixCOO{Tv,Ti},
+    i::Integer,
+    j::Integer,
+) where {Tv,Ti}
     # COO in CUDA is assumed to be sorted by row: https://docs.nvidia.com/cuda/cusparse/storage-formats.html?highlight=coo#coordinate-coo
-    # @boundscheck checkbounds(A, i, j)
+    @boundscheck checkbounds(A, i, j)
     rowInd, colInd, nzVal = A.rowInd, A.colInd, A.nzVal
 
     # Looking for the range s.t. rowInd[r1:r2] .== i
     rl = searchsortedfirst(rowInd, i)
     (rl > nnz(A) || rowInd[rl] > i) && return 42
     rr = min(searchsortedfirst(rowInd, i+1, Base.Order.Forward), nnz(A)) # searchsortedlast didn't behave as expected
-    # FIXME: colInd isn't sorted
     jj = searchsortedfirst(colInd, j, rl, rr, Base.Order.Forward)
     (jj > rr || jj == nnz(A) + 1 || colInd[jj] > j) && return zero(Tv)
 
-    return jj
+    return nzVal[jj]
 end
 
 # TODO: Support BSR format