basic support for sparse tensors (2) - #3201
Draft
simeonschaub wants to merge 11 commits into
Draft
simeonschaub wants to merge 11 commits into
simeonschaub wants to merge 11 commits into
Conversation
…ipSPARSE
Introduces an opaque `Reactant.CSRMatrix` (convertible from SparseMatrixCSC
via the SparseArrays ext). Inside traced code `A * x`, `A * B`, and `mul!`
emit `sparse_tensor.assemble` producing a CSR-encoded tensor consumed by
`stablehlo.dot_general`; `Compiler.lower_sparse_ops!` rewrites the pair to
`stablehlo.custom_call @reactant_csr_matmul` on the raw buffers before any
pass pipeline runs, so XLA only ever sees dense types. The custom call is
served by new cuSPARSE ("CUDA") and hipSPARSE ("ROCM") typed-FFI handlers
in ReactantExtra (requires a local jll rebuild for execution).
`@code_hlo optimize=:none` keeps the sparse_tensor IR for inspection.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replaces the Julia-side IR rewrite (src/compiler/SparseLowering.jl) with the new `lower-sparse-csr` MLIR pass in Enzyme-JAX, invoked via run_pass_pipeline! only when the traced module actually contains sparse_tensor ops (so older jlls without the pass keep working for dense code). Also switches the CSR buffers to 0-based indices as required by the sparse_tensor dialect; the pass emits index_base = 0 and the FFI handlers already support both bases. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
`mul!(C, A, B, alpha, beta)` (and `A * B`) on a CSRMatrix now emit the semantic `enzymexla.sparse.spmm alpha, A, B, beta, C` op instead of a bare dot_general, so lower-sparse-csr can fuse constant alpha/beta into a single accumulating cuSPARSE/hipSPARSE call (reactant_csr_matmul_acc, with C aliased to the output and copied on-device when XLA does not donate the buffer). Runtime (traced) alpha/beta keep the unfused custom_call + multiply/add form, now materialized by the pass instead of the Julia frontend. The plain handler gains an f64 "alpha" attribute. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
ffi::ScratchAllocator is move-only; the entry points were copying it into the shared impl. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simeonschaub
force-pushed
the
sds/sparse_csr2
branch
2 times, most recently
from
August 20, 2026 12:46
79fb079 to
74cf002
Compare
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Pre-existing violations on main that make the format-check-cpp job fail. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simeonschaub
force-pushed
the
sds/sparse_csr2
branch
from
August 20, 2026 12:47
74cf002 to
6baf30a
Compare
Documenter aborted with :missing_docs because the two new sparse docstrings were not included in any @docs block. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
TPUs emulate Float64, so `SparseMatrixCSC(to_rarray(A)) == A` fails there by a few ULPs. Keep nnz exact and compare values with isapprox. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simeonschaub
force-pushed
the
sds/sparse_csr2
branch
from
August 20, 2026 19:45
6baf30a to
4ac6ca7
Compare
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Local builds (deps/build_local.jl) were still fetching the upstream Enzyme-JAX pin, which lacks enzymexla.sparse.spmm and the lower-enzymexla-sparse pass, so the emitted op was unregistered at compile time. Introduces ENZYMEXLA_REPO so the download URL follows the pin. Pinned to simeonschaub/Enzyme-JAX@8508ac2a (sds/sparse_csr), which is rebased on current EnzymeAD/Enzyme-JAX main (includes 51a4cd6, the pin on Reactant main) plus the lower-enzymexla-sparse pass, enzymexla.sparse.spmm, the getSHLOLayout overload disambiguation fix, and clang-format fixes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simeonschaub
force-pushed
the
sds/sparse_csr2
branch
from
August 21, 2026 10:23
4ac6ca7 to
0fc45b7
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Companion to EnzymeAD/Enzyme-JAX#2922. Replaces #3198 and moves the lowering to Enzyme-JAX, which gains a
sparse.spmmop. I also added support for emitting a single cuSPARSE/hipSPARSE call, even ifmul!is called with different alpha and beta (they have to be constant)