Fix Enzyme sparse matrix sparsity pattern corruption (issue #835)

[ChrisRackauckas-Claude]

Summary

• Fixes the issue where Enzyme AD with sparse matrices corrupts the primal matrix's sparsity pattern (rowval, colptr)
• Adds sparse-safe helper functions that operate directly on nzval arrays instead of using broadcast operations that can change sparsity
• For dense matrices, falls back to standard broadcast operations

Root Cause

Enzyme.make_zero creates shadow sparse matrices that share the structural arrays (rowval, colptr) with the primal matrix. When the reverse rule executes broadcast operations like dA .-= z * transpose(y), these can change the sparsity pattern of the shadow matrix, and because the structural arrays are shared, this inadvertently corrupts the primal matrix's structure.

Evidence from investigation:

Primal rowval === shadow rowval? true
Primal colptr === shadow colptr? true

Solution

Add sparse-safe helper functions that dispatch on AbstractSparseMatrix:

1. _safe_add!(dst, src): For sparse matrices, adds values via dst.nzval .+= src.nzval
2. _safe_zero!(A): For sparse matrices, zeros via fill!(A.nzval, 0)
3. _sparse_outer_sub!(dA, z, y): For sparse matrices, only accumulates gradients into existing non-zero positions using a direct loop over the CSC structure

These preserve the sparsity pattern by operating only on the nzval array.

Test plan

• Verify sparsity pattern is preserved after AD (rowval, colptr, nnz unchanged)
• Verify matrix can still be displayed without AssertionError: _goodbuffers(S) crash
• Verify dense matrix Enzyme tests still pass (note: existing tests have pre-existing failures unrelated to this change)

Fixes #835

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

Updated based on feedback

Addressed two concerns:

1. GPU Safety

Changed the implementation from nested loops with scalar indexing to vectorized operations:

# Build column indices vector from colptr
col_indices = _expand_colptr_to_col_indices(colptr, n_cols, nnz_count)

# Vectorized update instead of nested loop
vals .-= z[rows] .* y[col_indices]

The _expand_colptr_to_col_indices helper still uses a loop internally (dispatched on Vector{Ti}), but this is CPU-only since SparseMatrixCSC itself is CPU-only. GPU sparse matrices (like CuSparseMatrixCSC) are different types and would need handling in their respective extensions (LinearSolveCUDAExt, etc.).

2. Non-zeros concern

The dense fallback dA .-= z * transpose(y) only applies to truly dense AbstractArray types. Sparse matrices dispatch to the SparseMatrixCSC method which only updates existing positions in the sparsity pattern. This is mathematically correct for sparse AD - gradients are only meaningful at positions where the matrix can be modified.

[ChrisRackauckas-Claude]

Updated implementation

Based on feedback, the implementation has been revised to use a non-allocating loop for sparse matrix operations:

function _sparse_outer_sub!(dA::SparseMatrixCSC, z::AbstractVector, y::AbstractVector)
    rows = rowvals(dA)
    vals = nonzeros(dA)
    colptr = getcolptr(dA)

    # Non-allocating loop over CSC structure
    # This is efficient and cache-friendly (column-major order)
    @inbounds for col in 1:size(dA, 2)
        y_col = y[col]
        for idx in colptr[col]:(colptr[col + 1] - 1)
            vals[idx] -= z[rows[idx]] * y_col
        end
    end

    return dA
end

Key points:

• Zero allocations: The loop operates directly on the existing arrays
• Cache-friendly: Column-major order iteration matches CSC storage
• CPU-only by design: SparseMatrixCSC is a CPU-only type; GPU sparse matrices (like CuSparseMatrixCSC) have their own types and would need handling in their respective extensions (CUDA, AMDGPU, etc.)

The SparseArrays package has been added as a dependency for the Enzyme extension.

Note: The dense matrix Enzyme tests are failing on main as a pre-existing issue (unrelated to this PR).

[ChrisRackauckas-Claude]

Updated: Added GPU sparse matrix support

Now supports both CPU and GPU sparse matrices:

CPU sparse (SparseMatrixCSC)

Uses non-allocating loop:

function _sparse_outer_sub!(dA::SparseMatrixCSC, z::AbstractVector, y::AbstractVector)
    rows = rowvals(dA)
    vals = nonzeros(dA)
    colptr = getcolptr(dA)

    @inbounds for col in 1:size(dA, 2)
        y_col = y[col]
        for idx in colptr[col]:(colptr[col + 1] - 1)
            vals[idx] -= z[rows[idx]] * y_col
        end
    end
    return dA
end

GPU sparse (AbstractSparseMatrixCSC - CuSparseMatrixCSC, ROCSparseMatrixCSC, etc.)

Uses vectorized operations with O(nnz) allocation:

function _sparse_outer_sub!(dA::AbstractSparseMatrixCSC, z::AbstractVector, y::AbstractVector)
    rows = rowvals(dA)
    vals = nonzeros(dA)
    colptr = getcolptr(dA)
    n_cols = size(dA, 2)

    # Build column indices (copies colptr to CPU, builds indices, copies back to GPU)
    col_indices = _expand_colptr_to_col_indices(rows, colptr, n_cols)

    # Vectorized update - works on GPU via broadcasting
    vals .-= z[rows] .* y[col_indices]
    return dA
end

The GPU path:

1. Copies colptr to CPU (one GPU→CPU transfer)
2. Builds column indices on CPU (fast)
3. Copies result to GPU (one CPU→GPU transfer)
4. Uses vectorized broadcasting (GPU-native)

This is 2 memory transfers total, which is acceptable for gradient computation.