baedeker_core/runtime/
gpu.rs

1// Copyright (C) 2026 Industrial Algebra
2// SPDX-License-Identifier: Apache-2.0
3
4//! Optional GPU backend slot for bulk SIMD offload.
5//!
6//! The trait defines the contract a backend (Borsalino/Metal on iOS, a
7//! wgpu-based desktop backend, or a test fake) implements so the runtime
8//! can dispatch bulk SIMD work to GPU compute. Backends are identified by
9//! opaque handles, keeping the trait object-safe and FFI-friendly.
10//!
11//! See `docs/borsalino-integration.md` for the integration design.
12
13use alloc::string::String;
14use alloc::vec::Vec;
15
16/// Opaque handle to a compiled compute kernel, owned by the backend.
17pub type GpuKernelId = u64;
18
19/// Opaque handle to a GPU buffer, owned by the backend.
20pub type GpuBufferId = u64;
21
22/// A GPU compute backend for bulk SIMD offload (Borsalino Level 1).
23///
24/// The runtime uses this to offload large vector operations: upload WASM
25/// linear-memory regions to buffers, dispatch a pre-compiled WGSL kernel,
26/// and read results back. Below the offload threshold the register IR
27/// executes element-wise on CPU.
28pub trait GpuBackend: core::fmt::Debug {
29    /// Human-readable backend name for diagnostics.
30    fn name(&self) -> &str;
31
32    /// Compile a WGSL compute kernel (once, cached by the backend).
33    fn compile(&mut self, name: &str, wgsl: &str) -> Result<GpuKernelId, GpuError>;
34
35    /// Create a GPU buffer initialized with `data`.
36    fn create_buffer(&mut self, data: &[u8]) -> Result<GpuBufferId, GpuError>;
37
38    /// Create an uninitialized GPU buffer of `byte_len` bytes.
39    fn create_buffer_uninit(&mut self, byte_len: usize) -> Result<GpuBufferId, GpuError>;
40
41    /// Dispatch a kernel over `workgroups` (x, y, z) with `buffers` bound
42    /// in order.
43    ///
44    /// Each workgroup runs the backend's default thread count. Prefer
45    /// [`dispatch_verified`](GpuBackend::dispatch_verified) when the workgroup
46    /// size is known, so non-default thread counts dispatch correctly.
47    fn dispatch(
48        &mut self,
49        kernel: GpuKernelId,
50        buffers: &[GpuBufferId],
51        workgroups: [u32; 3],
52    ) -> Result<(), GpuError>;
53
54    /// Dispatch a kernel with an explicit per-workgroup thread count.
55    ///
56    /// Like [`dispatch`](GpuBackend::dispatch), but each workgroup runs
57    /// `threads_per_group` threads rather than a backend-implied default.
58    /// This matters for kernels whose WGSL declares a non-default
59    /// `@workgroup_size`: `dispatch` silently uses the backend's default
60    /// (256 for Borsalino), which mis-dispatches such kernels.
61    ///
62    /// Backends that verify workgroup divisibility (Borsalino's
63    /// `dispatch_verified`) construct their proof from
64    /// `(workgroups, threads_per_group)` here; the divisibility check runs
65    /// on the x-dimension product `workgroups[0] * threads_per_group[0]`,
66    /// matching Borsalino's 1-D proof scope.
67    ///
68    /// The default implementation forwards to [`dispatch`](GpuBackend::dispatch),
69    /// ignoring `threads_per_group`. Concrete backends that honour explicit
70    /// workgroup sizes override this.
71    fn dispatch_verified(
72        &mut self,
73        kernel: GpuKernelId,
74        buffers: &[GpuBufferId],
75        workgroups: [u32; 3],
76        threads_per_group: [u32; 3],
77    ) -> Result<(), GpuError> {
78        let _ = threads_per_group;
79        self.dispatch(kernel, buffers, workgroups)
80    }
81
82    /// Read a buffer's full contents back to host memory.
83    fn read_buffer(&mut self, buffer: GpuBufferId) -> Result<Vec<u8>, GpuError>;
84}
85
86/// GPU backend failure.
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct GpuError {
89    pub kind: GpuErrorKind,
90    pub message: String,
91}
92
93/// The category of a GPU backend failure.
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum GpuErrorKind {
96    /// No GPU or driver available on this host.
97    Unavailable,
98    /// Kernel compilation failed.
99    CompileFailed,
100    /// Buffer allocation failed.
101    OutOfMemory,
102    /// Kernel dispatch failed.
103    DispatchFailed,
104    /// Buffer readback failed.
105    ReadbackFailed,
106}
107
108#[cfg(test)]
109mod tests {
110    use alloc::boxed::Box;
111
112    use super::*;
113    use crate::lower::RegModule;
114
115    /// Shared counters for observing a fake backend's activity.
116    #[derive(Debug, Default)]
117    struct FakeStats {
118        compiled: usize,
119        dispatches: usize,
120    }
121
122    /// A recording fake backend for slot tests.
123    #[derive(Debug)]
124    struct FakeBackend {
125        stats: alloc::rc::Rc<core::cell::RefCell<FakeStats>>,
126        buffers: Vec<(GpuBufferId, usize)>,
127    }
128
129    impl FakeBackend {
130        fn new(stats: alloc::rc::Rc<core::cell::RefCell<FakeStats>>) -> Self {
131            Self {
132                stats,
133                buffers: Vec::new(),
134            }
135        }
136    }
137
138    impl GpuBackend for FakeBackend {
139        fn name(&self) -> &str {
140            "fake"
141        }
142
143        fn compile(&mut self, _name: &str, _wgsl: &str) -> Result<GpuKernelId, GpuError> {
144            let mut stats = self.stats.borrow_mut();
145            stats.compiled += 1;
146            Ok(stats.compiled as u64)
147        }
148
149        fn create_buffer(&mut self, data: &[u8]) -> Result<GpuBufferId, GpuError> {
150            let id = self.buffers.len() as u64 + 1;
151            self.buffers.push((id, data.len()));
152            Ok(id)
153        }
154
155        fn create_buffer_uninit(&mut self, byte_len: usize) -> Result<GpuBufferId, GpuError> {
156            let id = self.buffers.len() as u64 + 1;
157            self.buffers.push((id, byte_len));
158            Ok(id)
159        }
160
161        fn dispatch(
162            &mut self,
163            _kernel: GpuKernelId,
164            _buffers: &[GpuBufferId],
165            _workgroups: [u32; 3],
166        ) -> Result<(), GpuError> {
167            self.stats.borrow_mut().dispatches += 1;
168            Ok(())
169        }
170
171        fn read_buffer(&mut self, buffer: GpuBufferId) -> Result<Vec<u8>, GpuError> {
172            let (_, size) = self
173                .buffers
174                .iter()
175                .find(|(id, _)| *id == buffer)
176                .expect("buffer exists");
177            Ok(alloc::vec![0xAA; *size])
178        }
179    }
180
181    fn empty_module() -> RegModule {
182        RegModule {
183            funcs: Vec::new(),
184            exports: Vec::new(),
185            imported_func_count: 0,
186            imported_funcs: Vec::new(),
187            imported_memories: Vec::new(),
188            imported_globals: Vec::new(),
189            imported_tables: Vec::new(),
190            start: None,
191            memories: Vec::new(),
192            globals: Vec::new(),
193            tables: Vec::new(),
194            elements: Vec::new(),
195            types: Vec::new(),
196            data: Vec::new(),
197            imported_memory_count: 0,
198            imported_global_count: 0,
199            imported_table_count: 0,
200        }
201    }
202
203    #[test]
204    fn store_has_no_gpu_by_default() {
205        let module = empty_module();
206        let store = crate::runtime::Store::instantiate(&module).unwrap();
207        assert!(!store.has_gpu());
208    }
209
210    #[test]
211    fn gpu_slot_accepts_and_exercises_a_backend() {
212        let module = empty_module();
213        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
214        let stats = alloc::rc::Rc::new(core::cell::RefCell::new(FakeStats::default()));
215        store.set_gpu(Box::new(FakeBackend::new(stats.clone())));
216        assert!(
217            store
218                .with_gpu_mut(|gpu| gpu.name() == "fake")
219                .expect("backend installed")
220        );
221
222        let (kernel, buf_a, buf_out) = store
223            .with_gpu_mut(|gpu| {
224                let kernel = gpu.compile("vadd_f32x4", "@compute fn vadd() {}")?;
225                let buf_a = gpu.create_buffer(&[1, 2, 3, 4])?;
226                let buf_out = gpu.create_buffer_uninit(16)?;
227                Ok::<_, GpuError>((kernel, buf_a, buf_out))
228            })
229            .expect("backend installed")
230            .unwrap();
231        store
232            .with_gpu_mut(|gpu| gpu.dispatch(kernel, &[buf_a, buf_out], [1, 1, 1]))
233            .expect("backend installed")
234            .unwrap();
235        let result = store
236            .with_gpu_mut(|gpu| gpu.read_buffer(buf_out))
237            .expect("backend installed")
238            .unwrap();
239        assert_eq!(result, alloc::vec![0xAA; 16]);
240
241        store.clear_gpu();
242        assert!(!store.has_gpu());
243    }
244
245    #[test]
246    fn dispatch_verified_default_delegates_to_dispatch() {
247        // A backend that only implements `dispatch` (the required method)
248        // inherits the trait default for `dispatch_verified`, which must
249        // forward to `dispatch` and accept an arbitrary `threads_per_group`.
250        let module = empty_module();
251        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
252        let stats = alloc::rc::Rc::new(core::cell::RefCell::new(FakeStats::default()));
253        store.set_gpu(Box::new(FakeBackend::new(stats.clone())));
254        let (kernel, buf) = store
255            .with_gpu_mut(|gpu| {
256                let kernel = gpu.compile("k", "@compute fn k() {}")?;
257                let buf = gpu.create_buffer(&[0; 4])?;
258                Ok::<_, GpuError>((kernel, buf))
259            })
260            .expect("backend installed")
261            .unwrap();
262        // Non-default threads_per_group must be accepted by the default impl.
263        store
264            .with_gpu_mut(|gpu| gpu.dispatch_verified(kernel, &[buf], [4, 1, 1], [128, 1, 1]))
265            .expect("backend installed")
266            .unwrap();
267        assert_eq!(
268            stats.borrow().dispatches,
269            1,
270            "default forwarded to dispatch"
271        );
272    }
273
274    fn memory_module() -> RegModule {
275        let mut module = empty_module();
276        module.memories.push(crate::types::MemType {
277            limits: crate::types::Limits { min: 1, max: None },
278        });
279        module
280    }
281
282    /// Write an f32 into the store's memory 0.
283    fn store_f32(store: &mut crate::runtime::Store, addr: u32, value: f32) {
284        store
285            .with_memory_mut(0, |mem| {
286                mem[addr as usize..addr as usize + 4].copy_from_slice(&value.to_le_bytes());
287            })
288            .expect("memory 0");
289    }
290
291    fn read_f32(store: &crate::runtime::Store, addr: u32) -> f32 {
292        store
293            .with_memory(0, |mem| {
294                f32::from_le_bytes(mem[addr as usize..addr as usize + 4].try_into().unwrap())
295            })
296            .expect("memory 0")
297    }
298
299    #[test]
300    fn f32_add_region_below_threshold_uses_cpu() {
301        let module = memory_module();
302        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
303        store_f32(&mut store, 0, 1.5);
304        store_f32(&mut store, 64, 2.25);
305        let stats = alloc::rc::Rc::new(core::cell::RefCell::new(FakeStats::default()));
306        store.set_gpu(Box::new(FakeBackend::new(stats.clone())));
307
308        store.f32_add_region(0, 64, 128, 1).unwrap();
309
310        assert_eq!(read_f32(&store, 128), 3.75);
311        assert_eq!(
312            stats.borrow().dispatches,
313            0,
314            "below threshold must not dispatch"
315        );
316    }
317
318    #[test]
319    fn f32_add_region_above_threshold_dispatches() {
320        let module = memory_module();
321        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
322        store.set_offload_threshold(4);
323        let stats = alloc::rc::Rc::new(core::cell::RefCell::new(FakeStats::default()));
324        store.set_gpu(Box::new(FakeBackend::new(stats.clone())));
325
326        // 8 elements = 32 bytes per region, all within one page.
327        store.f32_add_region(0, 256, 512, 8).unwrap();
328
329        // The fake's readback fills the output region with 0xAA.
330        let all_filled = store
331            .with_memory(0, |mem| mem[512..512 + 32].iter().all(|byte| *byte == 0xAA))
332            .expect("memory 0");
333        assert!(all_filled);
334        assert_eq!(stats.borrow().dispatches, 1);
335        assert_eq!(stats.borrow().compiled, 1, "kernel compiled once");
336
337        // A second call reuses the cached kernel.
338        store.f32_add_region(0, 256, 512, 8).unwrap();
339        assert_eq!(stats.borrow().compiled, 1, "kernel not recompiled");
340        assert_eq!(stats.borrow().dispatches, 2);
341    }
342
343    #[test]
344    fn f32_add_region_without_gpu_uses_cpu() {
345        let module = memory_module();
346        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
347        store_f32(&mut store, 0, 1.0);
348        store_f32(&mut store, 64, 2.0);
349        store.set_offload_threshold(0);
350
351        store.f32_add_region(0, 64, 128, 1).unwrap();
352        assert_eq!(read_f32(&store, 128), 3.0);
353    }
354
355    #[test]
356    fn f32_add_region_out_of_bounds_traps() {
357        let module = memory_module();
358        let mut store = crate::runtime::Store::instantiate(&module).unwrap();
359        let error = store.f32_add_region(0, 65533, 128, 1).unwrap_err();
360        assert_eq!(
361            error.kind,
362            crate::runtime::RuntimeErrorKind::Trap(
363                crate::runtime::RuntimeTrap::OutOfBoundsMemoryAccess
364            )
365        );
366    }
367}
baedeker_core/runtime/gpu.rs

baedeker_core/runtime/
gpu.rs
