| use bytemuck::{Pod, Zeroable}; |
| use std::borrow::Cow; |
| use wgpu::util::DeviceExt; |
|
|
| #[repr(C)] |
| #[derive(Copy, Clone, Pod, Zeroable)] |
| struct Params { |
| n: u32, c: u32, din: u32, hin: u32, win: u32, |
| scale: u32, dout: u32, hout: u32, wout: u32, |
| _p0: u32, _p1: u32, _p2: u32, |
| } |
|
|
| const WGSL: &str = r#" |
| struct Params { |
| n: u32, c: u32, din: u32, hin: u32, win: u32, |
| scale: u32, dout: u32, hout: u32, wout: u32, |
| p0: u32, p1: u32, p2: u32, |
| }; |
| @group(0) @binding(0) var<storage, read> x: array<f32>; |
| @group(0) @binding(1) var<storage, read_write> y: array<f32>; |
| @group(0) @binding(2) var<uniform> p: Params; |
| |
| @compute @workgroup_size(64) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let idx = gid.x; |
| let total = p.n * p.c * p.dout * p.hout * p.wout; |
| if (idx >= total) { return; } |
| var t = idx; |
| let ow = t % p.wout; t = t / p.wout; |
| let oh = t % p.hout; t = t / p.hout; |
| let od = t % p.dout; t = t / p.dout; |
| let co = t % p.c; |
| let bn = t / p.c; |
| let iw = ow / p.scale; |
| let ih = oh / p.scale; |
| let id = od / p.scale; |
| let src = (((bn * p.c + co) * p.din + id) * p.hin + ih) * p.win + iw; |
| y[idx] = x[src]; |
| } |
| "#; |
|
|
| fn cpu_up(x: &[f32], p: &Params) -> Vec<f32> { |
| let mut y = vec![0.0f32; (p.n * p.c * p.dout * p.hout * p.wout) as usize]; |
| for idx in 0..y.len() { |
| let mut t = idx; |
| let ow = t % p.wout as usize; t /= p.wout as usize; |
| let oh = t % p.hout as usize; t /= p.hout as usize; |
| let od = t % p.dout as usize; t /= p.dout as usize; |
| let co = t % p.c as usize; |
| let bn = t / p.c as usize; |
| let iw = ow / p.scale as usize; |
| let ih = oh / p.scale as usize; |
| let id = od / p.scale as usize; |
| y[idx] = x[(((bn * p.c as usize + co) * p.din as usize + id) * p.hin as usize + ih) * p.win as usize + iw]; |
| } |
| y |
| } |
|
|
| fn main() { |
| let p = Params { n: 1, c: 2, din: 2, hin: 2, win: 2, scale: 2, dout: 4, hout: 4, wout: 4, |
| _p0: 0, _p1: 0, _p2: 0 }; |
| let xn = (p.n * p.c * p.din * p.hin * p.win) as usize; |
| let yn = (p.n * p.c * p.dout * p.hout * p.wout) as usize; |
| let mut x = vec![0.0f32; xn]; |
| for i in 0..xn { x[i] = ((i * 13 % 71) as f32 - 35.0) / 19.0; } |
| let expect = cpu_up(&x, &p); |
|
|
| let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor { backends: wgpu::Backends::VULKAN, ..Default::default() }); |
| let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { |
| power_preference: wgpu::PowerPreference::HighPerformance, compatible_surface: None, force_fallback_adapter: false, |
| })).expect("no Vulkan adapter"); |
| let info = adapter.get_info(); |
| println!("adapter: {} | backend: {:?}", info.name, info.backend); |
| assert_eq!(info.backend, wgpu::Backend::Vulkan); |
|
|
| let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { |
| label: Some("upsample3d"), required_features: wgpu::Features::empty(), |
| required_limits: wgpu::Limits::downlevel_defaults(), ..Default::default() |
| })).expect("device"); |
|
|
| let x_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| label: None, contents: bytemuck::cast_slice(&x), usage: wgpu::BufferUsages::STORAGE, |
| }); |
| let y_buf = device.create_buffer(&wgpu::BufferDescriptor { |
| label: None, size: (yn * 4) as u64, |
| usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, mapped_at_creation: false, |
| }); |
| let readback = device.create_buffer(&wgpu::BufferDescriptor { |
| label: None, size: (yn * 4) as u64, |
| usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false, |
| }); |
| let p_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| label: None, contents: bytemuck::bytes_of(&p), usage: wgpu::BufferUsages::UNIFORM, |
| }); |
| let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { |
| label: Some("upsample3d"), source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(WGSL)), |
| }); |
| let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { |
| label: Some("upsample3d"), layout: None, module: &module, entry_point: Some("main"), |
| compilation_options: Default::default(), cache: None, |
| }); |
| let bg = device.create_bind_group(&wgpu::BindGroupDescriptor { |
| label: None, layout: &pipeline.get_bind_group_layout(0), |
| entries: &[ |
| wgpu::BindGroupEntry { binding: 0, resource: x_buf.as_entire_binding() }, |
| wgpu::BindGroupEntry { binding: 1, resource: y_buf.as_entire_binding() }, |
| wgpu::BindGroupEntry { binding: 2, resource: p_buf.as_entire_binding() }, |
| ], |
| }); |
| let mut enc = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); |
| { |
| let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None }); |
| pass.set_pipeline(&pipeline); pass.set_bind_group(0, &bg, &[]); |
| pass.dispatch_workgroups(((yn + 63) / 64) as u32, 1, 1); |
| } |
| enc.copy_buffer_to_buffer(&y_buf, 0, &readback, 0, (yn * 4) as u64); |
| queue.submit(Some(enc.finish())); |
| let slice = readback.slice(..); |
| slice.map_async(wgpu::MapMode::Read, |r| r.expect("map")); |
| device.poll(wgpu::PollType::wait_indefinitely()).expect("poll"); |
| let data = slice.get_mapped_range(); |
| let got: &[f32] = bytemuck::cast_slice(&data); |
| let mut max_err = 0.0f32; |
| for (a, b) in got.iter().zip(expect.iter()) { max_err = max_err.max((a - b).abs()); } |
| println!("upsample3d N={} C={} D={}->{} scale={} max_abs_err={max_err:.3e}", |
| p.n, p.c, p.din, p.dout, p.scale); |
| assert!(max_err < 1e-6, "upsample3d mismatch"); |
| println!("VULKAN UPSAMPLE3D OK"); |
| }
|
|
|
