feat(wgpu): add global stable sorting

examples/wgpu_global_sort_bench.rs

@@ -0,0 +1,64 @@
+//! WGPU global-sort wall-time diagnostic.
+//!
+//! Run with:
+//! `cargo run --release --example wgpu_global_sort_bench --features wgpu`
+
+#[cfg(feature = "wgpu")]
+fn main() -> numr::error::Result<()> {
+    use numr::ops::SortingOps;
+    use numr::runtime::Runtime;
+    use numr::runtime::RuntimeClient;
+    use numr::runtime::wgpu::{WgpuDevice, WgpuRuntime};
+    use numr::tensor::Tensor;
+    use std::time::Instant;
+
+    let device = WgpuDevice::new(0);
+    let client = WgpuRuntime::default_client(&device);
+    let sizes = [513usize, 4_097, 65_537, 1_000_003];
+
+    println!("boundary=public sort call + queue completion; input upload/output readback excluded");
+    for size in sizes {
+        let data: Vec<u32> = (0..size as u32)
+            .map(|index| index.wrapping_mul(747_796_405).wrapping_add(2_891_336_453))
+            .collect();
+        let input = Tensor::from_slice(&data, &[size], &device);
+
+        let validation: Vec<u32> = client.sort(&input, 0, false)?.to_vec();
+        assert!(validation.windows(2).all(|pair| pair[0] <= pair[1]));
+
+        for _ in 0..3 {
+            let output = client.sort(&input, 0, false)?;
+            client.synchronize();
+            std::hint::black_box(output);
+        }
+
+        let sample_count = if size >= 1_000_000 {
+            7
+        } else if size >= 65_000 {
+            11
+        } else {
+            21
+        };
+        let mut samples_ms = Vec::with_capacity(sample_count);
+        for _ in 0..sample_count {
+            let start = Instant::now();
+            let output = client.sort(&input, 0, false)?;
+            client.synchronize();
+            samples_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
+            std::hint::black_box(output);
+        }
+        samples_ms.sort_by(f64::total_cmp);
+        println!(
+            "size={size} samples={sample_count} median_ms={:.6} min_ms={:.6} max_ms={:.6}",
+            samples_ms[sample_count / 2],
+            samples_ms[0],
+            samples_ms[sample_count - 1]
+        );
+    }
+    Ok(())
+}
+
+#[cfg(not(feature = "wgpu"))]
+fn main() {
+    eprintln!("enable the `wgpu` feature");
+}

src/ops/wgpu/sorting.rs

@@ -12,7 +12,7 @@ use crate::runtime::wgpu::ops::helpers::{
     CountParams, FlatToMultiParams, SearchsortedParams, SortParams, TopkParams, UniqueCountsParams,
     alloc_output, create_params_buffer, get_tensor_buffer, pack_u32_array,
 };
-use crate::runtime::wgpu::shaders::sort;
+use crate::runtime::wgpu::shaders::{sort, sort_global};
 use crate::runtime::{RuntimeClient, ensure_contiguous, normalize_dim};
 use crate::tensor::Tensor;
 use wgpu::{Buffer, BufferDescriptor, BufferUsages, MapMode, PollType};
@@ -41,18 +41,6 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
         let dim_idx = normalize_dim(dim, ndim)?;
         let sort_size = shape[dim_idx];
 
-        // Check sort size limit (WebGPU bitonic sort in shared memory)
-        if sort_size > MAX_SHARED_SORT_SIZE {
-            return Err(Error::backend_limitation(
-                "WebGPU",
-                "sort",
-                format!(
-                    "max {} elements per dimension, got {}",
-                    MAX_SHARED_SORT_SIZE, sort_size
-                ),
-            ));
-        }
-
         // Compute strides
         let outer_size: usize = shape[..dim_idx].iter().product();
         let inner_size: usize = shape[dim_idx + 1..].iter().product();
@@ -64,9 +52,29 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
 
         // Allocate output
         let out = alloc_output(self, shape, dtype);
+        if a.numel() == 0 {
+            return Ok(out);
+        }
+
         let a_buf = get_tensor_buffer(&a_contig)?;
         let out_buf = get_tensor_buffer(&out)?;
 
+        if sort_size > MAX_SHARED_SORT_SIZE {
+            sort_global::launch_global_sort(
+                self.pipeline_cache(),
+                self.wgpu_queue(),
+                &a_buf,
+                Some(&out_buf),
+                None,
+                outer_size,
+                sort_size,
+                inner_size,
+                descending,
+                dtype,
+            )?;
+            return Ok(out);
+        }
+
         // Create params buffer
         let params = SortParams {
             outer_size: outer_size as u32,
@@ -76,14 +84,6 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
         };
         let params_buf = create_params_buffer(self, &params);
 
-        // Create dummy indices buffer
-        let dummy_indices_buf = self.wgpu_device().create_buffer(&wgpu::BufferDescriptor {
-            label: Some("dummy_sort_indices"),
-            size: 4,
-            usage: wgpu::BufferUsages::STORAGE,
-            mapped_at_creation: false,
-        });
-
         sort::launch_sort_values_only(
             self.pipeline_cache(),
             self.wgpu_queue(),
@@ -95,7 +95,6 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
             dtype,
         )?;
 
-        drop(dummy_indices_buf);
         Ok(out)
     }
 
@@ -125,17 +124,6 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
         let dim_idx = normalize_dim(dim, ndim)?;
         let sort_size = shape[dim_idx];
 
-        if sort_size > MAX_SHARED_SORT_SIZE {
-            return Err(Error::backend_limitation(
-                "WebGPU",
-                "sort_with_indices",
-                format!(
-                    "max {} elements per dimension, got {}",
-                    MAX_SHARED_SORT_SIZE, sort_size
-                ),
-            ));
-        }
-
         let outer_size: usize = shape[..dim_idx].iter().product();
         let inner_size: usize = shape[dim_idx + 1..].iter().product();
         let outer_size = outer_size.max(1);
@@ -146,10 +134,30 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
         let values_out = alloc_output(self, shape, dtype);
         let indices_out = alloc_output(self, shape, DType::I32);
 
+        if a.numel() == 0 {
+            return Ok((values_out, indices_out));
+        }
+
         let a_buf = get_tensor_buffer(&a_contig)?;
         let values_buf = get_tensor_buffer(&values_out)?;
         let indices_buf = get_tensor_buffer(&indices_out)?;
 
+        if sort_size > MAX_SHARED_SORT_SIZE {
+            sort_global::launch_global_sort(
+                self.pipeline_cache(),
+                self.wgpu_queue(),
+                &a_buf,
+                Some(&values_buf),
+                Some(&indices_buf),
+                outer_size,
+                sort_size,
+                inner_size,
+                descending,
+                dtype,
+            )?;
+            return Ok((values_out, indices_out));
+        }
+
         let params = SortParams {
             outer_size: outer_size as u32,
             sort_size: sort_size as u32,
@@ -198,17 +206,6 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
         let dim_idx = normalize_dim(dim, ndim)?;
         let sort_size = shape[dim_idx];
 
-        if sort_size > MAX_SHARED_SORT_SIZE {
-            return Err(Error::backend_limitation(
-                "WebGPU",
-                "argsort",
-                format!(
-                    "max {} elements per dimension, got {}",
-                    MAX_SHARED_SORT_SIZE, sort_size
-                ),
-            ));
-        }
-
         let outer_size: usize = shape[..dim_idx].iter().product();
         let inner_size: usize = shape[dim_idx + 1..].iter().product();
         let outer_size = outer_size.max(1);
@@ -218,9 +215,29 @@ impl SortingOps<WgpuRuntime> for WgpuClient {
 
         let indices_out = alloc_output(self, shape, DType::I32);
 
+        if a.numel() == 0 {
+            return Ok(indices_out);
+        }
+
         let a_buf = get_tensor_buffer(&a_contig)?;
         let indices_buf = get_tensor_buffer(&indices_out)?;
 
+        if sort_size > MAX_SHARED_SORT_SIZE {
+            sort_global::launch_global_sort(
+                self.pipeline_cache(),
+                self.wgpu_queue(),
+                &a_buf,
+                None,
+                Some(&indices_buf),
+                outer_size,
+                sort_size,
+                inner_size,
+                descending,
+                dtype,
+            )?;
+            return Ok(indices_out);
+        }
+
         let params = SortParams {
             outer_size: outer_size as u32,
             sort_size: sort_size as u32,

src/runtime/wgpu/shaders/mod.rs

@@ -21,6 +21,7 @@ pub mod quasirandom;
 pub mod shape;
 pub mod sort;
 pub mod sort_cmp;
+pub(crate) mod sort_global;
 pub mod special;
 pub mod statistics;
 

src/runtime/wgpu/shaders/pipeline.rs

@@ -5,12 +5,13 @@
 
 use parking_lot::Mutex;
 use std::collections::HashMap;
+use std::num::NonZeroU64;
 use std::sync::Arc;
 use wgpu::{
     BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor,
-    BindGroupLayoutEntry, BindingType, Buffer, BufferBindingType, ComputePipeline,
-    ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue, ShaderModule,
-    ShaderModuleDescriptor, ShaderSource, ShaderStages,
+    BindGroupLayoutEntry, BindingResource, BindingType, Buffer, BufferBinding, BufferBindingType,
+    ComputePipeline, ComputePipelineDescriptor, Device, PipelineLayoutDescriptor, Queue,
+    ShaderModule, ShaderModuleDescriptor, ShaderSource, ShaderStages,
 };
 
 use crate::dtype::DType;
@@ -38,6 +39,8 @@ pub struct PipelineCache {
     dynamic_pipelines: Mutex<HashMap<(String, String), Arc<ComputePipeline>>>,
     /// Cached bind group layouts by layout key
     layouts: Mutex<HashMap<LayoutKey, Arc<BindGroupLayout>>>,
+    /// Layouts whose final uniform binding uses a dynamic offset.
+    dynamic_uniform_layouts: Mutex<HashMap<(u32, u32), Arc<BindGroupLayout>>>,
 }
 
 /// Key for bind group layout cache
@@ -65,6 +68,7 @@ impl PipelineCache {
             pipelines: Mutex::new(HashMap::new()),
             dynamic_pipelines: Mutex::new(HashMap::new()),
             layouts: Mutex::new(HashMap::new()),
+            dynamic_uniform_layouts: Mutex::new(HashMap::new()),
         }
     }
 
@@ -255,6 +259,87 @@ impl PipelineCache {
         })
     }
 
+    /// Get or create a storage-buffer layout followed by one dynamic uniform.
+    pub fn get_or_create_dynamic_uniform_layout(
+        &self,
+        num_storage_buffers: u32,
+        num_readonly_storage: u32,
+    ) -> Arc<BindGroupLayout> {
+        let key = (num_storage_buffers, num_readonly_storage);
+        let mut layouts = self.dynamic_uniform_layouts.lock();
+        if let Some(layout) = layouts.get(&key) {
+            return layout.clone();
+        }
+
+        let mut entries = Vec::with_capacity(num_storage_buffers as usize + 1);
+        for binding in 0..num_storage_buffers {
+            entries.push(BindGroupLayoutEntry {
+                binding,
+                visibility: ShaderStages::COMPUTE,
+                ty: BindingType::Buffer {
+                    ty: BufferBindingType::Storage {
+                        read_only: binding < num_readonly_storage,
+                    },
+                    has_dynamic_offset: false,
+                    min_binding_size: None,
+                },
+                count: None,
+            });
+        }
+        entries.push(BindGroupLayoutEntry {
+            binding: num_storage_buffers,
+            visibility: ShaderStages::COMPUTE,
+            ty: BindingType::Buffer {
+                ty: BufferBindingType::Uniform,
+                has_dynamic_offset: true,
+                min_binding_size: None,
+            },
+            count: None,
+        });
+
+        let layout = Arc::new(
+            self.device
+                .create_bind_group_layout(&BindGroupLayoutDescriptor {
+                    label: Some("dynamic_uniform_compute_layout"),
+                    entries: &entries,
+                }),
+        );
+        layouts.insert(key, layout.clone());
+        layout
+    }
+
+    /// Create a bind group whose final binding is a dynamically offset uniform.
+    pub fn create_bind_group_with_dynamic_uniform(
+        &self,
+        layout: &BindGroupLayout,
+        storage_buffers: &[&Buffer],
+        uniform: &Buffer,
+        uniform_binding_size: u64,
+    ) -> BindGroup {
+        let mut entries: Vec<BindGroupEntry<'_>> = storage_buffers
+            .iter()
+            .enumerate()
+            .map(|(binding, buffer)| BindGroupEntry {
+                binding: binding as u32,
+                resource: buffer.as_entire_binding(),
+            })
+            .collect();
+        entries.push(BindGroupEntry {
+            binding: storage_buffers.len() as u32,
+            resource: BindingResource::Buffer(BufferBinding {
+                buffer: uniform,
+                offset: 0,
+                size: NonZeroU64::new(uniform_binding_size),
+            }),
+        });
+
+        self.device.create_bind_group(&BindGroupDescriptor {
+            label: Some("dynamic_uniform_compute_bind_group"),
+            layout,
+            entries: &entries,
+        })
+    }
+
     /// Get device reference
     pub fn device(&self) -> &Device {
         &self.device

src/runtime/wgpu/shaders/sort.rs

@@ -26,7 +26,6 @@ const SORT_SHADER_F32: &str = concat!(
 );
 const SORT_SHADER_I32: &str = include_str!("sort_i32.wgsl");
 const SORT_SHADER_U32: &str = include_str!("sort_u32.wgsl");
-
 // ============================================================================
 // Static shaders — topk/searchsorted (F32 only)
 // ============================================================================