[cudax] __sharded: sort -- the MGMN distributed sort as the engine behind the sharded name

[caugonnet]

Stacked on #5 (review this diff against sharded/communicator).

sharded::sort(group, data, comp) — a global, in-place sort of a sharded array — completes the two-tier picture introduced in #5:

• Tier 1 (container): owns placement, resources and bookkeeping. It manufactures the communicator / environment / iterator / size ranges with bind_engine — engine temporaries come from each rank's environment, so scratch lands on the place whose rank uses it.
• Tier 2 (engine): the __multi_gpu distributed sort runs over those ranges unmodified. The same MGMN construct that sorts across multi-process ranks sorts across in-process places, through the places communicator — the ladder's portability claim, exercised end to end.

Contract

The engine delivers to each rank its slice of the globally sorted sequence, redistributed back to the rank's original element count. That composes exactly with the container's fixed-size contract: shard sizes, offsets and capacities are unchanged by construction, and — unlike the size-mutating compaction algorithms — contiguous (allocate_contiguous) arrays are fully supported: after the sort, contiguous_data() reads as ONE globally sorted array through the base pointer. Sorting is not stable.

The engine slot is deliberately swappable behind the same name and contract: a future in-process, placement-aware specialization would be a performance change, not an API change.

Tests (cudax/test/sharded/algorithms/sort.cu, 2 locality domains)

• correctness vs std::sort of the concatenated shards (byte-exact) over mixed distributions — uniform random, all-equal, pre-sorted, reverse-sorted, duplicate-heavy ints — on uneven shards (first shard ~2x the others, odd total);
• custom comparator (descending) through the same path;
• shard metadata invariants: per-shard sizes, offsets and capacities preserved, total preserved, validate() clean after every sort;
• contiguous arrays: sorted in place and read back through contiguous_data() as one plain array, byte-exact vs the sorted host reference;
• repeated runs on identical input are byte-identical (checked empirically over 3 runs; for keys-only sorting this is also guaranteed structurally — the sorted multiset is unique);
• shape mismatch (shard count != place count) refused with std::invalid_argument before any engine work; empty arrays are a no-op.

All places + sharded tests green on GB300 (sm_103a, CUDA 13.4, 2 locality domains): 28/28 ctest, built with -Dcudax_ENABLE_NCCL=OFF.

Design review with Andrei.

🤖 Generated with Claude Code