cudax sharded: sharded_csr + opt-in cuSPARSE sparse products (spmv/spmm) with measured rebalance - #10970
Conversation
…free) One self-contained CSR shard per place_group place: the shard's nnz slice of values/colinds plus its rows+1 offsets REBASED to zero, stored in the shard's place. Each shard is a complete CSR operator over its row range, so any library that understands pointers and a stream can consume it with one ordinary per-place call. Split rows are caller-suppliable (nnz-balanced default); make_row_partitioned(n_cols[, contiguous]) builds matching output arrays including the VMM-backed contiguous form; fork_from/join_into order all backing arrays' streams. place_group gains a type-erased per-place library-state cache (lib_state: keyed by place index and state type, lazily constructed, torn down with the group through deleters captured at creation) so vendor layers can attach per-place state without __places or this header gaining vendor includes. The group must outlive containers built over it (the existing group contract, documented on the cache). Tests: sparse/sharded_csr.cu (construction/rebasing/boundaries vs a host CSR reference, empty rows, row-partitioned outputs incl. contiguous, lib_state create-once, fork/join ordering, host-side checks of the piecewise-rate model), places/place_group/lib_state.cu (generic cache, vendor-free), containers/stream_ordered_lifecycle.cu. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…red rebalance <cuda/experimental/sharded_sparse.cuh> (NOT in the umbrella; #errors without the cuSPARSE headers — the cufile.cuh precedent): spmv(group, A, x, y, alpha, beta) and spmm(group, A, B, C, n_cols, ...), FP64-first with a dtype trait. One confined cuSPARSE call per shard on the shard's place stream with the shard's exec place active; the row partition makes output row blocks disjoint, so there is never a combine step. ONE cusparseHandle_t per place drawn from the group's lib_state cache (stream rebound per call); descriptors, preprocessed plans and workspaces are container-owned, built lazily against the shard's fixed addresses, and retire with the storage. The plan-build warm-up launch writes a scratch buffer, never the caller's output. Measured rebalance: an nnz-balanced split is not a TIME-balanced split under per-place SM confinement. spmv/spmm_shard_times measure each shard solo through the exact call path; time_balanced_boundaries turns one measurement into a time-equalizing split via a piecewise-rate model. CMake: new option cudax_ENABLE_CUSPARSE (default OFF, the mathlibs precedent) gating the cuSPARSE-linked tests and the vendor header's header-test; CUDA::cusparse is linked for those targets only. Tests (gated): spmv_spmm.cu (host references on mixed/skewed matrices spanning both locality domains; whole-matrix single-call comparison — bitwise for SpMM/CSR_ALG3, tight-tolerance for SpMV/CSR_ALG2 whose reduction shape varies with the call's row count; alpha/beta; plan reuse; value mutation; contiguous outputs; shape refusal), rebalance.cu (boundary direction, equalized model prediction, decreasing measured skew), handle_lifecycle.cu (matrices share per-place handles, groups do not, containers die before the handle). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
@coderabbitai full review |
✅ Action performedFull review finished. |
| "cuda/experimental/__sharded/*" | ||
| # Opt-in vendor header: #errors without the cuSPARSE headers, compiled | ||
| # separately when cudax_ENABLE_CUSPARSE is set: | ||
| "cuda/experimental/sharded_sparse.cuh" |
There was a problem hiding this comment.
This should not be here, but part of __sharded
There was a problem hiding this comment.
Since it now lives in __sharded/, this is subsumed by the ungated __sharded/*.cuh header-test bucket — no separate cudax_ENABLE_CUSPARSE-gated header-test target needed at all. Its #error guard only checks __has_include(<cusparse.h>), which the CUDA toolkit always satisfies, and header tests are compile-only (OBJECT libraries, no link), so there was nothing to gate. Dropped the exclude and the now-fully-redundant gated block.
| $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--expt-relaxed-constexpr> | ||
| ) | ||
|
|
||
| if (cudax_ENABLE_CUSPARSE) |
There was a problem hiding this comment.
We already have some options for MATHLIBS, no need for one for CUSPARSE, one for CUBLAS, ...
| * Thread-safe: concurrent calls for the same slot yield the same object. | ||
| */ | ||
| template <typename _State, typename _Make> | ||
| _State& lib_state(size_t place_idx, _Make&& make) |
There was a problem hiding this comment.
is this the equivalent of raft handle ?
There was a problem hiding this comment.
this has to move to __sharded, we can't pollute the cudax dir like that
There was a problem hiding this comment.
Done: moved to cuda/experimental/__sharded/sparse.cuh (guard renamed, every include path and doc reference updated — sharded.cuh, sharded_csr.cuh, the three sparse tests, docs/cudax/sharded.rst). Verified it still compiles standalone on this node.
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesThe pull request adds a move-only, row-partitioned Sharded sparse functionality
Suggested reviewers: Merge Risk: 🟠 High · up to The new sparse-product path can hang during per-place state initialization and can release workspace before asynchronous GPU work finishes, risking stalled calls or incorrect results; its first-call synchronization also conflicts with the documented asynchronous and stream-capture behavior. These current-head issues should be fixed before merge. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
cudax/include/cuda/experimental/__sharded/sharded_csr.cuh (1)
593-609: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valuesuggestion:
lib_statecallsmake()and only afterwards wraps the raw pointer in ashared_ptr. Ifemplacethrows (allocation failure), the state object leaks. Build the owning pointer first.- _State* s = make(); - it = lib_state_ - .emplace(key, - ::std::shared_ptr<void>(s, - [](void* p) { - delete static_cast<_State*>(p); - })) - .first; + ::std::shared_ptr<void> owned(make(), [](void* p) { + delete static_cast<_State*>(p); + }); + it = lib_state_.emplace(key, mv(owned)).first;cudax/test/sharded/sparse/spmv_spmm.cu (1)
438-440: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winsuggestion: line 440 asserts bit-identical results between the sharded SpMM (one cuSPARSE call per row block) and one whole-matrix cuSPARSE call. The test itself explains at Line 371-376 why the same claim does not hold for SpMV: the reduction shape of the algorithm depends on the call's row count.
CUSPARSE_SPMM_CSR_ALG3currently happens to match, but that is a cuSPARSE implementation detail, not a documented guarantee, and it can change with the cuSPARSE version, the GPU architecture, or the number of places.expect_closeat line 439 already covers the correctness intent. Consider dropping thememcmp, or gate it and document it as version-dependent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bd252833-79bb-4ffc-90ab-4a0aefa2846f
📒 Files selected for processing (13)
cudax/CMakeLists.txtcudax/cmake/cudaxHeaderTesting.cmakecudax/include/cuda/experimental/__places/place_group.cuhcudax/include/cuda/experimental/__sharded/sharded_csr.cuhcudax/include/cuda/experimental/sharded.cuhcudax/include/cuda/experimental/sharded_sparse.cuhcudax/test/sharded/CMakeLists.txtcudax/test/sharded/containers/stream_ordered_lifecycle.cucudax/test/sharded/sparse/handle_lifecycle.cucudax/test/sharded/sparse/rebalance.cucudax/test/sharded/sparse/sharded_csr.cucudax/test/sharded/sparse/spmv_spmm.cudocs/cudax/sharded.rst
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| ::std::lock_guard<::std::mutex> lock(mutex_); | ||
| auto& slot = lib_state_cache_[place_idx]; | ||
| auto it = slot.find(::std::type_index(typeid(_State))); | ||
| if (it == slot.end()) | ||
| { | ||
| _State* s = make(); | ||
| it = slot | ||
| .emplace(::std::type_index(typeid(_State)), | ||
| ::std::shared_ptr<void>(s, | ||
| [](void* p) { | ||
| delete static_cast<_State*>(p); | ||
| })) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
important: Do not invoke make() while mutex_ is locked. A factory can call group.get_stream(place_idx) during library-state initialization. That call locks mutex_ again and blocks forever. Reserve an in-progress slot under the lock, release the lock for factory execution, then publish the state and notify waiters. Add a regression test where make() calls get_stream().
| wplace = sh.place; | ||
| const size_t wbytes = workspace_bytes == 0 ? 16 : workspace_bytes; | ||
| workspace = wplace.allocate(static_cast<::std::ptrdiff_t>(wbytes), stream); | ||
| cuda_safe_call(cudaStreamSynchronize(stream)); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
important: build() calls cudaStreamSynchronize(stream) twice. build() runs on the first spmv/spmm call, so that first call blocks the host, which contradicts the "ASYNCHRONOUS with respect to the host" contract documented at Line 611 and in docs/cudax/sharded.rst. It also makes the first call fail under an active stream capture, while docs/cudax/sharded.rst states that only allocation, host transfers and synchronization refuse during capture.
Either document the first-call blocking and capture behavior, or remove the syncs if the allocation and the warm-up are already stream-ordered on stream.
Also applies to: 253-253, 378-378, 414-414
| if (workspace) | ||
| { | ||
| const size_t wbytes = workspace_bytes == 0 ? 16 : workspace_bytes; | ||
| _CCCL_TRY | ||
| { | ||
| wplace.deallocate(workspace, wbytes, nullptr); | ||
| } | ||
| _CCCL_CATCH_ALL {} | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
important: the workspace is allocated stream-ordered on stream (line 222) but released on the null stream. spmv/spmm are asynchronous, so a matrix can be destroyed while the last SpMV that reads the workspace is still running on the shard stream. The null stream does not order against the group's shard streams if those streams were created with cudaStreamNonBlocking. The pool can then hand the bytes to another allocation before the kernel finishes.
Store the shard stream in the plan and free on it, so the free stays behind the enqueued work. The same applies to spmm_shard_plan at Line 482-490.
void* workspace = nullptr;
size_t workspace_bytes = 0;
data_place wplace; //!< place the workspace was drawn from
+ cudaStream_t wstream = nullptr; //!< stream the workspace was allocated on workspace = wplace.allocate(static_cast<::std::ptrdiff_t>(wbytes), stream);
+ wstream = stream;- wplace.deallocate(workspace, wbytes, nullptr);
+ wplace.deallocate(workspace, wbytes, wstream);| Each call runs one cuSPARSE call per shard on the shard's place stream | ||
| (``cusparseSetStream``). Per-(shard, operation) library state — handle, | ||
| descriptors, workspace, preprocessed plan — is created lazily on the first | ||
| call into the container's type-erased ``lib_state()`` slots, built once | ||
| against the shard's fixed addresses, and reused for the matrix's lifetime; | ||
| subsequent calls only rebind the dense pointers when they change. The row | ||
| partition makes the output row blocks disjoint, so there is never a combine | ||
| step, and outputs compose with ``allocate_contiguous`` backings unchanged. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
important: the sentence lists the handle as part of the per-(shard, operation) state held in the container's lib_state(). The implementation does the opposite. cudax/include/cuda/experimental/sharded_sparse.cuh stores the cusparseHandle_t per place in the place_group library-state cache, so it is shared by every matrix built over the group and destroyed with the group. cudax/test/sharded/sparse/handle_lifecycle.cu asserts that behavior. Correct the text so readers do not assume a handle dies with the matrix.
-Each call runs one cuSPARSE call per shard on the shard's place stream
-(``cusparseSetStream``). Per-(shard, operation) library state — handle,
-descriptors, workspace, preprocessed plan — is created lazily on the first
-call into the container's type-erased ``lib_state()`` slots, built once
-against the shard's fixed addresses, and reused for the matrix's lifetime;
-subsequent calls only rebind the dense pointers when they change. The row
+Each call runs one cuSPARSE call per shard on the shard's place stream
+(``cusparseSetStream``). The ``cusparseHandle_t`` is PER PLACE: it is created
+lazily in the ``place_group``'s library-state cache, shared by every matrix
+built over the group, and destroyed with the group. Per-(shard, operation)
+state — descriptors, workspace, preprocessed plan — is matrix-bound: it is
+created lazily on the first call into the container's type-erased
+``lib_state()`` slots, built once against the shard's fixed addresses, and
+reused for the matrix's lifetime; subsequent calls only rebind the dense
+pointers when they change. The row📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Each call runs one cuSPARSE call per shard on the shard's place stream | |
| (``cusparseSetStream``). Per-(shard, operation) library state — handle, | |
| descriptors, workspace, preprocessed plan — is created lazily on the first | |
| call into the container's type-erased ``lib_state()`` slots, built once | |
| against the shard's fixed addresses, and reused for the matrix's lifetime; | |
| subsequent calls only rebind the dense pointers when they change. The row | |
| partition makes the output row blocks disjoint, so there is never a combine | |
| step, and outputs compose with ``allocate_contiguous`` backings unchanged. | |
| Each call runs one cuSPARSE call per shard on the shard's place stream | |
| (``cusparseSetStream``). The ``cusparseHandle_t`` is PER PLACE: it is created | |
| lazily in the ``place_group``'s library-state cache, shared by every matrix | |
| built over the group, and destroyed with the group. Per-(shard, operation) | |
| state — descriptors, workspace, preprocessed plan — is matrix-bound: it is | |
| created lazily on the first call into the container's type-erased | |
| ``lib_state()`` slots, built once against the shard's fixed addresses, and | |
| reused for the matrix's lifetime; subsequent calls only rebind the dense | |
| pointers when they change. The row | |
| partition makes the output row blocks disjoint, so there is never a combine | |
| step, and outputs compose with ``allocate_contiguous`` backings unchanged. |
<cuda/experimental/sharded_sparse.cuh> lived flat in cuda/experimental/, following the cufile.cuh opt-in-vendor-header precedent -- but unlike cuFile, this header is not an independent feature: it is the sparse tier of the sharded rung, so it belongs alongside the rest of that layer rather than adding a loose top-level file for every future vendor-backed sharded product. Moved to <cuda/experimental/__sharded/sparse.cuh> (include guard renamed to match); every consumer's include path and doc-comment reference updated (sharded.cuh, sharded_csr.cuh, the three sparse tests, docs/cudax/sharded.rst). cudaxHeaderTesting.cmake: dropped the now-redundant top-level EXCLUDES entry (already covered by the __sharded/* wildcard), added an EXCLUDES on the ungated "Sharded headers" bucket so it does not try to compile the vendor header without cuSPARSE linked, and repointed the cudax_ENABLE_CUSPARSE-gated bucket's glob at the new path. Verified: the moved header compiles standalone against cusparse.h on this node. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sparse.cuh's #error guard only checks __has_include(<cusparse.h>), which the CUDA toolkit always satisfies -- and header tests are compile-only (OBJECT libraries, no link step), so the header never needed cudax_ENABLE_CUSPARSE gating in cudaxHeaderTesting.cmake at all: the ungated "Sharded headers" bucket already compiles it correctly once it lives in __sharded/. Drops the EXCLUDES entry that re-routed it and the now-fully-redundant gated header-test block below it (it linked nothing cuSPARSE-specific -- just cudax.compiler_interface again). Runtime tests (spmv_spmm.cu, rebalance.cu, handle_lifecycle.cu), which DO need CUDA::cusparse linked, keep their existing cudax_ENABLE_CUSPARSE gate in test/sharded/CMakeLists.txt -- unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sparse.cuh being ungated needs no explaining -- it is just an ordinary member of __sharded/*.cuh, nothing exceptional to call out at this exclude entry. cudaxHeaderTesting.cmake is now identical to NVIDIA#10957's version except for the two real deltas from U5's own header-test additions (the stale flat-path EXCLUDES entry, the redundant cudax_ENABLE_CUSPARSE-gated block) that this PR's move made unnecessary. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_spmm_state() keyed the container's lib_state cache by "cusparse_spmm:" + n_cols, unlike every other call-varying input in this file (x/y, B/C pointers), which is handled by rebinding a single retained plan in place. That made n_cols-keying an outlier: every distinct n_cols a caller ever passed permanently retained its own full spmm_state (one spmm_shard_plan -- descriptors + workspace -- per shard) for the container's lifetime, with no eviction, plus a std::string allocation on every single spmm() call just to do the lookup. n_cols is a shape, exactly like rows/cols/nnz -- it belongs to the plan, not to a cache key. spmm_shard_plan now tracks bound_n_cols alongside bound_B/bound_C; run() tears the descriptors/workspace down and rebuilds when n_cols no longer matches (dimensions are baked into mB/mC at creation, so unlike a pointer they cannot be rebound with SetValues) via a shared destroy_resources() helper also used by the destructor. get_spmm_state()/get_spmv_state() are now symmetric: one fixed-string-keyed slot each, no value-dependent key. Verified: the moved+fixed header still compiles standalone. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 participant
Description
Motivation
Sparse products are the workload that most wants placement: SpMV/SpMM
over a row-partitioned matrix decomposes into fully independent
per-place calls with disjoint output row blocks — no combine at all —
but CCCL offers no container that carries a CSR operator's placement,
and closed vendor libraries (cuSPARSE) cannot be taught placement from
the inside. The container has to make each piece look like an ordinary
complete matrix so the library never has to change.
This PR adds the sparse tier of the sharded rung: a row-partitioned CSR
container, and cuSPARSE-backed sparse products where a closed vendor
library is the engine — the container carries the placement, the
library never changes.
The gating split: container unconditional, vendor calls opt-in
sharded_csr<T>(in thesharded.cuhumbrella, vendor-free): oneself-contained CSR shard per
place_groupplace — the shard's nnzslice of values/colinds plus its rows+1 offsets rebased to zero,
stored in the shard's place. Because each shard is a complete CSR
operator over its row range, any library that understands pointers
and a stream can consume it with one ordinary per-place call. Split
rows are caller-suppliable (nnz-balanced default), and the container
attaches per-(shard, op) vendor state through the group's type-erased
lib_state()slots without this header ever including a vendorheader.
from_devicecopies device-to-device into per-place storage(snapshot semantics — nothing aliases; the caller may free after
return);
adopt-style zero-copy is the arrays' domain.fork_from/join_intoorder all backing arrays' streams against acaller stream.
<cuda/experimental/sharded_sparse.cuh>(opt-in, NOT included bythe umbrella):
sharded::spmv(group, A, x, y, alpha, beta)andsharded::spmm(group, A, B, C, n_cols, ...)(FP64-first, dtype-traittemplated). The header
#errors when the cuSPARSE headers are absent— the model of the in-tree
<cuda/experimental/cufile.cuh>precedent. On the CMake side, a new
cudax_ENABLE_CUSPARSEoption(default OFF, following the
cudax_ENABLE_CUDASTF_MATHLIBSprecedent) gates the cuSPARSE-linked tests and the vendor header's
header-test;
CUDA::cusparseis linked for those targets only.The call, handle and plan model
Each product issues one confined cuSPARSE call per shard, on the
shard's place stream, with the shard's exec place active. The row
partition makes the output row blocks disjoint, so there is never a
combine step. Library state is split by natural scope:
place_groupgains a vendor-free, type-erased per-place library-state cache (keyed by
place index and state type, lazily constructed, torn down with the
group), and the products draw ONE
cusparseHandle_tper place from it —so N matrices over one group share P handles, and the handle's stream is
rebound per call. Descriptors, preprocessed plans and workspaces remain
container-owned (matrix-bound, built lazily on first call against the
shard's fixed addresses, retiring with the storage); later calls only
rebind dense pointers when they change. The group must outlive
containers built over it (the existing group contract, documented on
the cache).
Outputs are row-partitioned
sharded_arrays fromA.make_row_partitioned(n_cols[, contiguous]), including the contiguous(VMM-backed) backing: the products write exact row ranges, and
contiguous_data()hands the result to unmodified single-pointerconsumers (positively tested).
Measured rebalance
An nnz-balanced split is not a TIME-balanced split: with per-place SM
confinement a call finishes at max(shard time), so a skewed row-length
distribution makes the default split pay the full skew.
spmv_shard_times/spmm_shard_timesmeasure each shard solothrough the exact call path the products use, and
sharded_csr::time_balanced_boundariesconverts one measurement into atime-equalizing split via a piecewise-rate model. One calibration round
is amortized over every subsequent call on the rebuilt matrix.
What composes later
Dense operands are plain device pointers in this PR. Which per-place
copies of a re-read operand should exist — and when a write makes
them stale — is a coherence question that belongs to the binding tier:
an STF
logical_datacan materialize a per-place instance and hand itspointer to these calls unchanged. The container deliberately does not
absorb that role (one home per byte is the container invariant).
Tests (all silent)
sparse/sharded_csr.cu(ungated): construction/rebasing/splitboundaries vs a host CSR reference (empty rows included), default
nnz balance, explicit/invalid boundaries, row-partitioned outputs
(separate + contiguous, exact offsets),
lib_statecreate-once,fork/join ordering, and deterministic host-side checks of the
piecewise-rate model.
sparse/spmv_spmm.cu(gated): FP64 correctness vs a host referenceon mixed and skewed matrices spanning both locality domains;
whole-matrix single-call comparison — bitwise for SpMM (CSR_ALG3),
tight-tolerance for SpMV (CSR_ALG2's reduction shape varies with the
call's row count; noted in the test); alpha/beta accumulation;
plan-reuse determinism; value mutation between calls (plans stay
valid, exact x2/x0.5 round trip restores byte-identical outputs);
contiguous outputs through the base pointer; mis-shaped refusal.
sparse/rebalance.cu(gated): on a deliberately time-skewedtwo-region matrix, asserts the boundary moves in the measured
direction, the model's predicted times equalize, and the re-measured
imbalance decreases.
sparse/handle_lifecycle.cu(gated): two matrices on one group shareper-place handles, two groups do not, container destruction leaves
the handle alive and usable. The generic cache is covered vendor-free
by
UNITTESTblocks next to it inplace_group.cuh(the__placesidiom), run by the unittested-header target.
containers/stream_ordered_lifecycle.cu: allocate/write/destroy-with-work-in-flight/reallocate on one external stream, single final sync.
Verified on a 2-locality-domain GB300 node (sm_103a, CUDA 13.4), both
cudax_ENABLE_CUSPARSE=ONand the default OFF variant (sparse targetsabsent, ungated tests green).
Launch build tally: 35/35 places+sharded ctest green with
cudax_ENABLE_CUSPARSE=ON; 32/32 with the default OFF (the gated targets are correctly absent and the ungated container tests stay green). Note for CI: the option defaults OFF, so upstream CI exercises the ungated path; the gated suites are validated on a 2-locality-domain GB300 (sm_103a, CUDA 13.4).Checklist
🤖 Generated with Claude Code