cudax: __sharded — placed containers and algorithms for the in-process rung

[caugonnet]

The framing: CUB's cooperation-scope ladder, extended

CUB's structural invariant: the primitive at scope N = scope N−1's primitive run locally + a combine using what scope N shares.

scope	what is shared	combine mechanism	primitives
warp	registers	shuffles	WarpReduce, WarpScan, ...
block	shared memory	smem staging	BlockReduce, BlockScan, ...
device	global memory	global staging	DeviceReduce, DeviceScan, ...
places (one process)	one VA, placed pages	in-place fold over the shared address space	sharded algorithms — this PR (__sharded)
ranks (multi-process / multi-node)	nothing	message passing	communicator algorithms (__multi_gpu MGMN)

Each rung owns its combine. Sharded containers exist exactly at the highest rung that still shares an address space; where nothing is shared, the MGMN communicator algorithms take over. The two surfaces are scoped by rung and compose: hierarchical multi-process + placed-per-process is ordinary rung nesting, the same way DeviceReduce runs BlockReduce runs WarpReduce.

What this PR adds

A new sibling capability cudax/include/cuda/experimental/__sharded/ (public umbrella cuda/experimental/sharded.cuh, namespace cuda::experimental::sharded), peer to __multi_gpu, matching the ladder 1:1. It builds only on the places layer (and place_group from the previous PR in this stack) — no other dependencies.

Container tier

• shard<T>: one placed piece — data/size/capacity/global_offset + data_place/exec_place/reference stream + global↔local index math
• sharded_array<T>: allocation from explicit (size, data_place, exec_place, stream) specs or from a place_group (even split or explicit sizes; reference streams from the group's pools); adoption of existing shards (non-owning views); allocate_like; host transfer (from_host, copy_from_host, copy_to_host); slice with place correspondence preserved (empty shards keep their position); each_shard visitation with exec-place activation; sync; validation; copy_between for arbitrary re-sharding
• Contiguous backing contract (allocate_contiguous): shards become views into ONE contiguous VA range whose physical pages are owned per place (VMM via places::localized_array). Logical shard boundaries are exact (shard(i).data == contiguous_data() + global_offset); physical ownership snaps to the allocation granularity (typically 2 MiB). contiguous_data() hands the whole array to unmodified single-pointer consumers. The range is mapped once, so shard sizes are fixed: size-mutating operations must refuse such arrays (the contract is documented on the method and enforced with a throw by the size-mutating algorithms copy_if/filter/remove_if and unique). One VA also means one physical home per byte — the container holds partitioned data only; replicated operands belong to the binding tier (STF logical_data).

Algorithm tier (native, per-place CUB + combine)

• Elementwise (no cross-place stage): fill, sequence, iota, tabulate, generate, for_each, transform (in-place / unary / binary, with cross-array stream ordering and layout checks)
• reduce / sum / min / max: per-place cub::DeviceReduce through the group's stream + memory-resource environments, then a combine of the P partials
• inclusive_scan / exclusive_scan (+ sum aliases): per-place cub::DeviceScan, host prefix of the P shard totals, in-place prefix fold into each shard over the shared address space
• adjacent_difference: local differences + exactly one boundary element per shard
• count / count_if: per-place cub::DeviceReduce::TransformReduce (predicate mapped to 0/1), per-place counts summed — read-only, so available on every array including contiguous ones
• histogram_even: per-place cub::DeviceHistogram::HistogramEven, per-bin sum of the per-place histograms — read-only; invalid bins/levels refused with std::invalid_argument
• copy_if / filter / remove_if: per-place in-place cub::DeviceSelect::If (stable compaction to the front of each shard), then shard sizes and global offsets are updated; capacities are unchanged, so reset_sizes_to_capacity() reuses the buffers
• unique (std::unique semantics across shard boundaries): per-place in-place cub::DeviceSelect::Unique, then duplicates straddling shard boundaries are trimmed with an O(1) size decrement per boundary (compared against the previous non-empty shard, so runs spanning empty shards are handled)

The size-mutating algorithms (copy_if/filter/remove_if, unique) enforce the contiguous-backing contract: they throw std::invalid_argument on allocate_contiguous arrays, leaving them untouched — shrinking shard sizes would leave gaps between shards' valid elements, falsifying the read-as-one-array contract of contiguous_data(), and compacting across the gaps would migrate elements across the placement the caller asked for.

Every algorithm takes (place_group&, containers...) — the mapping-tier signature; algorithm temporaries are drawn from each shard's own place.

Tests (cudax/test/sharded/)

• containers/sharded_array.cu: allocation (specs / group / uniform), index math, host roundtrips, copy_between incl. misaligned resharding, adoption + slicing, allocate_like layout+placement preservation, contract throws (size-count mismatch, check_compatible)
• containers/contiguous.cu: exact-boundary + density checks, per-shard placement correspondence, per-shard writes read by one whole-range kernel through contiguous_data() (and the reverse), backing release on clear()
• algorithms/{elementwise,reduce_scan}.cu: correctness vs host references over multiple places (locality domains), custom operators, layout-mismatch throws, empty-input behavior
• algorithms/count_histogram.cu: count/count_if and histogram_even vs host references, including a contiguous-array case showing read-only algorithms stay available there; invalid-argument throws
• algorithms/compaction.cu: copy_if vs host reference incl. empty-result shards and a keep-nothing case, remove_if/filter, unique vs std::unique with runs straddling shard boundaries plus the one-value-everywhere degenerate chain, and the refuse-on-contiguous negative tests (copy_if, remove_if, unique) verifying the array is left untouched
• include_only.cu + header-test targets for the new headers

Example

cudax/examples/places/sharded_reduce.cu: place_group::by_locality_domains() → sharded_array::allocate → iota → sum, 256M elements.

Verification

GB300 node (sm_103a, CUDA 13.4, 2 locality domains): 26/26 places+sharded tests and both examples pass.

cmake --preset=cudax -DCMAKE_CUDA_ARCHITECTURES=103a -Dcudax_ENABLE_NCCL=OFF
ninja -C build/cudax cudax.test.sharded cudax.test.places cudax.example.places
ctest --test-dir build/cudax -R "cudax\.(test\.(places|sharded)|example\.places)"

Sequenced next

• a places-scoped communicator conformant to the __multi_gpu communicator concept — already validated out of tree: cudax::reduce and the MGMN HSS sort run unmodified over it — as the bridge between the two rungs
• MGMN constructs as engines behind sharded names where they are the right tool (sort via HSS first)
• a CI bridge test keeping the cross-rung interop honest

---

Opened on the fork for design review with Andrei. Main review points: the __sharded-as-sibling placement, the container/algorithm split, and the contiguous-backing contract.

🤖 Generated with Claude Code

Update: sharded_array<T> gains fork_from(cudaStream_t) / join_into(cudaStream_t) — ordering declarations (not synchronizations; the host returns immediately) bridging a caller stream and the per-shard streams via a small container-owned event pool (lazily created disableTiming events, reused across calls; correct for adopted arrays with foreign streams). Both are capture-safe and documented as the composition idiom with a caller stream or graph. Tests cover the eager no-host-sync chain, an adopted-array variant, and a capture variant.