feat(wgpu): add global stable sorting

[SamJSui]

Summary

Closes #8.

• keep the existing shared-memory WGPU sort for axes up to 512 elements
• add a self-contained WGSL global path for larger axes: pack -> transformed stable keys -> tiled/global bitonic stages -> scatter
• share the implementation across sort, sort_with_indices, and argsort
• support U32, I32, and F32, both orders, and arbitrary axes
• preserve stable ties, including NaN ties and -0.0/+0.0
• detect NaNs by bit pattern and pad in transformed-key space so real infinities are retained

The branch is rebased onto the NaN-ordering fix from #9 and leaves the <=512 comparator owned by the existing shared sort_cmp.rs machinery. topk remains out of scope.

Validation

• cargo test --offline --release --features wgpu --test backend_parity_all backend_parity::sort -- --nocapture: 22 passed, 1 ignored
• explicit ignored 1,000,003-element physical-GPU validator: passed
• cargo test --offline --features wgpu --lib: 913 passed
• strict Clippy for the library and benchmark example: passed
• cargo fmt --all --check: passed

Coverage includes CPU/WGPU parity for all three supported dtypes, ascending and descending order, arbitrary axes, duplicate stability, NaNs, signed zero, and all three sorting APIs.

WGPU timing diagnostic

RTX 4070 Ti SUPER, driver 591.86. Five independent processes. Each process validates output before timing. Boundary is the public sort call plus queue completion; input upload and output readback are excluded.

Elements	Process medians (ms)	Median of medians
513	0.2130, 0.2019, 0.1994, 0.1978, 0.1960	0.1994 ms
4,097	0.5028, 0.5048, 0.5077, 0.5062, 0.5070	0.5062 ms
65,537	1.1490, 1.1754, 1.1547, 1.1538, 1.7023	1.1547 ms
1,000,003	5.9926, 3.2025, 3.3688, 5.9496, 5.9621	5.9496 ms

The one-million case is visibly bimodal, so these results establish scaling/capability rather than a stable throughput or speedup claim. Before this patch, WGPU sorting axes above 512 returned BackendLimitation; there is no runnable large-axis WGPU baseline.

[farhan-syah]

Reviewed this properly — a sorting network isn't something you can eyeball, so I re-implemented the shader logic standalone (pack → tiled 512 pass → global stages → scatter, including the key transform, padding-key selection and the index tiebreak) and diffed it against a stable reference.

It's correct. u32/i32/f32 at 513, 700, 1024, 1025, 1536 and 2049, both orders, duplicate-heavy inputs, NaN/±0.0/±inf, and multi-segment [2,513,3] — all matched exactly.

Two things I looked at hard and was happy with: the tiled pass XNORing the stage direction with tile parity isn't the textbook per-stage direction, but it's endpoint-equivalent, since flipping every comparator gives you the reverse-sorted tile that the global k=512 stage expects. And the padding keys do collide with real extremes (0xffffffff against u32::MAX and NaN), but the axis_index >= sort_size tiebreak forces padding into the tail you discard. Both are the kind of thing that looks wrong until you work it through.

Two things I'd like fixed before merge — project conventions rather than bugs:

• sort.rs:312,315 — .expect("output exists") in library code. The invariant holds today, but it's enforced only by the Option::or two lines above, so a later edit to the temporary-buffer construction turns a logic slip into a panic inside a GPU op. Worth restructuring so the buffer isn't an Option that has to be proven non-None.
• sort.rs is now 1115 lines against our 500-line limit, and the convention here is that a new operation means new files rather than growing existing ones. launch_global_sort plus GlobalSortParams is self-contained — could it move to sort_global.rs, next to the .wgsl it drives?

Performance notes, not merge gates:

• sort() allocates and writes a full-size temporary indices buffer, and argsort() a values buffer, because scatter_global_sort unconditionally writes both. Values-only and indices-only scatter entry points would drop ~4 MB and half the scatter stores per 1M call.
• The workspace buffers are created and destroyed on every call with no reuse — roughly 12 MB per 1M-element invocation. I suspect that's what's behind the bimodal 1M timings in your description, so it may be a fixable cost rather than device noise.
• Every j sub-stage for k >= 1024 goes out as its own global dispatch, including the j <= 256 ones that would fit entirely in the 512-element tile you already implemented. A merge-tail variant of the tile kernel would cut roughly a third of the 165 dispatches.

One minor thing: the >512 guard in src/ops/wgpu/sorting.rs now runs after get_tensor_buffer, so a shape with a zero-length dimension like [0, 1024] surfaces Internal("Buffer not found in registry") where that path used to return a clean backend-limitation error.

Good work — this is the scope we settled on in #8 and it landed clean.

[SamJSui]

Addressed in 1a84747: moved the global launcher and GlobalSortParams into sort_global.rs; replaced the temporary-output Option/expect proof with structurally owned fallback buffers; and returned empty outputs before buffer lookup for zero-element tensors. Coverage includes sort, argsort, and sort_with_indices on [0, 1024]. Validation: 23 active sorting tests passed; the ignored 1,000,003-element validator passed; all 913 library tests passed; strict Clippy and formatting passed.