From e558acc752cacd6840fce84332ab7925561b9c19 Mon Sep 17 00:00:00 2001 From: SamJSui Date: Thu, 13 Aug 2026 18:14:23 -0500 Subject: [PATCH 1/2] feat(wgpu): add global stable sorting # Conflicts: # src/runtime/wgpu/shaders/sort_f32.wgsl # src/runtime/wgpu/shaders/sort_i32.wgsl # src/runtime/wgpu/shaders/sort_u32.wgsl --- examples/wgpu_global_sort_bench.rs | 64 +++++ src/ops/wgpu/sorting.rs | 91 +++---- src/runtime/wgpu/shaders/pipeline.rs | 91 ++++++- src/runtime/wgpu/shaders/sort.rs | 277 ++++++++++++++++++++++ src/runtime/wgpu/shaders/sort_global.wgsl | 245 +++++++++++++++++++ tests/backend_parity/sort.rs | 172 ++++++++++++++ 6 files changed, 894 insertions(+), 46 deletions(-) create mode 100644 examples/wgpu_global_sort_bench.rs create mode 100644 src/runtime/wgpu/shaders/sort_global.wgsl diff --git a/examples/wgpu_global_sort_bench.rs b/examples/wgpu_global_sort_bench.rs new file mode 100644 index 00000000..e1230cef --- /dev/null +++ b/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 = (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 = 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"); +} diff --git a/src/ops/wgpu/sorting.rs b/src/ops/wgpu/sorting.rs index 5c7e8857..fbb21797 100644 --- a/src/ops/wgpu/sorting.rs +++ b/src/ops/wgpu/sorting.rs @@ -41,18 +41,6 @@ impl SortingOps 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(); @@ -67,6 +55,22 @@ impl SortingOps for WgpuClient { let a_buf = get_tensor_buffer(&a_contig)?; let out_buf = get_tensor_buffer(&out)?; + if sort_size > MAX_SHARED_SORT_SIZE { + sort::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 +80,6 @@ impl SortingOps for WgpuClient { }; let params_buf = create_params_buffer(self, ¶ms); - // 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 +91,6 @@ impl SortingOps for WgpuClient { dtype, )?; - drop(dummy_indices_buf); Ok(out) } @@ -125,17 +120,6 @@ impl SortingOps 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); @@ -150,6 +134,22 @@ impl SortingOps for WgpuClient { let values_buf = get_tensor_buffer(&values_out)?; let indices_buf = get_tensor_buffer(&indices_out)?; + if sort_size > MAX_SHARED_SORT_SIZE { + sort::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 +198,6 @@ impl SortingOps 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); @@ -221,6 +210,22 @@ impl SortingOps for WgpuClient { let a_buf = get_tensor_buffer(&a_contig)?; let indices_buf = get_tensor_buffer(&indices_out)?; + if sort_size > MAX_SHARED_SORT_SIZE { + sort::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, diff --git a/src/runtime/wgpu/shaders/pipeline.rs b/src/runtime/wgpu/shaders/pipeline.rs index 5005807b..3050df4c 100644 --- a/src/runtime/wgpu/shaders/pipeline.rs +++ b/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>>, /// Cached bind group layouts by layout key layouts: Mutex>>, + /// Layouts whose final uniform binding uses a dynamic offset. + dynamic_uniform_layouts: Mutex>>, } /// 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 { + 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> = 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 diff --git a/src/runtime/wgpu/shaders/sort.rs b/src/runtime/wgpu/shaders/sort.rs index fd74f053..afebc73e 100644 --- a/src/runtime/wgpu/shaders/sort.rs +++ b/src/runtime/wgpu/shaders/sort.rs @@ -26,6 +26,24 @@ 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"); +const GLOBAL_SORT_SHADER: &str = include_str!("sort_global.wgsl"); + +#[repr(C)] +#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] +struct GlobalSortParams { + outer_size: u32, + sort_size: u32, + inner_size: u32, + padded_size: u32, + segment_count: u32, + dtype_tag: u32, + descending: u32, + k: u32, + j: u32, + total_padded: u32, + padding_0: u32, + padding_1: u32, +} // ============================================================================ // Static shaders — topk/searchsorted (F32 only) @@ -198,6 +216,265 @@ fn check_data_dtype(dtype: DType, op: &'static str) -> Result<()> { // Sort Operations // ============================================================================ +/// Launch the global-memory stable bitonic path used for sort dimensions above 512. +#[allow(clippy::too_many_arguments)] +pub fn launch_global_sort( + cache: &PipelineCache, + queue: &Queue, + input: &Buffer, + values_output: Option<&Buffer>, + indices_output: Option<&Buffer>, + outer_size: usize, + sort_size: usize, + inner_size: usize, + descending: bool, + dtype: DType, +) -> Result<()> { + let dtype_tag = match dtype { + DType::U32 => 0, + DType::I32 => 1, + DType::F32 => 2, + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "global_sort", + }); + } + }; + let padded_size = sort_size.checked_next_power_of_two().ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "sort dimension is too large") + })?; + let segment_count = outer_size.checked_mul(inner_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "segment count overflows usize") + })?; + let total_padded = segment_count.checked_mul(padded_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "global workspace size overflows usize") + })?; + let logical_total = segment_count.checked_mul(sort_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "output size overflows usize") + })?; + + let outer_size_u32 = u32::try_from(outer_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "outer dimension exceeds u32"))?; + let sort_size_u32 = u32::try_from(sort_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "sort dimension exceeds u32"))?; + let inner_size_u32 = u32::try_from(inner_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "inner dimension exceeds u32"))?; + let padded_size_u32 = u32::try_from(padded_size).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "padded sort dimension exceeds u32") + })?; + let segment_count_u32 = u32::try_from(segment_count) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "segment count exceeds u32"))?; + let total_padded_u32 = u32::try_from(total_padded).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "global workspace exceeds u32 elements") + })?; + let _logical_total_u32 = u32::try_from(logical_total) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "output exceeds u32 elements"))?; + + let scratch_bytes = (total_padded as u64).checked_mul(4).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "global workspace byte size overflows") + })?; + let limits = cache.device().limits(); + let binding_limit = limits.max_storage_buffer_binding_size; + let allocation_limit = limits.max_buffer_size; + let effective_limit = binding_limit.min(allocation_limit); + if scratch_bytes > effective_limit { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + format!( + "global workspace binding requires {scratch_bytes} bytes, device limit is {effective_limit}" + ), + )); + } + + let make_storage = |label: &'static str, size: u64| { + cache.device().create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size, + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }) + }; + let keys = make_storage("global_sort_keys", scratch_bytes); + let values = make_storage("global_sort_values", scratch_bytes); + let indices = make_storage("global_sort_indices", scratch_bytes); + let step_dummy = make_storage("global_sort_step_dummy", 4); + let logical_bytes = (logical_total as u64) * 4; + let temporary_values = values_output + .is_none() + .then(|| make_storage("global_sort_temporary_values_output", logical_bytes)); + let temporary_indices = indices_output + .is_none() + .then(|| make_storage("global_sort_temporary_indices_output", logical_bytes)); + let values_output = values_output + .or(temporary_values.as_ref()) + .expect("output exists"); + let indices_output = indices_output + .or(temporary_indices.as_ref()) + .expect("output exists"); + + let base_params = GlobalSortParams { + outer_size: outer_size_u32, + sort_size: sort_size_u32, + inner_size: inner_size_u32, + padded_size: padded_size_u32, + segment_count: segment_count_u32, + dtype_tag, + descending: u32::from(descending), + k: 0, + j: 0, + total_padded: total_padded_u32, + padding_0: 0, + padding_1: 0, + }; + let mut stage_params = vec![base_params]; + // k <= 512 is fused into one shared-memory tile dispatch. + let mut k = 1024u32; + while k <= padded_size_u32 { + let mut j = k >> 1; + while j > 0 { + stage_params.push(GlobalSortParams { + k, + j, + ..base_params + }); + j >>= 1; + } + k = k.checked_shl(1).unwrap_or(0); + if k == 0 { + break; + } + } + + let param_size = std::mem::size_of::(); + let alignment = limits.min_uniform_buffer_offset_alignment as usize; + let stride = param_size.div_ceil(alignment) * alignment; + let params_bytes_len = stride.checked_mul(stage_params.len()).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "parameter buffer size overflows") + })?; + let mut params_bytes = vec![0u8; params_bytes_len]; + for (stage_index, params) in stage_params.iter().enumerate() { + let bytes = bytemuck::bytes_of(params); + let offset = stage_index * stride; + params_bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + let params_buffer = cache.device().create_buffer(&wgpu::BufferDescriptor { + label: Some("global_sort_params"), + size: params_bytes_len as u64, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + queue.write_buffer(¶ms_buffer, 0, ¶ms_bytes); + + let module = cache.get_or_create_module("global_sort", GLOBAL_SORT_SHADER); + let layout = cache.get_or_create_dynamic_uniform_layout(4, 0); + let pack_pipeline = + cache.get_or_create_pipeline("global_sort_pack", "pack_global_sort", &module, &layout); + let tile_pipeline = + cache.get_or_create_pipeline("global_sort_tiles", "sort_global_tiles", &module, &layout); + let step_pipeline = + cache.get_or_create_pipeline("global_sort_step", "global_bitonic_step", &module, &layout); + let scatter_pipeline = cache.get_or_create_pipeline( + "global_sort_scatter", + "scatter_global_sort", + &module, + &layout, + ); + let uniform_binding_size = param_size as u64; + let pack_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[input, &keys, &values, &indices], + ¶ms_buffer, + uniform_binding_size, + ); + let step_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[&keys, &values, &indices, &step_dummy], + ¶ms_buffer, + uniform_binding_size, + ); + let scatter_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[&values, &indices, values_output, indices_output], + ¶ms_buffer, + uniform_binding_size, + ); + + let dispatch_grid = |items: usize| -> Result<(u32, u32)> { + let groups = items.div_ceil(256); + let x = groups.clamp(1, 65_535); + let y = groups.div_ceil(x); + if y > 65_535 { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + "global sort dispatch exceeds WebGPU's 2-D dispatch grid", + )); + } + Ok((x as u32, y as u32)) + }; + let (padded_x, padded_y) = dispatch_grid(total_padded)?; + let (logical_x, logical_y) = dispatch_grid(logical_total)?; + let tile_groups = total_padded / 512; + let tile_x = tile_groups.clamp(1, 65_535); + let tile_y = tile_groups.div_ceil(tile_x); + if tile_y > 65_535 { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + "global sort tile dispatch exceeds WebGPU's 2-D dispatch grid", + )); + } + let mut encoder = cache + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("global_sort"), + }); + + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_pack"), + timestamp_writes: None, + }); + pass.set_pipeline(&pack_pipeline); + pass.set_bind_group(0, Some(&pack_bind_group), &[0]); + pass.dispatch_workgroups(padded_x, padded_y, 1); + } + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_tiles"), + timestamp_writes: None, + }); + pass.set_pipeline(&tile_pipeline); + pass.set_bind_group(0, Some(&step_bind_group), &[0]); + pass.dispatch_workgroups(tile_x as u32, tile_y as u32, 1); + } + for stage_index in 1..stage_params.len() { + let dynamic_offset = u32::try_from(stage_index * stride).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "dynamic uniform offset exceeds u32") + })?; + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_step"), + timestamp_writes: None, + }); + pass.set_pipeline(&step_pipeline); + pass.set_bind_group(0, Some(&step_bind_group), &[dynamic_offset]); + pass.dispatch_workgroups(padded_x, padded_y, 1); + } + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_scatter"), + timestamp_writes: None, + }); + pass.set_pipeline(&scatter_pipeline); + pass.set_bind_group(0, Some(&scatter_bind_group), &[0]); + pass.dispatch_workgroups(logical_x, logical_y, 1); + } + queue.submit(std::iter::once(encoder.finish())); + Ok(()) +} + /// Launch sort with indices kernel pub fn launch_sort( cache: &PipelineCache, diff --git a/src/runtime/wgpu/shaders/sort_global.wgsl b/src/runtime/wgpu/shaders/sort_global.wgsl new file mode 100644 index 00000000..5c00f9ec --- /dev/null +++ b/src/runtime/wgpu/shaders/sort_global.wgsl @@ -0,0 +1,245 @@ +const WORKGROUP_SIZE: u32 = 256u; + +struct GlobalSortParams { + outer_size: u32, + sort_size: u32, + inner_size: u32, + padded_size: u32, + segment_count: u32, + dtype_tag: u32, + descending: u32, + k: u32, + j: u32, + total_padded: u32, + _padding_0: u32, + _padding_1: u32, +} + +@group(0) @binding(0) var buffer_0: array; +@group(0) @binding(1) var buffer_1: array; +@group(0) @binding(2) var buffer_2: array; +@group(0) @binding(3) var buffer_3: array; +@group(0) @binding(4) var params: GlobalSortParams; + +var tile_keys: array; +var tile_values: array; +var tile_indices: array; + +fn flat_invocation_index( + workgroup_id: vec3, + local_id: vec3, + num_workgroups: vec3, + total_items: u32, +) -> u32 { + let flat_group = workgroup_id.y * num_workgroups.x + workgroup_id.x; + let valid_groups = total_items / WORKGROUP_SIZE + + select(0u, 1u, total_items % WORKGROUP_SIZE != 0u); + if (flat_group >= valid_groups) { + return total_items; + } + return flat_group * WORKGROUP_SIZE + local_id.x; +} + +fn transformed_key(raw: u32) -> u32 { + if (params.dtype_tag == 0u) { + return raw; + } + if (params.dtype_tag == 1u) { + return raw ^ 0x80000000u; + } + + // Normalize signed zero for comparison while preserving the original bits + // in the value buffer. Canonicalize every NaN to the largest key. + let magnitude = raw & 0x7fffffffu; + if (magnitude == 0u) { + return 0x80000000u; + } + if (magnitude > 0x7f800000u) { + return 0xffffffffu; + } + return select(raw ^ 0x80000000u, ~raw, (raw & 0x80000000u) != 0u); +} + +@compute @workgroup_size(256) +fn pack_global_sort( + @builtin(workgroup_id) workgroup_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(num_workgroups) num_workgroups: vec3, +) { + let packed_index = flat_invocation_index( + workgroup_id, + local_id, + num_workgroups, + params.total_padded, + ); + if (packed_index >= params.total_padded) { + return; + } + + let segment = packed_index / params.padded_size; + let axis_index = packed_index % params.padded_size; + if (segment >= params.segment_count) { + return; + } + + if (axis_index < params.sort_size) { + let outer = segment / params.inner_size; + let inner = segment % params.inner_size; + let source_index = outer * params.sort_size * params.inner_size + + axis_index * params.inner_size + inner; + let raw = buffer_0[source_index]; + buffer_1[packed_index] = transformed_key(raw); + buffer_2[packed_index] = raw; + buffer_3[packed_index] = axis_index; + } else { + buffer_1[packed_index] = select(0xffffffffu, 0u, params.descending != 0u); + buffer_2[packed_index] = 0u; + buffer_3[packed_index] = axis_index; + } +} + +fn comes_before(key_a: u32, index_a: u32, key_b: u32, index_b: u32) -> bool { + if (key_a == key_b) { + return index_a < index_b; + } + if (params.descending != 0u) { + return key_a > key_b; + } + return key_a < key_b; +} + +fn tile_compare_and_swap(left: u32, right: u32, first_before: bool) { + let key_left = tile_keys[left]; + let key_right = tile_keys[right]; + let value_left = tile_values[left]; + let value_right = tile_values[right]; + let index_left = tile_indices[left]; + let index_right = tile_indices[right]; + let left_before_right = comes_before(key_left, index_left, key_right, index_right); + let swap = select(left_before_right, !left_before_right, first_before); + if (swap) { + tile_keys[left] = key_right; + tile_keys[right] = key_left; + tile_values[left] = value_right; + tile_values[right] = value_left; + tile_indices[left] = index_right; + tile_indices[right] = index_left; + } +} + +// Replaces the first 45 global bitonic stages (k <= 512) with one dispatch. +// Adjacent tiles alternate direction, exactly matching the k=512 network state. +@compute @workgroup_size(256) +fn sort_global_tiles( + @builtin(workgroup_id) workgroup_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(num_workgroups) num_workgroups: vec3, +) { + let tile = workgroup_id.y * num_workgroups.x + workgroup_id.x; + let tile_count = params.total_padded / 512u; + if (tile >= tile_count) { + return; + } + let tile_base = tile * 512u; + let first = local_id.x; + let second = first + 256u; + tile_keys[first] = buffer_0[tile_base + first]; + tile_keys[second] = buffer_0[tile_base + second]; + tile_values[first] = buffer_1[tile_base + first]; + tile_values[second] = buffer_1[tile_base + second]; + tile_indices[first] = buffer_2[tile_base + first]; + tile_indices[second] = buffer_2[tile_base + second]; + workgroupBarrier(); + + let tiles_per_segment = params.padded_size / 512u; + let tile_in_segment = tile % tiles_per_segment; + let tile_before = tile_in_segment % 2u == 0u; + for (var k = 2u; k <= 512u; k = k << 1u) { + for (var j = k >> 1u; j > 0u; j = j >> 1u) { + let pair_left = (local_id.x / j) * 2u * j + (local_id.x % j); + let pair_right = pair_left + j; + let stage_before = ((pair_left / k) % 2u == 0u) == tile_before; + tile_compare_and_swap(pair_left, pair_right, stage_before); + workgroupBarrier(); + } + } + + buffer_0[tile_base + first] = tile_keys[first]; + buffer_0[tile_base + second] = tile_keys[second]; + buffer_1[tile_base + first] = tile_values[first]; + buffer_1[tile_base + second] = tile_values[second]; + buffer_2[tile_base + first] = tile_indices[first]; + buffer_2[tile_base + second] = tile_indices[second]; +} + +@compute @workgroup_size(256) +fn global_bitonic_step( + @builtin(workgroup_id) workgroup_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(num_workgroups) num_workgroups: vec3, +) { + let packed_index = flat_invocation_index( + workgroup_id, + local_id, + num_workgroups, + params.total_padded, + ); + if (packed_index >= params.total_padded) { + return; + } + + let axis_index = packed_index % params.padded_size; + let partner_axis = axis_index ^ params.j; + if (partner_axis <= axis_index || partner_axis >= params.padded_size) { + return; + } + + let segment_base = packed_index - axis_index; + let partner_index = segment_base + partner_axis; + let key_a = buffer_0[packed_index]; + let key_b = buffer_0[partner_index]; + let value_a = buffer_1[packed_index]; + let value_b = buffer_1[partner_index]; + let index_a = buffer_2[packed_index]; + let index_b = buffer_2[partner_index]; + + let a_before_b = comes_before(key_a, index_a, key_b, index_b); + let first_half = (axis_index & params.k) == 0u; + let swap = select(a_before_b, !a_before_b, first_half); + if (swap) { + buffer_0[packed_index] = key_b; + buffer_0[partner_index] = key_a; + buffer_1[packed_index] = value_b; + buffer_1[partner_index] = value_a; + buffer_2[packed_index] = index_b; + buffer_2[partner_index] = index_a; + } +} + +@compute @workgroup_size(256) +fn scatter_global_sort( + @builtin(workgroup_id) workgroup_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(num_workgroups) num_workgroups: vec3, +) { + let logical_total = params.segment_count * params.sort_size; + let output_linear = flat_invocation_index( + workgroup_id, + local_id, + num_workgroups, + logical_total, + ); + if (output_linear >= logical_total) { + return; + } + + let segment = output_linear / params.sort_size; + let axis_index = output_linear % params.sort_size; + let packed_index = segment * params.padded_size + axis_index; + let outer = segment / params.inner_size; + let inner = segment % params.inner_size; + let output_index = outer * params.sort_size * params.inner_size + + axis_index * params.inner_size + inner; + buffer_2[output_index] = buffer_0[packed_index]; + buffer_3[output_index] = buffer_1[packed_index]; +} diff --git a/tests/backend_parity/sort.rs b/tests/backend_parity/sort.rs index cf0c63cc..37f41a6c 100644 --- a/tests/backend_parity/sort.rs +++ b/tests/backend_parity/sort.rs @@ -112,6 +112,178 @@ fn test_argsort_parity() { } } +#[cfg(feature = "wgpu")] +#[test] +fn test_wgpu_global_argsort_is_stable_past_shared_memory_limit() { + const LEN: usize = 4097; + let data: Vec = (0..LEN).map(|index| (index % 17) as f64).collect(); + let shape = vec![LEN]; + + with_wgpu_backend(|wgpu_client, wgpu_device| { + let tensor = tensor_from_f64( + &data, + &shape, + numr::dtype::DType::U32, + &wgpu_device, + &wgpu_client, + ) + .expect("create duplicate-heavy WGPU input"); + let indices: Vec = wgpu_client + .argsort(&tensor, 0, false) + .expect("global WGPU argsort") + .to_vec(); + + let mut expected: Vec = (0..LEN as i32).collect(); + expected.sort_by_key(|&index| (index as usize % 17, index)); + assert_eq!(indices, expected); + }); +} + +#[cfg(feature = "wgpu")] +#[test] +fn test_wgpu_global_sort_family_matches_cpu_on_arbitrary_axis() { + use numr::dtype::DType; + + let shape = vec![2, 513, 3]; + for dtype in [DType::U32, DType::I32, DType::F32] { + let data: Vec = (0..shape.iter().product()) + .map(|index| { + let value = ((index * 37 + index / 11) % 101) as f64; + if dtype == DType::U32 { + value + } else { + value - 50.0 + } + }) + .collect(); + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = tensor_from_f64(&data, &shape, dtype, &cpu_device, &cpu_client) + .expect("create CPU global-sort input"); + + for descending in [false, true] { + let cpu_values = cpu_client + .sort(&cpu_tensor, 1, descending) + .expect("CPU sort"); + let cpu_argsort: Vec = cpu_client + .argsort(&cpu_tensor, 1, descending) + .expect("CPU argsort") + .to_vec(); + let (cpu_values_with_indices, cpu_indices): (_, Vec) = { + let (values, indices) = cpu_client + .sort_with_indices(&cpu_tensor, 1, descending) + .expect("CPU sort_with_indices"); + (values, indices.to_vec()) + }; + + with_wgpu_backend(|wgpu_client, wgpu_device| { + let wgpu_tensor = tensor_from_f64(&data, &shape, dtype, &wgpu_device, &wgpu_client) + .expect("create WGPU global-sort input"); + let wgpu_values = wgpu_client + .sort(&wgpu_tensor, 1, descending) + .expect("WGPU global sort"); + let wgpu_argsort: Vec = wgpu_client + .argsort(&wgpu_tensor, 1, descending) + .expect("WGPU global argsort") + .to_vec(); + let (wgpu_values_with_indices, wgpu_indices) = wgpu_client + .sort_with_indices(&wgpu_tensor, 1, descending) + .expect("WGPU global sort_with_indices"); + let wgpu_indices: Vec = wgpu_indices.to_vec(); + + assert_tensor_allclose( + &wgpu_values, + &cpu_values, + dtype, + &format!("global sort WGPU vs CPU [{dtype:?}, descending={descending}]"), + ); + assert_tensor_allclose( + &wgpu_values_with_indices, + &cpu_values_with_indices, + dtype, + &format!( + "global sort_with_indices WGPU vs CPU [{dtype:?}, descending={descending}]" + ), + ); + assert_eq!( + wgpu_argsort + .iter() + .map(|&value| i64::from(value)) + .collect::>(), + cpu_argsort, + "global argsort indices [{dtype:?}, descending={descending}]" + ); + assert_eq!( + wgpu_indices + .iter() + .map(|&value| i64::from(value)) + .collect::>(), + cpu_indices, + "global sort_with_indices indices [{dtype:?}, descending={descending}]" + ); + }); + } + } +} + +#[cfg(feature = "wgpu")] +#[test] +fn test_wgpu_global_f32_orders_nans_and_stabilizes_signed_zero() { + use numr::tensor::Tensor; + + let mut data: Vec = (0..513).map(|index| (index % 23) as f32 - 11.0).collect(); + data[3] = f32::NAN; + data[200] = f32::from_bits(0xffc0_0001); + data[17] = -0.0; + data[41] = 0.0; + + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = Tensor::from_slice(&data, &[data.len()], &cpu_device); + with_wgpu_backend(|wgpu_client, wgpu_device| { + let wgpu_tensor = Tensor::from_slice(&data, &[data.len()], &wgpu_device); + for descending in [false, true] { + let expected: Vec = cpu_client + .argsort(&cpu_tensor, 0, descending) + .expect("CPU global f32 argsort") + .to_vec(); + let actual: Vec = wgpu_client + .argsort(&wgpu_tensor, 0, descending) + .expect("WGPU global f32 argsort") + .to_vec::() + .into_iter() + .map(i64::from) + .collect(); + assert_eq!( + actual, expected, + "f32 global argsort WGPU vs CPU descending={descending}" + ); + } + }); +} + +#[cfg(feature = "wgpu")] +#[test] +#[ignore = "large physical-GPU validation"] +fn test_wgpu_global_sort_one_million_elements() { + use numr::tensor::Tensor; + + const LEN: usize = 1_000_003; + let data: Vec = (0..LEN as u32).rev().collect(); + with_wgpu_backend(|wgpu_client, wgpu_device| { + let tensor = Tensor::from_slice(&data, &[LEN], &wgpu_device); + let sorted: Vec = wgpu_client + .sort(&tensor, 0, false) + .expect("one-million-element WGPU global sort") + .to_vec(); + assert_eq!(sorted.len(), LEN); + assert!( + sorted + .iter() + .enumerate() + .all(|(index, &value)| value == index as u32) + ); + }); +} + #[test] fn test_topk_parity() { let data = vec![3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0, 6.0]; From 1a84747d542c6885aea88b7ad1f6b8b6d7847453 Mon Sep 17 00:00:00 2001 From: SamJSui Date: Thu, 13 Aug 2026 19:51:05 -0500 Subject: [PATCH 2/2] fix(wgpu): address global sort review feedback --- src/ops/wgpu/sorting.rs | 20 +- src/runtime/wgpu/shaders/mod.rs | 1 + src/runtime/wgpu/shaders/sort.rs | 278 ----------------------- src/runtime/wgpu/shaders/sort_global.rs | 290 ++++++++++++++++++++++++ tests/backend_parity/sort.rs | 25 ++ 5 files changed, 332 insertions(+), 282 deletions(-) create mode 100644 src/runtime/wgpu/shaders/sort_global.rs diff --git a/src/ops/wgpu/sorting.rs b/src/ops/wgpu/sorting.rs index fbb21797..aa300a84 100644 --- a/src/ops/wgpu/sorting.rs +++ b/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}; @@ -52,11 +52,15 @@ impl SortingOps 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::launch_global_sort( + sort_global::launch_global_sort( self.pipeline_cache(), self.wgpu_queue(), &a_buf, @@ -130,12 +134,16 @@ impl SortingOps 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::launch_global_sort( + sort_global::launch_global_sort( self.pipeline_cache(), self.wgpu_queue(), &a_buf, @@ -207,11 +215,15 @@ impl SortingOps 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::launch_global_sort( + sort_global::launch_global_sort( self.pipeline_cache(), self.wgpu_queue(), &a_buf, diff --git a/src/runtime/wgpu/shaders/mod.rs b/src/runtime/wgpu/shaders/mod.rs index 56187f06..41b8ad7a 100644 --- a/src/runtime/wgpu/shaders/mod.rs +++ b/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; diff --git a/src/runtime/wgpu/shaders/sort.rs b/src/runtime/wgpu/shaders/sort.rs index afebc73e..ed63fb80 100644 --- a/src/runtime/wgpu/shaders/sort.rs +++ b/src/runtime/wgpu/shaders/sort.rs @@ -26,25 +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"); -const GLOBAL_SORT_SHADER: &str = include_str!("sort_global.wgsl"); - -#[repr(C)] -#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] -struct GlobalSortParams { - outer_size: u32, - sort_size: u32, - inner_size: u32, - padded_size: u32, - segment_count: u32, - dtype_tag: u32, - descending: u32, - k: u32, - j: u32, - total_padded: u32, - padding_0: u32, - padding_1: u32, -} - // ============================================================================ // Static shaders — topk/searchsorted (F32 only) // ============================================================================ @@ -216,265 +197,6 @@ fn check_data_dtype(dtype: DType, op: &'static str) -> Result<()> { // Sort Operations // ============================================================================ -/// Launch the global-memory stable bitonic path used for sort dimensions above 512. -#[allow(clippy::too_many_arguments)] -pub fn launch_global_sort( - cache: &PipelineCache, - queue: &Queue, - input: &Buffer, - values_output: Option<&Buffer>, - indices_output: Option<&Buffer>, - outer_size: usize, - sort_size: usize, - inner_size: usize, - descending: bool, - dtype: DType, -) -> Result<()> { - let dtype_tag = match dtype { - DType::U32 => 0, - DType::I32 => 1, - DType::F32 => 2, - _ => { - return Err(Error::UnsupportedDType { - dtype, - op: "global_sort", - }); - } - }; - let padded_size = sort_size.checked_next_power_of_two().ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "sort dimension is too large") - })?; - let segment_count = outer_size.checked_mul(inner_size).ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "segment count overflows usize") - })?; - let total_padded = segment_count.checked_mul(padded_size).ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "global workspace size overflows usize") - })?; - let logical_total = segment_count.checked_mul(sort_size).ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "output size overflows usize") - })?; - - let outer_size_u32 = u32::try_from(outer_size) - .map_err(|_| Error::backend_limitation("WebGPU", "sort", "outer dimension exceeds u32"))?; - let sort_size_u32 = u32::try_from(sort_size) - .map_err(|_| Error::backend_limitation("WebGPU", "sort", "sort dimension exceeds u32"))?; - let inner_size_u32 = u32::try_from(inner_size) - .map_err(|_| Error::backend_limitation("WebGPU", "sort", "inner dimension exceeds u32"))?; - let padded_size_u32 = u32::try_from(padded_size).map_err(|_| { - Error::backend_limitation("WebGPU", "sort", "padded sort dimension exceeds u32") - })?; - let segment_count_u32 = u32::try_from(segment_count) - .map_err(|_| Error::backend_limitation("WebGPU", "sort", "segment count exceeds u32"))?; - let total_padded_u32 = u32::try_from(total_padded).map_err(|_| { - Error::backend_limitation("WebGPU", "sort", "global workspace exceeds u32 elements") - })?; - let _logical_total_u32 = u32::try_from(logical_total) - .map_err(|_| Error::backend_limitation("WebGPU", "sort", "output exceeds u32 elements"))?; - - let scratch_bytes = (total_padded as u64).checked_mul(4).ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "global workspace byte size overflows") - })?; - let limits = cache.device().limits(); - let binding_limit = limits.max_storage_buffer_binding_size; - let allocation_limit = limits.max_buffer_size; - let effective_limit = binding_limit.min(allocation_limit); - if scratch_bytes > effective_limit { - return Err(Error::backend_limitation( - "WebGPU", - "sort", - format!( - "global workspace binding requires {scratch_bytes} bytes, device limit is {effective_limit}" - ), - )); - } - - let make_storage = |label: &'static str, size: u64| { - cache.device().create_buffer(&wgpu::BufferDescriptor { - label: Some(label), - size, - usage: wgpu::BufferUsages::STORAGE, - mapped_at_creation: false, - }) - }; - let keys = make_storage("global_sort_keys", scratch_bytes); - let values = make_storage("global_sort_values", scratch_bytes); - let indices = make_storage("global_sort_indices", scratch_bytes); - let step_dummy = make_storage("global_sort_step_dummy", 4); - let logical_bytes = (logical_total as u64) * 4; - let temporary_values = values_output - .is_none() - .then(|| make_storage("global_sort_temporary_values_output", logical_bytes)); - let temporary_indices = indices_output - .is_none() - .then(|| make_storage("global_sort_temporary_indices_output", logical_bytes)); - let values_output = values_output - .or(temporary_values.as_ref()) - .expect("output exists"); - let indices_output = indices_output - .or(temporary_indices.as_ref()) - .expect("output exists"); - - let base_params = GlobalSortParams { - outer_size: outer_size_u32, - sort_size: sort_size_u32, - inner_size: inner_size_u32, - padded_size: padded_size_u32, - segment_count: segment_count_u32, - dtype_tag, - descending: u32::from(descending), - k: 0, - j: 0, - total_padded: total_padded_u32, - padding_0: 0, - padding_1: 0, - }; - let mut stage_params = vec![base_params]; - // k <= 512 is fused into one shared-memory tile dispatch. - let mut k = 1024u32; - while k <= padded_size_u32 { - let mut j = k >> 1; - while j > 0 { - stage_params.push(GlobalSortParams { - k, - j, - ..base_params - }); - j >>= 1; - } - k = k.checked_shl(1).unwrap_or(0); - if k == 0 { - break; - } - } - - let param_size = std::mem::size_of::(); - let alignment = limits.min_uniform_buffer_offset_alignment as usize; - let stride = param_size.div_ceil(alignment) * alignment; - let params_bytes_len = stride.checked_mul(stage_params.len()).ok_or_else(|| { - Error::backend_limitation("WebGPU", "sort", "parameter buffer size overflows") - })?; - let mut params_bytes = vec![0u8; params_bytes_len]; - for (stage_index, params) in stage_params.iter().enumerate() { - let bytes = bytemuck::bytes_of(params); - let offset = stage_index * stride; - params_bytes[offset..offset + bytes.len()].copy_from_slice(bytes); - } - let params_buffer = cache.device().create_buffer(&wgpu::BufferDescriptor { - label: Some("global_sort_params"), - size: params_bytes_len as u64, - usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - queue.write_buffer(¶ms_buffer, 0, ¶ms_bytes); - - let module = cache.get_or_create_module("global_sort", GLOBAL_SORT_SHADER); - let layout = cache.get_or_create_dynamic_uniform_layout(4, 0); - let pack_pipeline = - cache.get_or_create_pipeline("global_sort_pack", "pack_global_sort", &module, &layout); - let tile_pipeline = - cache.get_or_create_pipeline("global_sort_tiles", "sort_global_tiles", &module, &layout); - let step_pipeline = - cache.get_or_create_pipeline("global_sort_step", "global_bitonic_step", &module, &layout); - let scatter_pipeline = cache.get_or_create_pipeline( - "global_sort_scatter", - "scatter_global_sort", - &module, - &layout, - ); - let uniform_binding_size = param_size as u64; - let pack_bind_group = cache.create_bind_group_with_dynamic_uniform( - &layout, - &[input, &keys, &values, &indices], - ¶ms_buffer, - uniform_binding_size, - ); - let step_bind_group = cache.create_bind_group_with_dynamic_uniform( - &layout, - &[&keys, &values, &indices, &step_dummy], - ¶ms_buffer, - uniform_binding_size, - ); - let scatter_bind_group = cache.create_bind_group_with_dynamic_uniform( - &layout, - &[&values, &indices, values_output, indices_output], - ¶ms_buffer, - uniform_binding_size, - ); - - let dispatch_grid = |items: usize| -> Result<(u32, u32)> { - let groups = items.div_ceil(256); - let x = groups.clamp(1, 65_535); - let y = groups.div_ceil(x); - if y > 65_535 { - return Err(Error::backend_limitation( - "WebGPU", - "sort", - "global sort dispatch exceeds WebGPU's 2-D dispatch grid", - )); - } - Ok((x as u32, y as u32)) - }; - let (padded_x, padded_y) = dispatch_grid(total_padded)?; - let (logical_x, logical_y) = dispatch_grid(logical_total)?; - let tile_groups = total_padded / 512; - let tile_x = tile_groups.clamp(1, 65_535); - let tile_y = tile_groups.div_ceil(tile_x); - if tile_y > 65_535 { - return Err(Error::backend_limitation( - "WebGPU", - "sort", - "global sort tile dispatch exceeds WebGPU's 2-D dispatch grid", - )); - } - let mut encoder = cache - .device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("global_sort"), - }); - - { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("global_sort_pack"), - timestamp_writes: None, - }); - pass.set_pipeline(&pack_pipeline); - pass.set_bind_group(0, Some(&pack_bind_group), &[0]); - pass.dispatch_workgroups(padded_x, padded_y, 1); - } - { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("global_sort_tiles"), - timestamp_writes: None, - }); - pass.set_pipeline(&tile_pipeline); - pass.set_bind_group(0, Some(&step_bind_group), &[0]); - pass.dispatch_workgroups(tile_x as u32, tile_y as u32, 1); - } - for stage_index in 1..stage_params.len() { - let dynamic_offset = u32::try_from(stage_index * stride).map_err(|_| { - Error::backend_limitation("WebGPU", "sort", "dynamic uniform offset exceeds u32") - })?; - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("global_sort_step"), - timestamp_writes: None, - }); - pass.set_pipeline(&step_pipeline); - pass.set_bind_group(0, Some(&step_bind_group), &[dynamic_offset]); - pass.dispatch_workgroups(padded_x, padded_y, 1); - } - { - let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("global_sort_scatter"), - timestamp_writes: None, - }); - pass.set_pipeline(&scatter_pipeline); - pass.set_bind_group(0, Some(&scatter_bind_group), &[0]); - pass.dispatch_workgroups(logical_x, logical_y, 1); - } - queue.submit(std::iter::once(encoder.finish())); - Ok(()) -} - /// Launch sort with indices kernel pub fn launch_sort( cache: &PipelineCache, diff --git a/src/runtime/wgpu/shaders/sort_global.rs b/src/runtime/wgpu/shaders/sort_global.rs new file mode 100644 index 00000000..5959c56a --- /dev/null +++ b/src/runtime/wgpu/shaders/sort_global.rs @@ -0,0 +1,290 @@ +//! Global-memory stable bitonic sort for dimensions above the shared-memory limit. + +use wgpu::{Buffer, Queue}; + +use super::pipeline::PipelineCache; +use crate::dtype::DType; +use crate::error::{Error, Result}; + +const GLOBAL_SORT_SHADER: &str = include_str!("sort_global.wgsl"); + +#[repr(C)] +#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] +struct GlobalSortParams { + outer_size: u32, + sort_size: u32, + inner_size: u32, + padded_size: u32, + segment_count: u32, + dtype_tag: u32, + descending: u32, + k: u32, + j: u32, + total_padded: u32, + padding_0: u32, + padding_1: u32, +} + +/// Launch the global-memory stable bitonic path used for sort dimensions above 512. +#[allow(clippy::too_many_arguments)] +pub fn launch_global_sort( + cache: &PipelineCache, + queue: &Queue, + input: &Buffer, + values_output: Option<&Buffer>, + indices_output: Option<&Buffer>, + outer_size: usize, + sort_size: usize, + inner_size: usize, + descending: bool, + dtype: DType, +) -> Result<()> { + let dtype_tag = match dtype { + DType::U32 => 0, + DType::I32 => 1, + DType::F32 => 2, + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "global_sort", + }); + } + }; + let padded_size = sort_size.checked_next_power_of_two().ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "sort dimension is too large") + })?; + let segment_count = outer_size.checked_mul(inner_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "segment count overflows usize") + })?; + let total_padded = segment_count.checked_mul(padded_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "global workspace size overflows usize") + })?; + let logical_total = segment_count.checked_mul(sort_size).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "output size overflows usize") + })?; + + let outer_size_u32 = u32::try_from(outer_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "outer dimension exceeds u32"))?; + let sort_size_u32 = u32::try_from(sort_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "sort dimension exceeds u32"))?; + let inner_size_u32 = u32::try_from(inner_size) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "inner dimension exceeds u32"))?; + let padded_size_u32 = u32::try_from(padded_size).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "padded sort dimension exceeds u32") + })?; + let segment_count_u32 = u32::try_from(segment_count) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "segment count exceeds u32"))?; + let total_padded_u32 = u32::try_from(total_padded).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "global workspace exceeds u32 elements") + })?; + let _logical_total_u32 = u32::try_from(logical_total) + .map_err(|_| Error::backend_limitation("WebGPU", "sort", "output exceeds u32 elements"))?; + + let scratch_bytes = (total_padded as u64).checked_mul(4).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "global workspace byte size overflows") + })?; + let limits = cache.device().limits(); + let binding_limit = limits.max_storage_buffer_binding_size; + let allocation_limit = limits.max_buffer_size; + let effective_limit = binding_limit.min(allocation_limit); + if scratch_bytes > effective_limit { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + format!( + "global workspace binding requires {scratch_bytes} bytes, device limit is {effective_limit}" + ), + )); + } + + let make_storage = |label: &'static str, size: u64| { + cache.device().create_buffer(&wgpu::BufferDescriptor { + label: Some(label), + size, + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }) + }; + let keys = make_storage("global_sort_keys", scratch_bytes); + let values = make_storage("global_sort_values", scratch_bytes); + let indices = make_storage("global_sort_indices", scratch_bytes); + let step_dummy = make_storage("global_sort_step_dummy", 4); + let logical_bytes = (logical_total as u64) * 4; + + let temporary_values; + let values_output = match values_output { + Some(output) => output, + None => { + temporary_values = make_storage("global_sort_temporary_values_output", logical_bytes); + &temporary_values + } + }; + let temporary_indices; + let indices_output = match indices_output { + Some(output) => output, + None => { + temporary_indices = make_storage("global_sort_temporary_indices_output", logical_bytes); + &temporary_indices + } + }; + + let base_params = GlobalSortParams { + outer_size: outer_size_u32, + sort_size: sort_size_u32, + inner_size: inner_size_u32, + padded_size: padded_size_u32, + segment_count: segment_count_u32, + dtype_tag, + descending: u32::from(descending), + k: 0, + j: 0, + total_padded: total_padded_u32, + padding_0: 0, + padding_1: 0, + }; + let mut stage_params = vec![base_params]; + // k <= 512 is fused into one shared-memory tile dispatch. + let mut k = 1024u32; + while k <= padded_size_u32 { + let mut j = k >> 1; + while j > 0 { + stage_params.push(GlobalSortParams { + k, + j, + ..base_params + }); + j >>= 1; + } + k = k.checked_shl(1).unwrap_or(0); + if k == 0 { + break; + } + } + + let param_size = std::mem::size_of::(); + let alignment = limits.min_uniform_buffer_offset_alignment as usize; + let stride = param_size.div_ceil(alignment) * alignment; + let params_bytes_len = stride.checked_mul(stage_params.len()).ok_or_else(|| { + Error::backend_limitation("WebGPU", "sort", "parameter buffer size overflows") + })?; + let mut params_bytes = vec![0u8; params_bytes_len]; + for (stage_index, params) in stage_params.iter().enumerate() { + let bytes = bytemuck::bytes_of(params); + let offset = stage_index * stride; + params_bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + let params_buffer = cache.device().create_buffer(&wgpu::BufferDescriptor { + label: Some("global_sort_params"), + size: params_bytes_len as u64, + usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + queue.write_buffer(¶ms_buffer, 0, ¶ms_bytes); + + let module = cache.get_or_create_module("global_sort", GLOBAL_SORT_SHADER); + let layout = cache.get_or_create_dynamic_uniform_layout(4, 0); + let pack_pipeline = + cache.get_or_create_pipeline("global_sort_pack", "pack_global_sort", &module, &layout); + let tile_pipeline = + cache.get_or_create_pipeline("global_sort_tiles", "sort_global_tiles", &module, &layout); + let step_pipeline = + cache.get_or_create_pipeline("global_sort_step", "global_bitonic_step", &module, &layout); + let scatter_pipeline = cache.get_or_create_pipeline( + "global_sort_scatter", + "scatter_global_sort", + &module, + &layout, + ); + let uniform_binding_size = param_size as u64; + let pack_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[input, &keys, &values, &indices], + ¶ms_buffer, + uniform_binding_size, + ); + let step_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[&keys, &values, &indices, &step_dummy], + ¶ms_buffer, + uniform_binding_size, + ); + let scatter_bind_group = cache.create_bind_group_with_dynamic_uniform( + &layout, + &[&values, &indices, values_output, indices_output], + ¶ms_buffer, + uniform_binding_size, + ); + + let dispatch_grid = |items: usize| -> Result<(u32, u32)> { + let groups = items.div_ceil(256); + let x = groups.clamp(1, 65_535); + let y = groups.div_ceil(x); + if y > 65_535 { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + "global sort dispatch exceeds WebGPU's 2-D dispatch grid", + )); + } + Ok((x as u32, y as u32)) + }; + let (padded_x, padded_y) = dispatch_grid(total_padded)?; + let (logical_x, logical_y) = dispatch_grid(logical_total)?; + let tile_groups = total_padded / 512; + let tile_x = tile_groups.clamp(1, 65_535); + let tile_y = tile_groups.div_ceil(tile_x); + if tile_y > 65_535 { + return Err(Error::backend_limitation( + "WebGPU", + "sort", + "global sort tile dispatch exceeds WebGPU's 2-D dispatch grid", + )); + } + let mut encoder = cache + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("global_sort"), + }); + + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_pack"), + timestamp_writes: None, + }); + pass.set_pipeline(&pack_pipeline); + pass.set_bind_group(0, Some(&pack_bind_group), &[0]); + pass.dispatch_workgroups(padded_x, padded_y, 1); + } + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_tiles"), + timestamp_writes: None, + }); + pass.set_pipeline(&tile_pipeline); + pass.set_bind_group(0, Some(&step_bind_group), &[0]); + pass.dispatch_workgroups(tile_x as u32, tile_y as u32, 1); + } + for stage_index in 1..stage_params.len() { + let dynamic_offset = u32::try_from(stage_index * stride).map_err(|_| { + Error::backend_limitation("WebGPU", "sort", "dynamic uniform offset exceeds u32") + })?; + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_step"), + timestamp_writes: None, + }); + pass.set_pipeline(&step_pipeline); + pass.set_bind_group(0, Some(&step_bind_group), &[dynamic_offset]); + pass.dispatch_workgroups(padded_x, padded_y, 1); + } + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("global_sort_scatter"), + timestamp_writes: None, + }); + pass.set_pipeline(&scatter_pipeline); + pass.set_bind_group(0, Some(&scatter_bind_group), &[0]); + pass.dispatch_workgroups(logical_x, logical_y, 1); + } + queue.submit(std::iter::once(encoder.finish())); + Ok(()) +} diff --git a/tests/backend_parity/sort.rs b/tests/backend_parity/sort.rs index 37f41a6c..6e59198c 100644 --- a/tests/backend_parity/sort.rs +++ b/tests/backend_parity/sort.rs @@ -260,6 +260,31 @@ fn test_wgpu_global_f32_orders_nans_and_stabilizes_signed_zero() { }); } +#[cfg(feature = "wgpu")] +#[test] +fn test_wgpu_global_sort_family_handles_empty_outer_dimension() { + use numr::dtype::DType; + use numr::tensor::Tensor; + + with_wgpu_backend(|wgpu_client, wgpu_device| { + let input = Tensor::zeros(&[0, 1024], DType::U32, &wgpu_device); + let sorted = wgpu_client + .sort(&input, 1, false) + .expect("sort empty tensor"); + let argsorted = wgpu_client + .argsort(&input, 1, false) + .expect("argsort empty tensor"); + let (values, indices) = wgpu_client + .sort_with_indices(&input, 1, false) + .expect("sort_with_indices empty tensor"); + + assert_eq!(sorted.shape(), &[0, 1024]); + assert_eq!(argsorted.shape(), &[0, 1024]); + assert_eq!(values.shape(), &[0, 1024]); + assert_eq!(indices.shape(), &[0, 1024]); + }); +} + #[cfg(feature = "wgpu")] #[test] #[ignore = "large physical-GPU validation"]