Support bounded multi-row GPU gray dispatch

CHANGELOG.md

@@ -21,6 +21,10 @@ removed no sooner than the next major (see `docs/API_STABILITY.md`).
 
 ### Changed
 
+- GPU RGB-to-gray supports images beyond one dispatch axis by distributing
+  flat 256-lane workgroups across bounded rows. Single-row work retains its
+  original shader; dimensions, counters, device limits and profiler capacity
+  reject before allocation, with unchanged luma bytes and explicit transfers.
 - FPFH point-cloud descriptors share one KD-tree across bounded scoped CPU
   workers for both SPFH construction and neighbor weighting. Small/one-worker
   clouds stay serial; output order and exact descriptor words are preserved

crates/spatialrust-gpu/src/image/kernels/gray.rs

@@ -3,6 +3,7 @@
 use bytemuck::{Pod, Zeroable};
 use spatialrust_core::{SpatialError, SpatialResult};
 use spatialrust_image::{ColorSpace, ImageMetadata};
+use std::sync::OnceLock;
 use wgpu::util::DeviceExt;
 
 use crate::image::gpu_image::{create_texture, GpuImage, GpuImageReceipt};
@@ -15,8 +16,53 @@ const WORKGROUP: u32 = 256;
 struct GrayParams {
     width: u32,
     height: u32,
-    _pad0: u32,
-    _pad1: u32,
+    groups_per_row: u32,
+    group_count: u32,
+}
+
+#[derive(Clone, Copy, Debug)]
+struct GrayPlan {
+    groups_x: u32,
+    groups_y: u32,
+    group_count: u32,
+}
+
+fn gray_plan(width: u32, height: u32, limits: &wgpu::Limits) -> SpatialResult<GrayPlan> {
+    let invalid = || {
+        SpatialError::InvalidArgument(
+        "GPU gray requires positive dimensions, u32 pixel counts and supported texture/compute/binding limits".to_owned(),
+    )
+    };
+    let pixels = u64::from(width) * u64::from(height);
+    if width == 0
+        || height == 0
+        || width > limits.max_texture_dimension_2d
+        || height > limits.max_texture_dimension_2d
+        || width > i32::MAX as u32
+        || height > i32::MAX as u32
+        || pixels > u64::from(u32::MAX)
+        || limits.max_compute_invocations_per_workgroup < WORKGROUP
+        || limits.max_compute_workgroup_size_x < WORKGROUP
+        || limits.max_compute_workgroup_size_y == 0
+        || limits.max_compute_workgroup_size_z == 0
+        || limits.max_compute_workgroups_per_dimension == 0
+        || limits.max_sampled_textures_per_shader_stage == 0
+        || limits.max_storage_textures_per_shader_stage == 0
+        || limits.max_uniform_buffers_per_shader_stage == 0
+        || limits.max_uniform_buffer_binding_size < std::mem::size_of::<GrayParams>() as u32
+        || limits.max_buffer_size < std::mem::size_of::<GrayParams>() as u64
+        || limits.max_bind_groups == 0
+        || limits.max_bindings_per_bind_group < 3
+    {
+        return Err(invalid());
+    }
+    let group_count = pixels.div_ceil(u64::from(WORKGROUP)) as u32;
+    let groups_y = group_count.div_ceil(limits.max_compute_workgroups_per_dimension);
+    if groups_y > limits.max_compute_workgroups_per_dimension {
+        return Err(invalid());
+    }
+    let groups_x = group_count.div_ceil(groups_y);
+    Ok(GrayPlan { groups_x, groups_y, group_count })
 }
 
 const GRAY_WGSL: &str = r#"
@@ -43,7 +89,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 pub(crate) struct GrayPipeline {
     bind_group_layout: wgpu::BindGroupLayout,
-    pipeline: wgpu::ComputePipeline,
+    pipeline: OnceLock<wgpu::ComputePipeline>,
+    grid_pipeline: OnceLock<wgpu::ComputePipeline>,
 }
 
 pub(crate) fn create_gray_pipeline(device: &wgpu::Device) -> GrayPipeline {
@@ -82,58 +129,99 @@ pub(crate) fn create_gray_pipeline(device: &wgpu::Device) -> GrayPipeline {
             },
         ],
     });
+    GrayPipeline { bind_group_layout, pipeline: OnceLock::new(), grid_pipeline: OnceLock::new() }
+}
+
+fn create_pipeline(
+    device: &wgpu::Device,
+    bind_group_layout: &wgpu::BindGroupLayout,
+    source: &str,
+) -> wgpu::ComputePipeline {
     let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
         label: Some("gpu-image-gray-shader"),
-        source: wgpu::ShaderSource::Wgsl(GRAY_WGSL.into()),
+        source: wgpu::ShaderSource::Wgsl(source.into()),
     });
     let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
         label: Some("gpu-image-gray-pl"),
-        bind_group_layouts: &[&bind_group_layout],
+        bind_group_layouts: &[bind_group_layout],
         push_constant_ranges: &[],
     });
-    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+    device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
         label: Some("gpu-image-gray-pipeline"),
         layout: Some(&pipeline_layout),
         module: &shader,
         entry_point: Some("main"),
         compilation_options: wgpu::PipelineCompilationOptions::default(),
         cache: None,
-    });
-    GrayPipeline { bind_group_layout, pipeline }
+    })
 }
 
 /// Converts an RGB `GpuImage` to a gray `GpuImage` using BT.601 fixed-point luma.
+///
+/// Large images distribute flat workgroups across bounded dispatch rows. Runtime,
+/// dimensions and device limits are checked before allocation or pipeline setup.
 pub fn rgb_to_gray_gpu(runtime: &WgpuRuntime, source: &GpuImage) -> SpatialResult<GpuImage> {
     rgb_to_gray_gpu_with_profiler(runtime, source, None)
 }
 
 pub(crate) fn rgb_to_gray_gpu_with_profiler(
     runtime: &WgpuRuntime,
     source: &GpuImage,
-    mut profiler: Option<&mut crate::GpuStageProfiler>,
+    profiler: Option<&mut crate::GpuStageProfiler>,
 ) -> SpatialResult<GpuImage> {
     source.validate_runtime(runtime)?;
     if source.channels() != 3 {
         return Err(SpatialError::InvalidArgument(
             "rgb_to_gray_gpu requires a 3-channel RGB GpuImage".to_owned(),
         ));
     }
-    let pixel_count = (source.width() as usize).saturating_mul(source.height() as usize);
+    let plan = gray_plan(source.width(), source.height(), &runtime.device().limits())?;
+    if let Some(profiler) = &profiler {
+        profiler.validate(runtime, 1)?;
+    }
+    let cache = runtime.image_gray_pipeline.get_or_init(|| create_gray_pipeline(runtime.device()));
+    let pipeline = if plan.groups_y == 1 {
+        cache
+            .pipeline
+            .get_or_init(|| create_pipeline(runtime.device(), &cache.bind_group_layout, GRAY_WGSL))
+    } else {
+        cache.grid_pipeline.get_or_init(|| {
+            create_pipeline(
+                runtime.device(),
+                &cache.bind_group_layout,
+                include_str!("gray_grid.wgsl"),
+            )
+        })
+    };
+    gray_dispatch(runtime, source, plan, cache, pipeline, profiler)
+}
+
+fn gray_dispatch(
+    runtime: &WgpuRuntime,
+    source: &GpuImage,
+    plan: GrayPlan,
+    cache: &GrayPipeline,
+    pipeline: &wgpu::ComputePipeline,
+    mut profiler: Option<&mut crate::GpuStageProfiler>,
+) -> SpatialResult<GpuImage> {
     let output = create_texture(runtime, source.width(), source.height(), "gpu-image-gray-out");
     let out_bytes = u64::from(source.width()) * u64::from(source.height()) * 4;
-    let params = GrayParams { width: source.width(), height: source.height(), _pad0: 0, _pad1: 0 };
+    let params = GrayParams {
+        width: source.width(),
+        height: source.height(),
+        groups_per_row: plan.groups_x,
+        group_count: plan.group_count,
+    };
     let uniform = runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor {
         label: Some("gpu-image-gray-params"),
         contents: bytemuck::bytes_of(&params),
         usage: wgpu::BufferUsages::UNIFORM,
     });
-    let pipeline =
-        runtime.image_gray_pipeline.get_or_init(|| create_gray_pipeline(runtime.device()));
     let source_view = source.view();
     let output_view = output.create_view(&wgpu::TextureViewDescriptor::default());
     let bind_group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
         label: Some("gpu-image-gray-bg"),
-        layout: &pipeline.bind_group_layout,
+        layout: &cache.bind_group_layout,
         entries: &[
             wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
             wgpu::BindGroupEntry {
@@ -154,9 +242,9 @@ pub(crate) fn rgb_to_gray_gpu_with_profiler(
             label: Some("gpu-image-gray-pass"),
             timestamp_writes: crate::profiling::timestamp_writes(&mut profiler, "rgb_to_gray_gpu"),
         });
-        pass.set_pipeline(&pipeline.pipeline);
+        pass.set_pipeline(pipeline);
         pass.set_bind_group(0, &bind_group, &[]);
-        pass.dispatch_workgroups(pixel_count.div_ceil(WORKGROUP as usize) as u32, 1, 1);
+        pass.dispatch_workgroups(plan.groups_x, plan.groups_y, 1);
     }
     runtime.queue().submit(Some(encoder.finish()));
     let mut receipt = GpuImageReceipt::default();
@@ -165,3 +253,7 @@ pub(crate) fn rgb_to_gray_gpu_with_profiler(
     let metadata = ImageMetadata { color_space: ColorSpace::Gray, ..source.metadata() };
     GpuImage::from_parts(runtime, source.width(), source.height(), 1, output, metadata, receipt)
 }
+
+#[cfg(all(test, not(target_arch = "wasm32")))]
+#[path = "gray_tests.rs"]
+mod tests;

crates/spatialrust-gpu/src/image/kernels/gray_benchmark.rs

@@ -0,0 +1,64 @@
+use super::*;
+
+#[test]
+#[ignore = "explicit native resident GPU gray diagnostic"]
+fn paired_frozen_current_and_large_gray() {
+    let runtime =
+        WgpuRuntime::new_headless_with_timestamps(crate::WgpuPowerPreference::HighPerformance)
+            .unwrap();
+    let cache = runtime.image_gray_pipeline.get_or_init(|| create_gray_pipeline(runtime.device()));
+    let frozen = create_pipeline(
+        runtime.device(),
+        &cache.bind_group_layout,
+        include_str!("../../../tests/fixtures/gray_previous.wgsl"),
+    );
+    eprintln!("backend={}; device_type={}; 2 warmups + 20 alternating rounds; shared host executor; timings include validation/allocation/encoding/submit/wait, exclude upload/pipeline setup/profiler creation/query collection/readback/checks/recycle",runtime.adapter_info().backend,runtime.adapter_info().device_type);
+    for (w, h) in
+        [(17, 19), (128, 128), (640, 480), (1920, 1080), (3840, 2160), (4096, 4095), (4096, 4096)]
+    {
+        let image = fixture(w, h, 0);
+        let expected = oracle(image.as_slice());
+        let source = GpuImage::upload_u8(&runtime, image.view()).unwrap();
+        let plan = gray_plan(w as u32, h as u32, &runtime.device().limits()).unwrap();
+        // The frozen linear kernel has no valid dispatch for the large case.
+        let modes = if plan.groups_y == 1 { 2 } else { 1 };
+        let mut stage: [Vec<f64>; 2] = std::array::from_fn(|_| Vec::new());
+        let mut wall: [Vec<f64>; 2] = std::array::from_fn(|_| Vec::new());
+        for round in 0..22 {
+            for offset in 0..modes {
+                let mode = (round + offset) % modes;
+                let current = modes == 1 || mode == 1;
+                let mut profiler = crate::GpuStageProfiler::new(&runtime, 1).unwrap();
+                let start = std::time::Instant::now();
+                let output = if current {
+                    rgb_to_gray_gpu_with_profiler(&runtime, &source, Some(&mut profiler)).unwrap()
+                } else {
+                    source.validate_runtime(&runtime).unwrap();
+                    gray_plan(w as u32, h as u32, &runtime.device().limits()).unwrap();
+                    profiler.validate(&runtime, 1).unwrap();
+                    gray_dispatch(&runtime, &source, plan, cache, &frozen, Some(&mut profiler))
+                        .unwrap()
+                };
+                runtime.wait_idle();
+                let ms = start.elapsed().as_secs_f64() * 1000.0;
+                let report = profiler.collect(&runtime).unwrap();
+                assert_eq!(report.stages.len(), 1);
+                assert_eq!(report.device_to_host_bytes, 16);
+                assert_eq!(output.receipt().stages(), ["upload_u8_texture", "rgb_to_gray_gpu"]);
+                assert_eq!(output.receipt().gpu_to_gpu_bytes(), (w * h * 4) as u64);
+                assert_eq!(output.receipt().device_to_host_bytes(), 0);
+                assert_rgba(&runtime, &output, &expected);
+                assert_eq!(output.metadata().color_space, ColorSpace::Gray);
+                assert_eq!(output.metadata().color_range, image.metadata().color_range);
+                if round >= 2 {
+                    let i = usize::from(current);
+                    stage[i].push(report.stages[0].elapsed_ns / 1000.0);
+                    wall[i].push(ms);
+                }
+                output.recycle(&runtime);
+            }
+        }
+        eprintln!("GRAY {{\"width\":{w},\"height\":{h},\"groups_x\":{},\"groups_y\":{},\"previous_us\":{:?},\"current_us\":{:?},\"previous_ms\":{:?},\"current_ms\":{:?}}}",plan.groups_x,plan.groups_y,stage[0],stage[1],wall[0],wall[1]);
+        source.recycle(&runtime);
+    }
+}

crates/spatialrust-gpu/src/image/kernels/gray_grid.wgsl

@@ -0,0 +1,29 @@
+struct Params {
+    width: u32,
+    height: u32,
+    groups_per_row: u32,
+    group_count: u32,
+};
+
+@group(0) @binding(0) var<uniform> params: Params;
+@group(0) @binding(1) var input_px: texture_2d<u32>;
+@group(0) @binding(2) var output_px: texture_storage_2d<rgba8uint, write>;
+
+fn flat_pixel(group: vec3<u32>, lane: u32) -> u32 {
+    let flat_group = group.y * params.groups_per_row + group.x;
+    // A balanced rectangle can contain padding groups. Reject them before
+    // multiplying by 256 so the u32::MAX pixel boundary cannot wrap.
+    if (flat_group >= params.group_count) { return 0xffffffffu; }
+    return flat_group * 256u + lane;
+}
+
+@compute @workgroup_size(256)
+fn main(@builtin(workgroup_id) group: vec3<u32>, @builtin(local_invocation_index) lane: u32) {
+    let index = flat_pixel(group, lane);
+    let pixel_count = params.width * params.height;
+    if (index >= pixel_count) { return; }
+    let xy = vec2<i32>(i32(index % params.width), i32(index / params.width));
+    let rgb = textureLoad(input_px, xy, 0).rgb;
+    let gray = (77u * rgb.r + 150u * rgb.g + 29u * rgb.b + 128u) >> 8u;
+    textureStore(output_px, xy, vec4<u32>(gray, 0u, 0u, 0u));
+}