kopitiam_gpu/ops/vector_add.rs
1//! Elementwise vector add — the demonstrator op that proves the cascade.
2//!
3//! `out[i] = a[i] + b[i]`. Trivial arithmetic on purpose: the point of this op
4//! is not the maths but to exercise the whole GPU->CPU machinery end to end —
5//! host data into storage buffers, a WGSL compute dispatch, readback through a
6//! staging buffer, and a pure-Rust twin that gives the identical answer when
7//! there is no GPU.
8//!
9//! Both paths return `Vec<f32>` of the same length as the inputs, and for
10//! elementwise `f32` addition the GPU and CPU results are **bit-for-bit equal**
11//! (IEEE-754 addition is exact and deterministic here — no reordering, no
12//! fused-multiply-add), which is why the tests can assert plain equality when a
13//! GPU is present.
14
15use crate::context::GpuContext;
16use crate::executor::{ComputeOp, GpuOpError};
17use wgpu::util::DeviceExt;
18
19/// Must match `@workgroup_size(64)` in `shaders/vector_add.wgsl`. The host
20/// dispatches `ceil(n / WORKGROUP_SIZE)` workgroups; the shader guards the tail.
21const WORKGROUP_SIZE: u32 = 64;
22
23/// The two input vectors, borrowed. Equal length is a precondition (checked in
24/// both paths); mismatched lengths are a caller bug, not something to paper over.
25pub struct VectorAddInput<'a> {
26 pub a: &'a [f32],
27 pub b: &'a [f32],
28}
29
30/// Elementwise vector addition. Zero-sized: it carries no state, it just names
31/// the operation so the [`ComputeOp`] impl and its two kernels hang off a type.
32pub struct VectorAdd;
33
34impl ComputeOp for VectorAdd {
35 type Input<'a> = VectorAddInput<'a>;
36 type Output = Vec<f32>;
37
38 fn compute_gpu(
39 &self,
40 ctx: &GpuContext,
41 input: &Self::Input<'_>,
42 ) -> Result<Self::Output, GpuOpError> {
43 vector_add_gpu(ctx, input.a, input.b)
44 }
45
46 fn compute_cpu(&self, input: &Self::Input<'_>) -> Self::Output {
47 vector_add_cpu(input.a, input.b)
48 }
49}
50
51/// The pure-Rust floor of the cascade. Infallible and correct.
52///
53/// Length rule: the output is as long as the SHORTER input (`zip` stops there).
54/// The GPU path enforces equal lengths and errors on mismatch; to keep the two
55/// twins returning the same thing, callers should only ever pass equal-length
56/// slices. Given equal lengths, this and the GPU kernel agree exactly.
57pub fn vector_add_cpu(a: &[f32], b: &[f32]) -> Vec<f32> {
58 a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()
59}
60
61/// The wgpu compute path. Fallible: any wgpu-level problem returns `Err` so the
62/// [`crate::Executor`] can fall back to [`vector_add_cpu`].
63///
64/// The dispatch shape, binding by binding, mirrors `shaders/vector_add.wgsl`:
65///
66/// 1. **Upload** `a` and `b` into read-only STORAGE buffers via
67/// `create_buffer_init` (bytemuck reinterprets `&[f32]` as the `&[u8]` wgpu
68/// wants — same bytes, no copy beyond the upload itself).
69/// 2. **Allocate** an `out` STORAGE buffer with `COPY_SRC` (the shader writes it;
70/// we then copy it to a mappable staging buffer — you cannot map a STORAGE
71/// buffer directly for reading).
72/// 3. **Bind** all three at group 0, bindings 0/1/2, matching the WGSL exactly.
73/// 4. **Dispatch** `ceil(n / 64)` workgroups in x; the shader's bounds guard
74/// handles the rounded-up tail invocations.
75/// 5. **Read back**: copy `out` -> staging, submit, `map_async` + `poll(Wait)`
76/// to block until the GPU is done, then `bytemuck`-cast the mapped bytes back
77/// to `f32`. `poll(Wait)` is the synchronous join point — without it the map
78/// callback may not have fired yet.
79pub fn vector_add_gpu(ctx: &GpuContext, a: &[f32], b: &[f32]) -> Result<Vec<f32>, GpuOpError> {
80 if a.len() != b.len() {
81 return Err(GpuOpError::InvalidInput(format!(
82 "vector_add needs equal-length inputs, got {} and {}",
83 a.len(),
84 b.len()
85 )));
86 }
87 let n = a.len();
88 // An empty dispatch is legal but pointless (and a zero-sized buffer is a
89 // validation error on some backends), so short-circuit it.
90 if n == 0 {
91 return Ok(Vec::new());
92 }
93
94 let device = ctx.device();
95 let queue = ctx.queue();
96 // Bytes needed for one vector: a and b are equal length (checked above), so
97 // size_of_val(a) is the byte length of each buffer. u64 for BufferAddress.
98 let byte_len = std::mem::size_of_val(a) as wgpu::BufferAddress;
99
100 // (1) input buffers, uploaded at creation.
101 let a_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
102 label: Some("vector_add.a"),
103 contents: bytemuck::cast_slice(a),
104 usage: wgpu::BufferUsages::STORAGE,
105 });
106 let b_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
107 label: Some("vector_add.b"),
108 contents: bytemuck::cast_slice(b),
109 usage: wgpu::BufferUsages::STORAGE,
110 });
111
112 // (2) output buffer: the shader writes it (STORAGE) and we copy it out
113 // afterwards (COPY_SRC). Not mappable itself — hence the staging buffer.
114 let out_buf = device.create_buffer(&wgpu::BufferDescriptor {
115 label: Some("vector_add.out"),
116 size: byte_len,
117 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
118 mapped_at_creation: false,
119 });
120 // Staging buffer we CAN map for reading: COPY_DST target + MAP_READ.
121 let staging = device.create_buffer(&wgpu::BufferDescriptor {
122 label: Some("vector_add.staging"),
123 size: byte_len,
124 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
125 mapped_at_creation: false,
126 });
127
128 // Shader + pipeline. `layout: None` lets wgpu derive the bind-group layout
129 // from the WGSL, so the binding indices come straight from the shader.
130 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
131 label: Some("vector_add.wgsl"),
132 source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/vector_add.wgsl").into()),
133 });
134 let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
135 label: Some("vector_add.pipeline"),
136 layout: None,
137 module: &shader,
138 entry_point: Some("main"),
139 compilation_options: wgpu::PipelineCompilationOptions::default(),
140 cache: None,
141 });
142
143 // (3) bind group: group 0, bindings 0/1/2 = a, b, out. Order and indices
144 // MUST match the WGSL @binding numbers.
145 let bind_group_layout = pipeline.get_bind_group_layout(0);
146 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
147 label: Some("vector_add.bind_group"),
148 layout: &bind_group_layout,
149 entries: &[
150 wgpu::BindGroupEntry {
151 binding: 0,
152 resource: a_buf.as_entire_binding(),
153 },
154 wgpu::BindGroupEntry {
155 binding: 1,
156 resource: b_buf.as_entire_binding(),
157 },
158 wgpu::BindGroupEntry {
159 binding: 2,
160 resource: out_buf.as_entire_binding(),
161 },
162 ],
163 });
164
165 // (4) encode the dispatch. ceil(n / WORKGROUP_SIZE) workgroups in x.
166 let mut encoder =
167 device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("vector_add") });
168 {
169 let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
170 label: Some("vector_add.pass"),
171 timestamp_writes: None,
172 });
173 pass.set_pipeline(&pipeline);
174 pass.set_bind_group(0, &bind_group, &[]);
175 let workgroups = (n as u32).div_ceil(WORKGROUP_SIZE);
176 pass.dispatch_workgroups(workgroups, 1, 1);
177 }
178 // (5) copy result out and submit.
179 encoder.copy_buffer_to_buffer(&out_buf, 0, &staging, 0, byte_len);
180 queue.submit(Some(encoder.finish()));
181
182 // Map the staging buffer and block until the GPU has actually finished.
183 let slice = staging.slice(..);
184 let (tx, rx) = std::sync::mpsc::channel();
185 slice.map_async(wgpu::MapMode::Read, move |res| {
186 // Ignore send errors: if the receiver is gone we're tearing down anyway.
187 let _ = tx.send(res);
188 });
189 // poll(Wait) is the synchronous join: it drives the GPU work and the map
190 // callback to completion. Without it the recv below could block forever.
191 device
192 .poll(wgpu::PollType::wait_indefinitely())
193 .map_err(|e| GpuOpError::Backend(format!("device poll failed: {e:?}")))?;
194 rx.recv()
195 .map_err(|e| GpuOpError::Backend(format!("map callback dropped: {e}")))?
196 .map_err(|e| GpuOpError::Backend(format!("buffer map failed: {e:?}")))?;
197
198 // wgpu 30's get_mapped_range is fallible (the map could have raced a device
199 // loss); propagate that as a Backend error so the cascade falls to CPU.
200 let data = slice
201 .get_mapped_range()
202 .map_err(|e| GpuOpError::Backend(format!("get_mapped_range failed: {e:?}")))?;
203 let result: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
204 // Drop the mapped view before unmapping (wgpu requires the view be gone).
205 drop(data);
206 staging.unmap();
207
208 Ok(result)
209}