Skip to content

basic support for sparse tensors (2) #3201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
simeonschaub wants to merge 11 commits into
base: main
Choose a base branch
from
Draft

basic support for sparse tensors (2) #3201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
862dac8
Add minimal sparse_tensor-dialect CSR spmv/spmm lowered to cuSPARSE/h…
simeonschaub Aug 18, 2026
6e6d784
fix FileCheck test
simeonschaub Aug 19, 2026
1890046
Move sparse CSR lowering into the Enzyme-JAX lower-sparse-csr pass
simeonschaub Aug 19, 2026
66397ab
Emit enzymexla.sparse.spmm from mul! for fused sparse alpha/beta
simeonschaub Aug 19, 2026
cd135d3
Pass ScratchAllocator by reference in CSR matmul impls
simeonschaub Aug 19, 2026
b150aa9
Apply clang-format and JuliaFormatter to sparse CSR sources
simeonschaub Aug 20, 2026
e87ab56
clang-format API.cpp
simeonschaub Aug 20, 2026
5dbced8
Document CSRMatrix and sparse_csr_spmm in the API reference
simeonschaub Aug 20, 2026
a65eaa5
Compare CSR round trip approximately in the sparse tests
simeonschaub Aug 20, 2026
71f7371
Follow the Enzyme-JAX rename to lower-enzymexla-sparse
simeonschaub Aug 21, 2026
0fc45b7
[do not merge] pin Enzyme-JAX to the sds/sparse_csr fork branch
simeonschaub Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions deps/ReactantExtra/API.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3566,7 +3566,7 @@ REACTANT_ABI void reactantXLAFree(LinkableRuntime **__restrict__ lrtP,
return;
auto lrt = *lrtP;
void *buffer = *(void **)buffer0;
auto erased = lrt->allocations.erase((void*)buffer0);
auto erased = lrt->allocations.erase((void *)buffer0);
assert(erased == 1);
(void)erased;
free(buffer0);
Expand Down Expand Up @@ -4035,7 +4035,8 @@ REACTANT_ABI void ReactantCreateLLVMMod(
llvm::LLVMContext **out_context, size_t *out_off, size_t *out_tmp_buf) {

std::string fn(fn_str ? std::string(fn_str, fn_len) : std::string());
llvm::StringRef source(source_str ? llvm::StringRef(source_str, source_len) : llvm::StringRef());
llvm::StringRef source(source_str ? llvm::StringRef(source_str, source_len)
: llvm::StringRef());

std::vector<llvm::SmallVector<int64_t>> out_shapes;
out_shapes.reserve(num_out_shapes);
Expand Down
7 changes: 7 additions & 0 deletions deps/ReactantExtra/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1390,6 +1390,7 @@ cc_library(
"@com_google_absl//absl/log:globals",
"@llvm-project//mlir:CAPIIRObjects",
"@llvm-project//mlir:CAPILLVMObjects",
"@llvm-project//mlir:CAPISparseTensorObjects",

# Broken upstream x/ref https://github.com/jax-ml/jax/issues/33344
# "@jax//jaxlib/mosaic:tpu_dialect_capi_objects",
Expand Down Expand Up @@ -1427,6 +1428,9 @@ cc_library(
"@xla//xla/stream_executor:cuda_platform",
"@xla//xla/stream_executor:kernel",
"@xla//xla/stream_executor/cuda:all_runtime",
# cuSPARSE (dlopen stub) + headers for the reactant_csr_matmul FFI handler
"@xla//xla/tsl/cuda:cusparse",
"@local_config_cuda//cuda:cuda_headers",
]) + if_rocm([
"@xla//xla/service:gpu_plugin",
"@xla//xla/pjrt/c:pjrt_c_api_gpu",
Expand All @@ -1435,6 +1439,9 @@ cc_library(
"@xla//xla/stream_executor:rocm_platform",
"@xla//xla/service/gpu:amdgpu_compiler",
"@xla//xla/backends/profiler/gpu:device_tracer",
# hipSPARSE + headers for the reactant_csr_matmul FFI handler
"@local_config_rocm//rocm:hipsparse",
"@local_config_rocm//rocm:rocm_headers",
]) + select({
# gloo tcp transport only builds on linux
"@xla//xla/tsl:macos": [
Expand Down
9 changes: 7 additions & 2 deletions deps/ReactantExtra/WORKSPACE
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,15 @@ NSYNC_COMMIT = "82b118aa7ace3132e517e2c467f8732978cf4023"

NSYNC_SHA256 = ""

ENZYMEXLA_COMMIT = "51a4cd6c58b1950f28b650167d1e946eec8ec0d5"
# [do not merge] sds/sparse_csr branch of simeonschaub/Enzyme-JAX (see
# ENZYMEXLA_REPO below); switch back to an EnzymeAD/Enzyme-JAX commit once
# https://github.com/EnzymeAD/Enzyme-JAX merges the sparse support.
ENZYMEXLA_COMMIT = "8508ac2af05574b59653c079a6c907e0a985cc9e"

ENZYMEXLA_SHA256 = ""

ENZYMEXLA_REPO = "https://github.com/simeonschaub/Enzyme-JAX"

http_archive(
name = "nsync",
sha256 = NSYNC_SHA256,
Expand Down Expand Up @@ -48,7 +53,7 @@ sed -i.bak0 "s,//:patches,@enzyme_ad//:patches,g" third_party/*/workspace.bzl
],
sha256 = ENZYMEXLA_SHA256,
strip_prefix = "Enzyme-JAX-" + ENZYMEXLA_COMMIT,
urls = ["https://github.com/EnzymeAD/Enzyme-JAX/archive/{commit}.tar.gz".format(commit = ENZYMEXLA_COMMIT)],
urls = ["{repo}/archive/{commit}.tar.gz".format(commit = ENZYMEXLA_COMMIT, repo = ENZYMEXLA_REPO)],
)

NEW_XLA_PATCHES = []
Expand Down
Loading
Loading