Merge pull request JuliaGPU#3101 from JuliaGPU/tb/coo

CUSPARSE: fix COO scalar getindex crossing row boundaries (JuliaGPU#3100)

Author: maleadt

lib/cusparse/src/array.jl

@@ -462,39 +462,57 @@ Base.getindex(A::CuSparseMatrixCSR, ::Colon, j::Integer) = CuSparseVector(sparse
 
 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]
+    result = zero(Tv)
+    for k in 1:nnz(A)
+        A.iPtr[k] == i && (result = sum_duplicate(result, A.nzVal[k]))
+    end
+    return result
 end
 
+# Scalar getindex methods linear-scan the minor axis rather than binary-searching
+# and sum across matching entries. cuSPARSE formats don't guarantee sorted indices
+# within a major-axis slice (e.g. SpGEMM output may leave CSR columns unsorted
+# within a row, and COO is only guaranteed row-sorted), nor uniqueness — duplicate
+# (i, j) entries are permitted and their values sum, matching the convention of
+# Julia's `sparse()` constructor and SciPy/CuPy. For Bool we OR instead of sum,
+# also matching `sparse()`, since Bool + Bool doesn't stay Bool.
+sum_duplicate(a, b) = a + b
+sum_duplicate(a::Bool, b::Bool) = a | b
+
 function Base.getindex(A::CuSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T
     @boundscheck checkbounds(A, i0, i1)
     r1 = Int(A.colPtr[i1])
     r2 = Int(A.colPtr[i1+1]-1)
-    (r1 > r2) && return zero(T)
-    r1 = searchsortedfirst(rowvals(A), i0, r1, r2, Base.Order.Forward)
-    (r1 > r2 || rowvals(A)[r1] != i0) && return zero(T)
-    nonzeros(A)[r1]
+    result = zero(T)
+    for k in r1:r2
+        rowvals(A)[k] == i0 && (result = sum_duplicate(result, nonzeros(A)[k]))
+    end
+    return result
 end
 
 function Base.getindex(A::CuSparseMatrixCSR{T}, i0::Integer, i1::Integer) where T
     @boundscheck checkbounds(A, i0, i1)
     c1 = Int(A.rowPtr[i0])
     c2 = Int(A.rowPtr[i0+1]-1)
-    (c1 > c2) && return zero(T)
-    c1 = searchsortedfirst(A.colVal, i1, c1, c2, Base.Order.Forward)
-    (c1 > c2 || A.colVal[c1] != i1) && return zero(T)
-    nonzeros(A)[c1]
+    result = zero(T)
+    for k in c1:c2
+        A.colVal[k] == i1 && (result = sum_duplicate(result, nonzeros(A)[k]))
+    end
+    return result
 end
 
 function Base.getindex(A::CuSparseMatrixCOO{T}, i0::Integer, i1::Integer) where T
     @boundscheck checkbounds(A, i0, i1)
+    # cuSPARSE only guarantees COO is sorted by row, so binary-search the row
+    # range but linear-scan for the column.
     r1 = searchsortedfirst(A.rowInd, i0, Base.Order.Forward)
     (r1 > length(A.rowInd) || A.rowInd[r1] > i0) && return zero(T)
-    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 || c1 == length(A.colInd) + 1 || A.colInd[c1] > i1) && return zero(T)
-    nonzeros(A)[c1]
+    r2 = searchsortedlast(A.rowInd, i0, Base.Order.Forward)
+    result = zero(T)
+    for k in r1:r2
+        A.colInd[k] == i1 && (result = sum_duplicate(result, nonzeros(A)[k]))
+    end
+    return result
 end
 
 function Base.getindex(A::CuSparseMatrixBSR{T}, i0::Integer, i1::Integer) where T
@@ -504,10 +522,11 @@ function Base.getindex(A::CuSparseMatrixBSR{T}, i0::Integer, i1::Integer) where
     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]
+    result = zero(T)
+    for k in c1:c2
+        A.colVal[k] == i1_block && (result = sum_duplicate(result, nonzeros(A)[k+block_idx]))
+    end
+    return result
 end
 
 # matrix slices

lib/cusparse/src/conversions.jl

@@ -327,6 +327,41 @@ for SparseMatrixType in (:CuSparseMatrixCSC, :CuSparseMatrixCSR, :CuSparseMatrix
     end
 end
 
+# Wrapper around `cusparseCsr2cscEx2` that works around a cuSPARSE 12.0 bug:
+# invoking the routine with one-based indexing silently corrupts `cscVal` for
+# matrices with certain shapes (e.g. even `m`). On affected versions we shift
+# the index arrays to zero-based around the call; zero-based indexing returns
+# correct results.
+function _csr2cscEx2!(m, n, nnz_, csrVal::CuVector{T}, csrRowPtr, csrColInd,
+                      cscVal::CuVector{T}, cscColPtr, cscRowInd,
+                      action, index, algo) where T
+    buggy = index == 'O' && v"12.0" <= version() < v"12.1"
+    if buggy
+        csrRowPtr = csrRowPtr .- one(Cint)
+        csrColInd = csrColInd .- one(Cint)
+        effidx = 'Z'
+    else
+        effidx = index
+    end
+    function bufferSize()
+        out = Ref{Csize_t}(1)
+        cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz_, csrVal,
+            csrRowPtr, csrColInd, cscVal, cscColPtr, cscRowInd,
+            T, action, effidx, algo, out)
+        return out[]
+    end
+    with_workspace(bufferSize) do buffer
+        cusparseCsr2cscEx2(handle(), m, n, nnz_, csrVal,
+            csrRowPtr, csrColInd, cscVal, cscColPtr, cscRowInd,
+            T, action, effidx, algo, buffer)
+    end
+    if buggy
+        cscColPtr .+= one(Cint)
+        cscRowInd .+= one(Cint)
+    end
+    return
+end
+
 # by flipping rows and columns, we can use that to get CSC to CSR too
 for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
     @eval begin
@@ -335,18 +370,8 @@ for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
             colPtr = CUDACore.zeros(Cint, n+1)
             rowVal = CUDACore.zeros(Cint, nnz(csr))
             nzVal = CUDACore.zeros($elty, nnz(csr))
-            function bufferSize()
-                out = Ref{Csize_t}(1)
-                cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz(csr), nonzeros(csr),
-                    csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
-                    $elty, action, index, algo, out)
-                return out[]
-            end
-            with_workspace(bufferSize) do buffer
-                cusparseCsr2cscEx2(handle(), m, n, nnz(csr), nonzeros(csr),
-                    csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
-                    $elty, action, index, algo, buffer)
-            end
+            _csr2cscEx2!(m, n, nnz(csr), nonzeros(csr), csr.rowPtr, csr.colVal,
+                         nzVal, colPtr, rowVal, action, index, algo)
             CuSparseMatrixCSC(colPtr,rowVal,nzVal,size(csr))
         end
         CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty, Ti}; index::SparseChar='O', action::cusparseAction_t=CUSPARSE_ACTION_NUMERIC, algo::cusparseCsr2CscAlg_t=CUSPARSE_CSR2CSC_ALG1) where {Ti} =
@@ -358,18 +383,8 @@ for elty in (:Float32, :Float64, :ComplexF32, :ComplexF64)
             rowPtr = CUDACore.zeros(Cint,m+1)
             colVal = CUDACore.zeros(Cint,nnz(csc))
             nzVal  = CUDACore.zeros($elty,nnz(csc))
-            function bufferSize()
-                out = Ref{Csize_t}(1)
-                cusparseCsr2cscEx2_bufferSize(handle(), n, m, nnz(csc), nonzeros(csc),
-                    csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
-                    $elty, action, index, algo, out)
-                return out[]
-            end
-            with_workspace(bufferSize) do buffer
-                cusparseCsr2cscEx2(handle(), n, m, nnz(csc), nonzeros(csc),
-                    csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
-                    $elty, action, index, algo, buffer)
-            end
+            _csr2cscEx2!(n, m, nnz(csc), nonzeros(csc), csc.colPtr, rowvals(csc),
+                         nzVal, rowPtr, colVal, action, index, algo)
             CuSparseMatrixCSR(rowPtr,colVal,nzVal,size(csc))
         end
         CuSparseMatrixCSR(csc::CuSparseMatrixCSC{$elty, Ti}; index::SparseChar='O', action::cusparseAction_t=CUSPARSE_ACTION_NUMERIC, algo::cusparseCsr2CscAlg_t=CUSPARSE_CSR2CSC_ALG1) where {Ti} =
@@ -390,18 +405,8 @@ for (elty, welty) in ((:Float16, :Float32),
             rowVal = CUDACore.zeros(Cint, nnz(csr))
             nzVal = CUDACore.zeros($elty, nnz(csr))
             if $elty == Float16 #broken for ComplexF16?
-                function bufferSize()
-                    out = Ref{Csize_t}(1)
-                    cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz(csr), nonzeros(csr),
-                        csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
-                        $elty, action, index, algo, out)
-                    return out[]
-                end
-                with_workspace(bufferSize) do buffer
-                    cusparseCsr2cscEx2(handle(), m, n, nnz(csr), nonzeros(csr),
-                        csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
-                        $elty, action, index, algo, buffer)
-                end
+                _csr2cscEx2!(m, n, nnz(csr), nonzeros(csr), csr.rowPtr, csr.colVal,
+                             nzVal, colPtr, rowVal, action, index, algo)
                 return CuSparseMatrixCSC(colPtr,rowVal,nzVal,size(csr))
             else
                 wide_csr = CuSparseMatrixCSR(csr.rowPtr, csr.colVal, convert(CuVector{$welty}, nonzeros(csr)), size(csr))
@@ -419,18 +424,8 @@ for (elty, welty) in ((:Float16, :Float32),
             colVal = CUDACore.zeros(Cint,nnz(csc))
             nzVal  = CUDACore.zeros($elty,nnz(csc))
             if $elty == Float16 #broken for ComplexF16?
-                function bufferSize()
-                    out = Ref{Csize_t}(1)
-                    cusparseCsr2cscEx2_bufferSize(handle(), n, m, nnz(csc), nonzeros(csc),
-                        csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
-                        $elty, action, index, algo, out)
-                    return out[]
-                end
-                with_workspace(bufferSize) do buffer
-                    cusparseCsr2cscEx2(handle(), n, m, nnz(csc), nonzeros(csc),
-                        csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
-                        $elty, action, index, algo, buffer)
-                end
+                _csr2cscEx2!(n, m, nnz(csc), nonzeros(csc), csc.colPtr, rowvals(csc),
+                             nzVal, rowPtr, colVal, action, index, algo)
                 return CuSparseMatrixCSR(rowPtr,colVal,nzVal,size(csc))
             else
                 wide_csc = CuSparseMatrixCSC(csc.colPtr, csc.rowVal, convert(CuVector{$welty}, nonzeros(csc)), size(csc))

lib/cusparse/test/array.jl

@@ -127,6 +127,111 @@
             @test Array(d_x[i, :]) == x[i, :]
         end
     end
+    # regression test for #3100: scalar getindex at every (i, j)
+    let dense = sparse(reshape(1:16, 4, 4)),
+        d_dense = CuSparseMatrixCOO(dense)
+        CUDACore.@allowscalar begin
+            for j in axes(dense, 2), i in axes(dense, 1)
+                @test d_dense[i, j] == dense[i, j]
+            end
+        end
+    end
+    # sparse case with empty rows and missing entries
+    let s = sparse([1, 1, 3, 4], [1, 3, 2, 4], [10, 20, 30, 40], 4, 4),
+        d_s = CuSparseMatrixCOO(s)
+        CUDACore.@allowscalar begin
+            for j in axes(s, 2), i in axes(s, 1)
+                @test d_s[i, j] == s[i, j]
+            end
+        end
+    end
+    # COO sorted by row but not by column within each row — cuSPARSE's
+    # documented invariant is row-sorted only, so getindex must handle this.
+    let dense = sparse(reshape(1:16, 4, 4)),
+        d_scrambled = CuSparseMatrixCOO(
+            CuArray(Int32[1,1,1,1, 2,2,2,2, 3,3,3,3, 4,4,4,4]),
+            CuArray(Int32[4,2,1,3, 2,4,3,1, 3,1,4,2, 1,4,2,3]),
+            CuArray([13,5,1,9, 6,14,10,2, 11,3,15,7, 4,16,8,12]),
+            (4, 4))
+        CUDACore.@allowscalar begin
+            for j in axes(dense, 2), i in axes(dense, 1)
+                @test d_scrambled[i, j] == dense[i, j]
+            end
+        end
+    end
+    # CSR with unsorted column indices within each row — can arise from
+    # SpGEMM and other operations, so getindex must not binary-search.
+    let dense = sparse(reshape(1:16, 4, 4)),
+        d_scrambled = CuSparseMatrixCSR(
+            CuArray(Int32[1, 5, 9, 13, 17]),
+            CuArray(Int32[4,2,1,3, 2,4,3,1, 3,1,4,2, 1,4,2,3]),
+            CuArray([13,5,1,9, 6,14,10,2, 11,3,15,7, 4,16,8,12]),
+            (4, 4))
+        CUDACore.@allowscalar begin
+            for j in axes(dense, 2), i in axes(dense, 1)
+                @test d_scrambled[i, j] == dense[i, j]
+            end
+        end
+    end
+    # CSC with unsorted row indices within each column
+    let dense = sparse(reshape(1:16, 4, 4)),
+        d_scrambled = CuSparseMatrixCSC(
+            CuArray(Int32[1, 5, 9, 13, 17]),
+            CuArray(Int32[4,2,1,3, 2,4,3,1, 3,1,4,2, 1,4,2,3]),
+            CuArray([4,2,1,3, 6,8,7,5, 11,9,12,10, 13,16,14,15]),
+            (4, 4))
+        CUDACore.@allowscalar begin
+            for j in axes(dense, 2), i in axes(dense, 1)
+                @test d_scrambled[i, j] == dense[i, j]
+            end
+        end
+    end
+    # Duplicate (i, j) entries sum, matching SciPy/CuPy and Julia's sparse().
+    # Each format stores three entries at (1, 1) whose values sum to 3.
+    CUDACore.@allowscalar begin
+        let d_csr = CuSparseMatrixCSR(
+                CuArray(Int32[1, 4, 4]),
+                CuArray(Int32[1, 1, 1]),
+                CuArray([10, -3, -4]),
+                (2, 2))
+            @test d_csr[1, 1] == 3
+            @test d_csr[2, 2] == 0
+            @test d_csr[1, 2] == 0
+        end
+        let d_csc = CuSparseMatrixCSC(
+                CuArray(Int32[1, 4, 4]),
+                CuArray(Int32[1, 1, 1]),
+                CuArray([10, -3, -4]),
+                (2, 2))
+            @test d_csc[1, 1] == 3
+            @test d_csc[2, 2] == 0
+            @test d_csc[2, 1] == 0
+        end
+        let d_coo = CuSparseMatrixCOO(
+                CuArray(Int32[1, 1, 1]),
+                CuArray(Int32[1, 1, 1]),
+                CuArray([10, -3, -4]),
+                (2, 2))
+            @test d_coo[1, 1] == 3
+            @test d_coo[2, 2] == 0
+            @test d_coo[1, 2] == 0
+        end
+        let d_vec = CuSparseVector{Int}(CuArray(Int32[3, 1, 3]), CuArray([10, 7, -4]), 4)
+            @test d_vec[1] == 7
+            @test d_vec[3] == 6
+            @test d_vec[2] == 0
+            @test d_vec[4] == 0
+        end
+        # Bool duplicates combine via OR, not integer sum.
+        let d_bool = CuSparseMatrixCSR(
+                CuArray(Int32[1, 3, 3]),
+                CuArray(Int32[1, 1]),
+                CuArray([true, true]),
+                (2, 2))
+            @test d_bool[1, 1] === true
+            @test d_bool[2, 2] === false
+        end
+    end
     y = sprand(k,n,0.2)
     d_y = CuSparseMatrixCOO(y)
     @test_throws ArgumentError copyto!(d_y,d_x)