Crate rlx_sparse 

Sparse linear algebra for RLX — CSR LU, mat-vec, Conjugate Gradient.

Downstream package modeled on jax.experimental.sparse. Registers against rlx’s custom-op scaffold without requiring any edits to the framework crates. Three ops + a SparseTensor boundary abstraction.

Usage

// At application startup, once.
rlx_sparse::register();

// Build graph as usual.
let mut g = Graph::new("photonics");
let v  = g.input("values",  Shape::new(&[nnz], DType::F64));
let ci = ...; // I32 col_idx (Op::Constant or Op::Input)
let rp = ...; // I32 row_ptr
let b  = g.input("b", Shape::new(&[n], DType::F64));

let a = rlx_sparse::SparseTensor::from_csr(v, ci, rp, n, n);
let x = a.solve(&mut g, b);                 // direct LU
let y = a.mat_vec(&mut g, x);               // sparse matvec
let z = a.cg_solve(&mut g, b, 200, 1e-12);  // iterative CG

What’s registered

• rlx_sparse.lu_solve — direct LU via host LAPACK dgesv. v1 densifies CSR before solving; performance not yet sparse-fast, semantics are correct. Swapping for SuiteSparse UMFPACK or KLU is a kernel-body change with zero IR diff.
• rlx_sparse.mat_vec — y = A·x over CSR.
• rlx_sparse.cg_solve — Conjugate Gradient for SPD systems with max_iter + tol baked into the op’s attrs blob.

Adjoint convention (v1)

All three ops assume A is symmetric. The closed-form adjoint dL/db = solve(Aᵀ, dL/dx) reuses the same CSR triplet as the forward call. Non-symmetric A requires an explicit transpose triplet — sketch in the vjp body of each op. dL/dvalues is non-differentiable in v1; it’s gather(-(dL/db) ⊗ x) and slots in as a separate gather op.

Backend support

Backend	Status
CPU	Full forward + autodiff. Real LAPACK.
Metal	Trait surface only — full executor dispatch is a follow-up.
MLX	Trait surface only — full executor dispatch is a follow-up.
Others	Op::Custom rejected at legalize; pin graph to Device::Cpu.

Structs

SparseTensor

CSR-format sparse matrix at the IR level. Bundles the three CSR NodeIds with structural shape info known at graph-build time.

Constants

SPARSE_BICGSTAB_SOLVE

BiCGSTAB iterative solver for general (non-symmetric) sparse A·x = b. 4 inputs (values, col_idx, row_ptr, b) + attrs encoding (max_iter: u32, tol: f64, transpose_a: u8). When transpose_a is set, the kernel solves Aᵀ·x = b — used by VJPs for adjoint solves.

SPARSE_CG_SOLVE

SPARSE_CHOLESKY_SOLVE

Direct sparse Cholesky for SPD A·x = b. Densifies A and uses LAPACK dpotrf + triangular solves. Same I/O contract as SPARSE_LU_SOLVE but only valid for SPD matrices — ½× factor cost of LU and numerically more stable.

SPARSE_GMRES_SOLVE

GMRES solve for non-symmetric systems. Iterative analog of CG for the asymmetric Maxwell / advection-diffusion regime. Same 7-input shape as SPARSE_LU_SOLVE_GENERAL: forward uses A, VJP routes the adjoint through Aᵀ.

SPARSE_ILU_PCG_SOLVE

ILU(0)-preconditioned CG. Factors A in-place over its existing sparsity pattern (zero fill-in) and applies the LU triangular solves as the preconditioner. Same 4-input shape as PCG; converges faster than Jacobi-PCG on stiff systems where row-row coupling dominates the off-diagonal.

SPARSE_LSQR_SOLVE

LSQR for sparse least-squares min_x ||A·x - b||₂. 4 inputs (values, col_idx, row_ptr, b) + attrs encoding (max_iter, tol, n_cols). Forward only in v1 — VJP returns empty (least-squares adjoint requires either AᵀA solve or a recursive LSQR call which is non-trivial; defer until a use case appears).

SPARSE_LU_SOLVE

SPARSE_LU_SOLVE_GENERAL

Non-symmetric LU solve. Forward x = A⁻¹·b (uses A only). VJP dL/db = solve(Aᵀ, dL/dx) — needs an explicit transpose triplet, supplied as the last 3 inputs to keep the IR self- contained. The 4-input SPARSE_LU_SOLVE is the symmetric specialization.

SPARSE_MAT_VEC

SPARSE_PCG_SOLVE

Jacobi-preconditioned CG. Same 4-input shape as SPARSE_CG_SOLVE (values, col_idx, row_ptr, b). The kernel extracts diag(A) internally and uses it as the diagonal preconditioner — dramatically faster than plain CG on ill-conditioned circuit matrices where the diagonal magnitudes vary by orders of magnitude. Convergence requires SPD A like CG.

SPARSE_SPGEMM

Sparse-sparse matrix multiply (CSR × CSR → CSR). 6 inputs: (a_values, a_col_idx, a_row_ptr, b_values, b_col_idx, b_row_ptr) plus attrs encoding (k: u32 = inner dim = b’s row count). Output is a packed buffer [c_values | c_col_idx_as_f64 | c_row_ptr_as_f64] with sizes encoded in attrs alongside k. v1 caps nnz output at max_nnz (attrs); allocate generously for known patterns.

SPARSE_TRANSPOSE_VALUES

Permute a CSR values vector into the values vector of Aᵀ. 5 inputs: (values_A, col_idx_A, row_ptr_A, col_idx_AT, row_ptr_AT). The transposed pattern (col_idx_AT, row_ptr_AT) is structural — depends only on the original pattern — so it can be precomputed once at graph-build time via csr_transpose_pattern and embedded as constants. Only the values get permuted per call. Useful for inverse-design Newton loops where the matrix entries change each iteration but the sparsity pattern is fixed.

SPARSE_VALUES_GRAD

Outer-product gather op (the dL/dvalues building block). Computes out[k] = u[row_of(k)] * v[col_idx[k]] for each non-zero position k in the CSR pattern. Used by SparseLu/SparseMatVec/ SparseCg/SparseGmres VJPs to gather the dense outer-product u ⊗ v at the matrix’s nonzero positions.

Functions

cg_solve

Register every sparse op’s IR-level extension and per-backend kernels enabled at compile time. Idempotent — the underlying registries already warn on overwrite. Call once at application startup. Host CG for SPD A·x = b given CSR (values, col_idx, row_ptr).

csr_transpose_pattern

Compute (col_idx_T, row_ptr_T) — the sparsity pattern of Aᵀ — from A’s pattern. This is the structural step that must happen before SPARSE_TRANSPOSE_VALUES can permute the values per Newton iteration. Result is independent of the values, so downstream callers compute it once and embed as Op::Constant.

encode_cg_attrs

Encode CG attrs into the opaque Vec<u8> blob carried on Op::Custom. Layout: [max_iter:u32 LE, tol:f64 LE] — 12 bytes.

register

spgemm_csr

CSR × CSR → CSR via Gustavson’s algorithm. Pure-Rust convenience wrapper around algos::spgemm_csr — exposed outside the IR because sparsity patterns are structural (and typically static across the differentiable training loop), so a one-shot multiply at graph-build time is the natural shape. Returns (c_values, c_col_idx, c_row_ptr) for the output CSR.