Fix scalar indexing issue and add sparse+sparse benchmarks

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

Author: Copilot

benchmarks/matrix_benchmarks.jl

@@ -228,6 +228,66 @@ function benchmark_sparse_dense_add!(
     return nothing
 end
 
+"""
+    benchmark_sparse_sparse_add!(SUITE, array_constructor, array_type_name; N=10000, T=Float64)
+
+Benchmark sparse + sparse matrix addition for CSC, CSR, and COO formats.
+
+# Arguments
+- `SUITE`: The BenchmarkGroup to add benchmarks to
+- `array_constructor`: Function to construct arrays (e.g., `Array`, `JLArray`)
+- `array_type_name`: String name for the array type (for display)
+
+# Keyword Arguments
+- `N`: Size of the matrix (default: 10000)
+- `T`: Element type (default: Float64)
+"""
+function benchmark_sparse_sparse_add!(
+    SUITE,
+    array_constructor,
+    array_type_name;
+    N = 10000,
+    T = Float64,
+)
+    # Create two sparse matrices with 1% density
+    sm_a_csc_std = sprand(T, N, N, 0.01)
+    sm_b_csc_std = sprand(T, N, N, 0.01)
+
+    # Convert to different formats
+    sm_a_csc = DeviceSparseMatrixCSC(sm_a_csc_std)
+    sm_b_csc = DeviceSparseMatrixCSC(sm_b_csc_std)
+    sm_a_csr = DeviceSparseMatrixCSR(sm_a_csc_std)
+    sm_b_csr = DeviceSparseMatrixCSR(sm_b_csc_std)
+    sm_a_coo = DeviceSparseMatrixCOO(sm_a_csc_std)
+    sm_b_coo = DeviceSparseMatrixCOO(sm_b_csc_std)
+
+    # Adapt to device
+    dsm_a_csc = adapt(array_constructor, sm_a_csc)
+    dsm_b_csc = adapt(array_constructor, sm_b_csc)
+    dsm_a_csr = adapt(array_constructor, sm_a_csr)
+    dsm_b_csr = adapt(array_constructor, sm_b_csr)
+    dsm_a_coo = adapt(array_constructor, sm_a_coo)
+    dsm_b_coo = adapt(array_constructor, sm_b_coo)
+
+    # Level 3: Format (CSC, CSR, COO - will be plotted together)
+    SUITE["Sparse + Sparse Addition"][array_type_name]["CSC"] = @benchmarkable begin
+        $dsm_a_csc + $dsm_b_csc
+        _synchronize_backend($dsm_a_csc)
+    end
+
+    SUITE["Sparse + Sparse Addition"][array_type_name]["CSR"] = @benchmarkable begin
+        $dsm_a_csr + $dsm_b_csr
+        _synchronize_backend($dsm_a_csr)
+    end
+
+    SUITE["Sparse + Sparse Addition"][array_type_name]["COO"] = @benchmarkable begin
+        $dsm_a_coo + $dsm_b_coo
+        _synchronize_backend($dsm_a_coo)
+    end
+
+    return nothing
+end
+
 """
     benchmark_kron!(SUITE, array_constructor, array_type_name; N=100, T=Float64)
 

benchmarks/runbenchmarks.jl

@@ -26,6 +26,7 @@ benchmark_matrix_vector_mul!(SUITE, Array, "Array")
 benchmark_matrix_matrix_mul!(SUITE, Array, "Array")
 benchmark_three_arg_dot!(SUITE, Array, "Array")
 benchmark_sparse_dense_add!(SUITE, Array, "Array")
+benchmark_sparse_sparse_add!(SUITE, Array, "Array")
 benchmark_kron!(SUITE, Array, "Array")
 benchmark_conversions!(SUITE, Array, "Array")
 
@@ -37,6 +38,7 @@ benchmark_matrix_vector_mul!(SUITE, JLArray, "JLArray")
 benchmark_matrix_matrix_mul!(SUITE, JLArray, "JLArray")
 benchmark_three_arg_dot!(SUITE, JLArray, "JLArray")
 benchmark_sparse_dense_add!(SUITE, JLArray, "JLArray")
+benchmark_sparse_sparse_add!(SUITE, JLArray, "JLArray")
 benchmark_kron!(SUITE, JLArray, "JLArray")
 benchmark_conversions!(SUITE, JLArray, "JLArray")
 

src/matrix_csc/matrix_csc.jl

@@ -368,11 +368,15 @@ function Base.:+(A::DeviceSparseMatrixCSC, B::DeviceSparseMatrixCSC)
         ndrange = (n,),
     )
 
-    # Build colptr for result matrix
+    # Build colptr for result matrix using cumsum
+    # colptr_C[i+1] = 1 + sum(nnz_per_col[1:i])
+    cumsum_nnz = _cumsum_AK(nnz_per_col)
     colptr_C = similar(getcolptr(A), n + 1)
-    colptr_C[1] = one(Ti)
-    colptr_C[2:end] .= _cumsum_AK(nnz_per_col)
+    # Set colptr_C[2:end] to cumsum + 1
+    colptr_C[2:end] .= cumsum_nnz
     colptr_C[2:end] .+= one(Ti)
+    # Set colptr_C[1] to 1 using broadcasting
+    colptr_C[1:1] .= one(Ti)
 
     # Allocate result arrays
     nnz_total = allowed_getindex(colptr_C, n + 1) - one(Ti)

src/matrix_csr/matrix_csr.jl

@@ -365,11 +365,15 @@ function Base.:+(A::DeviceSparseMatrixCSR, B::DeviceSparseMatrixCSR)
         ndrange = (m,),
     )
 
-    # Build rowptr for result matrix
+    # Build rowptr for result matrix using cumsum
+    # rowptr_C[i+1] = 1 + sum(nnz_per_row[1:i])
+    cumsum_nnz = _cumsum_AK(nnz_per_row)
     rowptr_C = similar(getrowptr(A), m + 1)
-    rowptr_C[1] = one(Ti)
-    rowptr_C[2:end] .= _cumsum_AK(nnz_per_row)
+    # Set rowptr_C[2:end] to cumsum + 1
+    rowptr_C[2:end] .= cumsum_nnz
     rowptr_C[2:end] .+= one(Ti)
+    # Set rowptr_C[1] to 1 using broadcasting
+    rowptr_C[1:1] .= one(Ti)
 
     # Allocate result arrays
     nnz_total = allowed_getindex(rowptr_C, m + 1) - one(Ti)

test/Project.toml

@@ -1,6 +1,7 @@
 [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"