Fix COO addition to merge duplicates and remove unnecessary test dependency

Co-authored-by: albertomercurio <61953577+albertomercurio@users.noreply.github.com>

Author: Copilot

src/matrix_coo/matrix_coo.jl

@@ -334,8 +334,8 @@ end
 Add two sparse matrices in COO format. Both matrices must have the same dimensions
 and be on the same backend (device).
 
-The result is a COO matrix with entries from both A and B concatenated. Note that
-duplicate entries (same row and column) are not combined, which is valid for COO format.
+The result is a COO matrix with entries from both A and B properly merged,
+with duplicate entries (same row and column) combined by summing their values.
 
 # Examples
 ```jldoctest
@@ -366,27 +366,68 @@ function Base.:+(A::DeviceSparseMatrixCOO, B::DeviceSparseMatrixCOO)
         throw(ArgumentError("Both matrices must have the same backend"))
 
     m, n = size(A)
+    Ti = eltype(getrowind(A))
     Tv = promote_type(eltype(nonzeros(A)), eltype(nonzeros(B)))
 
     # Concatenate the coordinate arrays
     nnz_A = nnz(A)
     nnz_B = nnz(B)
-    nnz_total = nnz_A + nnz_B
-
-    # Allocate result arrays
-    rowind_C = similar(getrowind(A), nnz_total)
-    colind_C = similar(getcolind(A), nnz_total)
-    nzval_C = similar(nonzeros(A), Tv, nnz_total)
-
-    # Copy entries from A
-    rowind_C[1:nnz_A] .= getrowind(A)
-    colind_C[1:nnz_A] .= getcolind(A)
-    nzval_C[1:nnz_A] .= nonzeros(A)
-
-    # Copy entries from B
-    rowind_C[(nnz_A+1):end] .= getrowind(B)
-    colind_C[(nnz_A+1):end] .= getcolind(B)
-    nzval_C[(nnz_A+1):end] .= nonzeros(B)
+    nnz_concat = nnz_A + nnz_B
+
+    # Allocate concatenated arrays
+    rowind_concat = similar(getrowind(A), nnz_concat)
+    colind_concat = similar(getcolind(A), nnz_concat)
+    nzval_concat = similar(nonzeros(A), Tv, nnz_concat)
+
+    # Copy entries from A and B
+    rowind_concat[1:nnz_A] .= getrowind(A)
+    colind_concat[1:nnz_A] .= getcolind(A)
+    nzval_concat[1:nnz_A] .= nonzeros(A)
+    rowind_concat[(nnz_A+1):end] .= getrowind(B)
+    colind_concat[(nnz_A+1):end] .= getcolind(B)
+    nzval_concat[(nnz_A+1):end] .= nonzeros(B)
+
+    # Sort by (row, col) using keys similar to COO->CSC conversion
+    backend = backend_A
+    keys = similar(rowind_concat, Ti, nnz_concat)
+    kernel_make_keys! = kernel_make_csc_keys!(backend)
+    kernel_make_keys!(keys, rowind_concat, colind_concat, m; ndrange = (nnz_concat,))
+
+    # Sort using AcceleratedKernels
+    perm = _sortperm_AK(keys)
+
+    # Apply permutation to get sorted arrays
+    rowind_sorted = rowind_concat[perm]
+    colind_sorted = colind_concat[perm]
+    nzval_sorted = nzval_concat[perm]
+
+    # Mark unique entries (first occurrence of each (row, col) pair)
+    keep_mask = similar(rowind_sorted, Bool, nnz_concat)
+    kernel_mark! = kernel_mark_unique_coo!(backend)
+    kernel_mark!(keep_mask, rowind_sorted, colind_sorted, nnz_concat; ndrange = (nnz_concat,))
+
+    # Compute write indices using cumsum
+    write_indices = _cumsum_AK(keep_mask)
+    nnz_final = allowed_getindex(write_indices, nnz_concat)
+
+    # Allocate final arrays
+    rowind_C = similar(getrowind(A), nnz_final)
+    colind_C = similar(getcolind(A), nnz_final)
+    nzval_C = similar(nonzeros(A), Tv, nnz_final)
+
+    # Compact: merge duplicates by summing
+    kernel_compact! = kernel_compact_coo!(backend)
+    kernel_compact!(
+        rowind_C,
+        colind_C,
+        nzval_C,
+        rowind_sorted,
+        colind_sorted,
+        nzval_sorted,
+        write_indices,
+        nnz_concat;
+        ndrange = (nnz_concat,),
+    )
 
     return DeviceSparseMatrixCOO(m, n, rowind_C, colind_C, nzval_C)
 end

src/matrix_coo/matrix_coo_kernels.jl

@@ -181,3 +181,55 @@ end
         nzval_C[idx] = val_A * val_B
     end
 end
+
+# Kernel for marking duplicate entries in sorted COO format
+# Returns a mask where mask[i] = true if entry i should be kept (first occurrence or sum)
+@kernel inbounds=true function kernel_mark_unique_coo!(
+    keep_mask,
+    @Const(rowind),
+    @Const(colind),
+    @Const(nnz_total),
+)
+    i = @index(Global)
+
+    if i == 1
+        # Always keep the first entry
+        keep_mask[i] = true
+    elseif i <= nnz_total
+        # Keep if different from previous entry
+        keep_mask[i] = (rowind[i] != rowind[i-1] || colind[i] != colind[i-1])
+    end
+end
+
+# Kernel for compacting COO by summing duplicate entries
+@kernel inbounds=true function kernel_compact_coo!(
+    rowind_out,
+    colind_out,
+    nzval_out,
+    @Const(rowind_in),
+    @Const(colind_in),
+    @Const(nzval_in),
+    @Const(write_indices),
+    @Const(nnz_in),
+)
+    i = @index(Global)
+
+    if i <= nnz_in
+        out_idx = write_indices[i]
+        
+        # If this is a new entry (or first of duplicates), write it
+        if i == 1 || (rowind_in[i] != rowind_in[i-1] || colind_in[i] != colind_in[i-1])
+            rowind_out[out_idx] = rowind_in[i]
+            colind_out[out_idx] = colind_in[i]
+            
+            # Sum all duplicates
+            val_sum = nzval_in[i]
+            j = i + 1
+            while j <= nnz_in && rowind_in[j] == rowind_in[i] && colind_in[j] == colind_in[i]
+                val_sum += nzval_in[j]
+                j += 1
+            end
+            nzval_out[out_idx] = val_sum
+        end
+    end
+end

test/Project.toml

@@ -1,7 +1,6 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-DeviceSparseArrays = "da3fe0eb-88a8-4d14-ae1a-857c283e9c70"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"