perf(rbd): one-launch segmented bitonic sort for small per-batch key counts

src_rbd/utils/radix_sort/mod.rs

@@ -2,7 +2,7 @@
 
 use crate::shaders::utils::radix_sort::{
     GpuInitSortBatched, GpuInitSortDispatch, GpuSortCount, GpuSortReduce, GpuSortScan,
-    GpuSortScanAdd, GpuSortScatter, SortUniforms,
+    GpuSortScanAdd, GpuSortScatter, GpuSortSmall, SMALL_SORT_MAX, SMALL_SORT_WG, SortUniforms,
 };
 use khal::backend::{GpuBackend, GpuBackendError, GpuPass};
 use khal::{BufferUsages, Shader};
@@ -30,6 +30,7 @@ pub struct RadixSort {
     scan: GpuSortScan,
     scan_add: GpuSortScanAdd,
     scatter: GpuSortScatter,
+    small: GpuSortSmall,
 }
 
 /// Workspace buffers for radix sort operations.
@@ -124,6 +125,38 @@ impl RadixSort {
             );
         }
 
+        // Fast path: small per-batch capacity → one segmented bitonic-sort
+        // launch (output bit-identical to the radix path) instead of the
+        // ~ceil(bits/4)×5-launch radix cascade.
+        let per_batch_max = input_keys.len() as u32 / num_batches.max(1);
+        let small_enabled = std::env::var("NEXUS_SMALL_SORT")
+            .map(|v| v != "0")
+            .unwrap_or(true);
+        if small_enabled && per_batch_max > 0 && per_batch_max <= SMALL_SORT_MAX {
+            workspace.pass_uniforms.clear();
+            workspace.pass_uniforms.push(Tensor::scalar(
+                backend,
+                SortUniforms {
+                    shift: 0,
+                    max_keys_per_batch: per_batch_max,
+                    has_aux: 0,
+                },
+                BufferUsages::STORAGE | BufferUsages::UNIFORM,
+            )?);
+            // `call` takes logical THREAD counts (div_ceil'd by the workgroup
+            // size) — one 64-wide workgroup per batch.
+            return self.small.call(
+                pass,
+                [num_batches * SMALL_SORT_WG, 1, 1],
+                &workspace.pass_uniforms[0],
+                n_sort,
+                input_keys,
+                input_values,
+                output_keys,
+                output_values,
+            );
+        }
+
         if num_batches <= 1 {
             self.dispatch_single_batch(
                 backend,

src_rbd_shaders/utils/radix_sort/mod.rs

@@ -10,6 +10,7 @@ mod sort_reduce;
 mod sort_scan;
 mod sort_scan_add;
 mod sort_scatter;
+mod sort_small;
 mod sorting;
 
 pub use init_batch_ids::*;
@@ -19,4 +20,5 @@ pub use sort_reduce::*;
 pub use sort_scan::*;
 pub use sort_scan_add::*;
 pub use sort_scatter::*;
+pub use sort_small::*;
 pub use sorting::*;

src_rbd_shaders/utils/radix_sort/sort_small.rs

@@ -0,0 +1,112 @@
+//! Small-batch segmented sort: a one-launch replacement for the multi-pass
+//! radix cascade when the per-batch capacity is small.
+//!
+//! One workgroup per batch loads the batch's keys into shared memory,
+//! bitonic-sorts them, and writes the result — no batch-id passes, no
+//! ping-pong/histogram workspace. Bitonic networks are not stable, so the
+//! comparator orders by `(key, original_index)`; the output is therefore
+//! bit-identical to the stable radix path's. Capacity: `SMALL_SORT_MAX`
+//! elements per batch, padded to a power of two with `u32::MAX` sentinels.
+
+use khal_std::glamx::UVec3;
+use khal_std::index::MaybeIndexUnchecked;
+use khal_std::macros::{spirv, spirv_bindgen};
+use khal_std::sync::workgroup_memory_barrier_with_group_sync;
+
+use super::sorting::SortUniforms;
+
+/// Maximum per-batch capacity handled by the small sort (shared-memory bound).
+pub const SMALL_SORT_MAX: u32 = 128;
+/// Workgroup width: each thread owns ≤ 2 slots of the padded array.
+pub const SMALL_SORT_WG: u32 = 64;
+
+/// Segmented bitonic sort: one workgroup per batch.
+///
+/// Reuses [`SortUniforms`]: `max_keys_per_batch` is the per-batch capacity
+/// (`shift` / `has_aux` are unused). Like the radix path, only the first
+/// `n_sort[batch]` elements of each batch participate; output slots past the
+/// count are left untouched (consumers never read past the count).
+#[spirv_bindgen]
+#[spirv(compute(threads(64)))]
+pub fn gpu_sort_small(
+    #[spirv(local_invocation_id)] lid: UVec3,
+    #[spirv(workgroup_id)] gid: UVec3,
+    #[spirv(uniform, descriptor_set = 0, binding = 0)] config: &SortUniforms,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] n_sort: &[u32],
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] src: &[u32],
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] values: &[u32],
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] out: &mut [u32],
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] out_values: &mut [u32],
+    #[spirv(workgroup)] keys: &mut [u32; 128],
+    #[spirv(workgroup)] idxs: &mut [u32; 128],
+) {
+    let batch = gid.x;
+    let tid = lid.x;
+    let n = n_sort
+        .read(batch as usize)
+        .min(config.max_keys_per_batch)
+        .min(SMALL_SORT_MAX);
+    let base = (batch * config.max_keys_per_batch) as usize;
+    if n == 0 {
+        return;
+    }
+
+    // Padded (power-of-two) problem size.
+    let mut p = 1u32;
+    while p < n {
+        p <<= 1;
+    }
+
+    // Load keys (+ identity index) into shared memory; pad with MAX sentinels.
+    let mut i = tid;
+    while i < p {
+        let k = if i < n {
+            src.read(base + i as usize)
+        } else {
+            u32::MAX
+        };
+        keys.write(i as usize, k);
+        idxs.write(i as usize, i);
+        i += SMALL_SORT_WG;
+    }
+    workgroup_memory_barrier_with_group_sync();
+
+    // Bitonic sort, ascending on (key, original_index) — the index tiebreak
+    // makes it stable.
+    let mut k = 2u32;
+    while k <= p {
+        let mut j = k >> 1;
+        while j > 0 {
+            let mut i = tid;
+            while i < p {
+                let partner = i ^ j;
+                if partner > i {
+                    let ascending = (i & k) == 0;
+                    let (ka, ia) = (keys.read(i as usize), idxs.read(i as usize));
+                    let (kb, ib) = (keys.read(partner as usize), idxs.read(partner as usize));
+                    // Lexicographic (key, index) "a greater than b".
+                    let a_gt_b = ka > kb || (ka == kb && ia > ib);
+                    if a_gt_b == ascending {
+                        keys.write(i as usize, kb);
+                        idxs.write(i as usize, ib);
+                        keys.write(partner as usize, ka);
+                        idxs.write(partner as usize, ia);
+                    }
+                }
+                i += SMALL_SORT_WG;
+            }
+            workgroup_memory_barrier_with_group_sync();
+            j >>= 1;
+        }
+        k <<= 1;
+    }
+
+    // Write back: keys directly, values gathered through the sorted index.
+    let mut i = tid;
+    while i < n {
+        out.write(base + i as usize, keys.read(i as usize));
+        let src_idx = idxs.read(i as usize) as usize;
+        out_values.write(base + i as usize, values.read(base + src_idx));
+        i += SMALL_SORT_WG;
+    }
+}