cudax sharded: sharded_csr + opt-in cuSPARSE sparse products (spmv/spmm) with measured rebalance

[caugonnet]

Stacked on #6 (base sharded/sort); review only the last commit.

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 the sharded.cuh umbrella, vendor-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. 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 owns type-erased lib_state() slots so vendor layers can attach per-(shard, op) state without this header ever including a vendor header.
• <cuda/experimental/sharded_sparse.cuh> (opt-in, NOT included by the umbrella): sharded::spmv(group, A, x, y, alpha, beta) and sharded::spmm(group, A, B, C, n_cols, ...) (FP64-first, dtype-trait templated). The header #errors when the cuSPARSE headers are absent — the exact model of the in-tree <cuda/experimental/cufile.cuh> precedent. On the CMake side, a new cudax_ENABLE_CUSPARSE option (default OFF, following the cudax_ENABLE_CUDASTF_MATHLIBS precedent) gates the cuSPARSE-linked tests and the vendor header's header-test; CUDA::cusparse is linked for those targets only.

The call and plan model

Each product issues one confined cuSPARSE call per shard, on the shard's place stream (cusparseSetStream), with the shard's exec place active. The row partition makes the output row blocks disjoint, so there is never a combine step. Per-(shard, operation) cuSPARSE state — handle, descriptors, workspace (drawn from the shard's place), preprocessed plan — is created lazily on first call into the container's lib_state() and reused for the matrix's lifetime; later calls only rebind the dense pointers when they change. Plans are built once against the shard's fixed addresses and die with the matrix they describe.

Outputs are row-partitioned sharded_arrays from A.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-pointer consumers (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_times measure each shard solo through the exact call path the products use, and sharded_csr::time_balanced_boundaries converts one measurement into a time-equalizing split via a piecewise-rate model. One calibration round is amortized over every subsequent call on the rebuilt matrix — the natural fit for iterative consumers that reuse one operator across many products (repeat the round, keeping the best measured split, when the skew is extreme).

What composes later

Dense operands are plain device pointers in this PR (e.g. one whole-device allocation readable from every place). 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_data can materialize and cache a per-place instance and hand its pointer to these calls unchanged. The container deliberately does not absorb that role (one home per byte is the container invariant; replicated operands are the runtime's); wiring the STF-replicated operand path is the natural follow-up tier.

Tests (all silent)

• sparse/sharded_csr.cu (ungated, no cuSPARSE): construction/rebasing/split boundaries vs a host CSR reference on a mixed row-length matrix (empty rows included), default-split nnz balance, explicit/invalid boundaries, row-partitioned outputs (separate + contiguous, exact logical offsets), lib_state create-once caching, and deterministic host-side checks of the piecewise-rate model (exact boundary math, direction, monotonicity in the skew).
• sparse/spmv_spmm.cu (gated): FP64 correctness vs a host reference on mixed (empty rows) and skewed matrices sized to span both locality domains; comparison against ONE whole-matrix single call — bitwise for SpMM (CSR_ALG3); tight-tolerance for SpMV (CSR_ALG2's reduction shape varies with the call's row count, so a row split can differ in the last bits — noted in the test); alpha/beta accumulation; plan-reuse determinism across repeated calls; contiguous outputs read back through the single base pointer; mis-shaped output refusal.
• sparse/rebalance.cu (gated): on a deliberately time-skewed two-region matrix, asserts the boundary moves in the measured direction (the slower shard shrinks), that the model's predicted times at the new boundaries are equalized, and that the re-measured imbalance decreases — direction/monotonicity, not exact values.

Verified on a 2-locality-domain device (sm_103a, CUDA 13.4): 33/33 ctest green across the places+sharded suites including the three new tests; the cudax_ENABLE_CUSPARSE=OFF (default) variant configures and builds with the sparse targets absent and the ungated container test green; header tests green for both the umbrella (now including sharded_csr.cuh) and the gated vendor header.

Error handling follows the house safe-call's optional-vendor-status precedent, scoped to the vendor header (cusparse_safe_call), leaving __stf untouched.

🤖 Generated with Claude Code

Update (0612022): two workload-shaped tests added — test_value_mutation (in-place transform-rescale of the values between calls: plans stay valid with no rebuild, scaled-reference + whole-matrix agreement, exact ×2/×0.5 round trip restores byte-identical outputs) and containers/stream_ordered_lifecycle.cu (full allocate/write/destroy-with-work-in-flight/reallocate cycle on one external non-pool stream, single final sync; ordering/correctness asserts only).

Update (f8ec211, 67a182b) — handle lifecycle: library state is now split by natural scope. place_group gains a type-erased per-place library-state cache (keyed by place index and state type, lazily constructed, torn down with the group; __places stays vendor-free — objects are destroyed through deleters captured at creation), and the sparse products draw ONE cusparseHandle_t per place from it instead of creating a handle per (matrix, shard) — the raft::handle_t precedent. Descriptors, plans and workspaces remain container-owned (matrix-bound, retire with the storage); the handle's stream is rebound per call since one handle now serves every matrix built over the group. The group must outlive containers built over it (the existing group contract, now documented on the cache). Tests: places/place_group/lib_state.cu covers the generic cache vendor-free (identity, laziness, type/group isolation, teardown-destroys-exactly-once, move); sharded/sparse/handle_lifecycle.cu covers the vendor wiring (two matrices on one group share per-place handles, two groups do not, container destruction leaves the handle alive and usable). The second commit also fixes the plan-build warm-up launch to write into a scratch buffer instead of the caller's output, whose contents must survive for a first call carrying beta != 0 (caught by the existing alpha/beta accumulation arm).

Update: sharded_csr gains the same fork_from/join_into ordering members as sharded_array (the backing arrays share the matrix's per-place streams, so ordering the shard streams orders every consumer), with an ordering test in the ungated container suite.

[caugonnet]

Added two workload-shaped tests in 0612022:

• test_value_mutation (sparse/spmv_spmm.cu): a transform kernel rescales the sharded_csr values in place between calls — the pattern solvers use for problem scaling. Structure is untouched, so the per-shard plans must stay valid with no rebuild or re-preprocess: results are checked against the scaled host reference and against one whole-matrix call on the scaled operator, and an exact ×2.0/×0.5 round trip must restore byte-identical outputs (pinning that no plan or engine state drifts across mutations). This is the contract that lets an application own the values array as ordinary mutable data while the container owns placement and plans.

• containers/stream_ordered_lifecycle.cu (new, ungated): a full sharded_array cycle — allocate → write → copy out → host-side destroy with work still in flight → reallocate → write — carried by ONE external stream (a plain cudaStreamCreate stream outside any place_group pool, i.e. the shape of an application's own compute stream), with a single final synchronization. Asserts ordering/correctness only; host-side completion timing of the underlying pool operations is implementation- and driver-defined. Also exercises the lifetime rule that containers must be destroyed before the streams their specs reference.

Suite state on GB300 (sm_103a, CUDA 13.4, 2 locality domains, cudax_ENABLE_CUSPARSE=ON): 32/32 sharded+places ctest green including the two additions.

[caugonnet]

ba948bc adds the two integration seams for applications whose matrices already live on the GPU and whose existing kernels rewrite values in place:

• sharded_csr::from_device(...) — adopt a device CSR: offsets make one small host round trip (the split and per-shard rebasing are host-side either way), colinds/values are sliced device-to-device into the shards' places. Host and device construction share one builder.
• contiguous = true on construction — values/colinds are exact nnz-slices, so they admit the one-VA-range backing; contiguous_values()/contiguous_colinds() hand the whole array to unmodified writers (e.g. an existing problem-scaling kernel) as one normal pointer while each slice keeps per-place physical placement.

Together with the earlier mutation test this closes the loop: an application can own the matrix values as ordinary mutable device data (single whole-array writer, no knowledge of shards) while the container owns placement and plans. Tests cover construction equivalence vs the host path, contiguous read-back + exact shard-view offsets, and SpMM through a device-adopted contiguous matrix with a whole-array ×2/×0.5 mutation round trip restoring byte-identical results. Suite: 32/32 sharded+places ctest green on GB300.