spatialrust_gpu/image/
ai_tensor.rs

1//! Device-resident planar floating-point tensors packed from GPU images.
2
3use bytemuck::{Pod, Zeroable};
4use spatialrust_core::{SpatialError, SpatialResult};
5
6use super::gpu_image::{read_staging_bytes, runtime_device_key, GpuImage, GpuImageReceipt};
7use crate::WgpuRuntime;
8
9const WORKGROUP_X: u32 = 16;
10const WORKGROUP_Y: u32 = 16;
11
12const SHADER: &str = r#"
13struct Params {
14    width: u32,
15    height: u32,
16    channels: u32,
17    _pad: u32,
18    scale: f32,
19    mean0: f32,
20    mean1: f32,
21    mean2: f32,
22    mean3: f32,
23    std0: f32,
24    std1: f32,
25    std2: f32,
26    std3: f32,
27    _tail0: f32,
28    _tail1: f32,
29    _tail2: f32,
30}
31@group(0) @binding(0) var<uniform> params: Params;
32@group(0) @binding(1) var source_px: texture_2d<u32>;
33@group(0) @binding(2) var<storage, read_write> output: array<f32>;
34
35@compute @workgroup_size(16, 16)
36fn pack_ai_chw(@builtin(global_invocation_id) gid: vec3<u32>) {
37    if (gid.x >= params.width || gid.y >= params.height) { return; }
38    let pixel = textureLoad(source_px, vec2<i32>(gid.xy), 0);
39    let values = vec4<f32>(pixel);
40    let means = vec4<f32>(params.mean0, params.mean1, params.mean2, params.mean3);
41    let stds = vec4<f32>(params.std0, params.std1, params.std2, params.std3);
42    let index = gid.y * params.width + gid.x;
43    let plane = params.width * params.height;
44    for (var channel = 0u; channel < params.channels; channel++) {
45        output[channel * plane + index] = (values[channel] * params.scale - means[channel]) / stds[channel];
46    }
47}
48"#;
49
50#[repr(C)]
51#[derive(Clone, Copy, Pod, Zeroable)]
52struct AiPackParams {
53    width: u32,
54    height: u32,
55    channels: u32,
56    _pad: u32,
57    scale: f32,
58    mean: [f32; 4],
59    std: [f32; 4],
60    _tail: [f32; 3],
61}
62
63pub(crate) struct AiPackPipeline {
64    layout: wgpu::BindGroupLayout,
65    pipeline: wgpu::ComputePipeline,
66}
67
68pub(crate) fn create_ai_pack_pipeline(device: &wgpu::Device) -> AiPackPipeline {
69    let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
70        label: Some("gpu-ai-pack-bgl"),
71        entries: &[
72            wgpu::BindGroupLayoutEntry {
73                binding: 0,
74                visibility: wgpu::ShaderStages::COMPUTE,
75                ty: wgpu::BindingType::Buffer {
76                    ty: wgpu::BufferBindingType::Uniform,
77                    has_dynamic_offset: false,
78                    min_binding_size: None,
79                },
80                count: None,
81            },
82            wgpu::BindGroupLayoutEntry {
83                binding: 1,
84                visibility: wgpu::ShaderStages::COMPUTE,
85                ty: wgpu::BindingType::Texture {
86                    sample_type: wgpu::TextureSampleType::Uint,
87                    view_dimension: wgpu::TextureViewDimension::D2,
88                    multisampled: false,
89                },
90                count: None,
91            },
92            wgpu::BindGroupLayoutEntry {
93                binding: 2,
94                visibility: wgpu::ShaderStages::COMPUTE,
95                ty: wgpu::BindingType::Buffer {
96                    ty: wgpu::BufferBindingType::Storage { read_only: false },
97                    has_dynamic_offset: false,
98                    min_binding_size: None,
99                },
100                count: None,
101            },
102        ],
103    });
104    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
105        label: Some("gpu-ai-pack-shader"),
106        source: wgpu::ShaderSource::Wgsl(SHADER.into()),
107    });
108    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
109        label: Some("gpu-ai-pack-layout"),
110        bind_group_layouts: &[&layout],
111        push_constant_ranges: &[],
112    });
113    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
114        label: Some("gpu-ai-pack-pipeline"),
115        layout: Some(&pipeline_layout),
116        module: &shader,
117        entry_point: Some("pack_ai_chw"),
118        compilation_options: wgpu::PipelineCompilationOptions::default(),
119        cache: None,
120    });
121    AiPackPipeline { layout, pipeline }
122}
123
124/// Device-resident planar `f32` tensor with explicit optional readback.
125pub struct GpuAiTensor {
126    buffer: wgpu::Buffer,
127    width: u32,
128    height: u32,
129    channels: u32,
130    byte_len: u64,
131    device_key: usize,
132    receipt: GpuImageReceipt,
133}
134
135impl GpuAiTensor {
136    /// Returns tensor dimensions as CHW.
137    #[must_use]
138    pub const fn shape(&self) -> [u32; 3] {
139        [self.channels, self.height, self.width]
140    }
141
142    /// Returns cumulative transfer and stage accounting.
143    #[must_use]
144    pub const fn receipt(&self) -> &GpuImageReceipt {
145        &self.receipt
146    }
147
148    /// Explicitly reads planar values back to host memory.
149    pub fn readback_f32(&mut self, runtime: &WgpuRuntime) -> SpatialResult<Vec<f32>> {
150        if self.device_key != runtime_device_key(runtime) {
151            return Err(SpatialError::InvalidArgument(
152                "GpuAiTensor belongs to a different runtime device".to_owned(),
153            ));
154        }
155        let staging = runtime.device().create_buffer(&wgpu::BufferDescriptor {
156            label: Some("gpu-ai-tensor-readback"),
157            size: self.byte_len,
158            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
159            mapped_at_creation: false,
160        });
161        let mut encoder =
162            runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
163                label: Some("gpu-ai-tensor-readback-encoder"),
164            });
165        encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging, 0, self.byte_len);
166        runtime.queue().submit(Some(encoder.finish()));
167        let bytes = read_staging_bytes(runtime.device(), &staging, self.byte_len as usize)?;
168        self.receipt.record_device_to_host(self.byte_len, "readback_ai_chw_f32");
169        Ok(bytes
170            .chunks_exact(std::mem::size_of::<f32>())
171            .map(|chunk| f32::from_ne_bytes(chunk.try_into().expect("four-byte f32 chunk")))
172            .collect())
173    }
174
175    /// Returns tensor storage to the runtime pool.
176    pub fn recycle(self, runtime: &WgpuRuntime) {
177        if self.device_key == runtime_device_key(runtime) {
178            runtime.recycle_storage(self.byte_len, self.buffer);
179        }
180    }
181}
182
183/// Packs a resident image into normalized planar CHW `f32` storage.
184pub fn pack_ai_chw_gpu(
185    runtime: &WgpuRuntime,
186    source: &GpuImage,
187    scale: f32,
188    mean: [f32; 4],
189    std: [f32; 4],
190) -> SpatialResult<GpuAiTensor> {
191    source.validate_runtime(runtime)?;
192    if !scale.is_finite()
193        || mean.iter().any(|value| !value.is_finite())
194        || std.iter().any(|value| !value.is_finite() || *value == 0.0)
195    {
196        return Err(SpatialError::InvalidArgument(
197            "AI packing scale/mean/std must be finite and std non-zero".to_owned(),
198        ));
199    }
200    let elements =
201        u64::from(source.width()) * u64::from(source.height()) * u64::from(source.channels());
202    let byte_len = elements * std::mem::size_of::<f32>() as u64;
203    let output = runtime.buffer_pool().acquire_storage(runtime, "gpu-ai-chw-output", byte_len);
204    let params = AiPackParams {
205        width: source.width(),
206        height: source.height(),
207        channels: source.channels(),
208        _pad: 0,
209        scale,
210        mean,
211        std,
212        _tail: [0.0; 3],
213    };
214    let uniform = runtime.device().create_buffer(&wgpu::BufferDescriptor {
215        label: Some("gpu-ai-pack-params"),
216        size: std::mem::size_of::<AiPackParams>() as u64,
217        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
218        mapped_at_creation: false,
219    });
220    runtime.queue().write_buffer(&uniform, 0, bytemuck::bytes_of(&params));
221    let pipeline =
222        runtime.image_ai_pack_pipeline.get_or_init(|| create_ai_pack_pipeline(runtime.device()));
223    let source_view = source.view();
224    let bind_group = runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
225        label: Some("gpu-ai-pack-bg"),
226        layout: &pipeline.layout,
227        entries: &[
228            wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
229            wgpu::BindGroupEntry {
230                binding: 1,
231                resource: wgpu::BindingResource::TextureView(&source_view),
232            },
233            wgpu::BindGroupEntry { binding: 2, resource: output.as_entire_binding() },
234        ],
235    });
236    let mut encoder = runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
237        label: Some("gpu-ai-pack-encoder"),
238    });
239    {
240        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
241            label: Some("pack_ai_chw_gpu"),
242            timestamp_writes: None,
243        });
244        pass.set_pipeline(&pipeline.pipeline);
245        pass.set_bind_group(0, &bind_group, &[]);
246        pass.dispatch_workgroups(
247            source.width().div_ceil(WORKGROUP_X),
248            source.height().div_ceil(WORKGROUP_Y),
249            1,
250        );
251    }
252    runtime.queue().submit(Some(encoder.finish()));
253    let mut receipt = source.receipt().clone();
254    receipt.record_gpu_to_gpu(byte_len, "pack_ai_chw_gpu");
255    Ok(GpuAiTensor {
256        buffer: output,
257        width: source.width(),
258        height: source.height(),
259        channels: source.channels(),
260        byte_len,
261        device_key: runtime_device_key(runtime),
262        receipt,
263    })
264}
spatialrust_gpu/image/ai_tensor.rs

spatialrust_gpu/image/
ai_tensor.rs
