[RFC] basic support for the sparse tensor dialect - #3198
Draft
simeonschaub wants to merge 2 commits into
Draft
simeonschaub wants to merge 2 commits into
simeonschaub wants to merge 2 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>
Collaborator
|
Hey Simeon! Thanks for the PR. Would it be possible to port the lowerings from this PR in Enzyme-JAX. That way we can also include other missing pieces such as AD rules for the sparse ops. You can create a subdialect |
Contributor
Author
|
Sure, I can give that a go! |
This was referenced Aug 20, 2026
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
2 participants
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.
Introduces
Reactant.CSRMatrix, hooks it up to https://mlir.llvm.org/docs/Dialects/SparseTensorOps/ and adds the ability to lower simple spmv and spmm ops to either cuSPARSE or hipSPARSE. I verified the cuSPARSE path locally.Putting this up for comments, since I'm not sure this is the right approach, but this would be really useful to us in CoolPDLP.jl!