metrale_gpu_runtime/
kernel_args.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `KernelLaunch`, the builder for kernel launches on any
4//! `GpuBackend`. Each argument is recorded as a buffer or as scalar bytes, so the
5//! Metal backend can bind buffers with `setBuffer:offset:atIndex:` and scalars with
6//! `setBytes:length:atIndex:`.
7//!
8//! Owner: gpu-runtime.
9//! Invariants:
10//! - Each `arg_*` call adds exactly one `KernelArg` to the launch; a
11//!   `CUtensorMap` occupies 16 consecutive storage slots but is still one argument.
12//! - Arguments reach `GpuBackend::launch_typed` in the order they were added.
13//!
14//! Usage:
15//!
16//! ```ignore
17//! KernelLaunch::new(gpu, kernel)
18//!     .grid([num_tokens, 1, 1])
19//!     .block([256, 1, 1])
20//!     .arg_ptr(input)
21//!     .arg_u32(hidden_size)
22//!     .arg_f32(eps)
23//!     .launch(stream)?;
24//! ```
25
26use anyhow::Result;
27
28use crate::gpu::{DevicePtr, GpuBackend, KernelArg, KernelHandle};
29
30/// 2026-09-25: Per-arg metadata: the first `storage` slot it occupies, and for a
31/// scalar its byte width. Buffer args set `is_buffer` and leave `byte_len` 0.
32struct ArgKind {
33    is_buffer: bool,
34    /// 2026-09-25: Byte count for scalar args: 4 for u32/i32/f32, 8 for u64, 128
35    /// for a `CUtensorMap`.
36    byte_len: u16,
37    /// 2026-09-25: Starting slot in `storage`. A `CUtensorMap` occupies 16
38    /// consecutive slots, so an arg's slot is not its position in `kinds`.
39    slot: u32,
40}
41
42/// 2026-09-25: Builder for a kernel launch: grid, block, dynamic shared memory
43/// and typed arguments. `launch()` passes the arguments as `&[KernelArg]` to
44/// `GpuBackend::launch_typed`. Grid and block default to `[1, 1, 1]`, shared
45/// memory to 0.
46pub struct KernelLaunch<'a> {
47    gpu: &'a dyn GpuBackend,
48    kernel: KernelHandle,
49    grid: [u32; 3],
50    block: [u32; 3],
51    shared_mem: u32,
52    /// 2026-09-25: Argument bytes in u64 slots: little-endian for scalars, the
53    /// device address for buffers.
54    storage: Vec<u64>,
55    /// 2026-09-25: One entry per argument, in order.
56    kinds: Vec<ArgKind>,
57}
58
59impl<'a> KernelLaunch<'a> {
60    pub fn new(gpu: &'a dyn GpuBackend, kernel: KernelHandle) -> Self {
61        Self {
62            gpu,
63            kernel,
64            grid: [1, 1, 1],
65            block: [1, 1, 1],
66            shared_mem: 0,
67            storage: Vec::with_capacity(16),
68            kinds: Vec::with_capacity(16),
69        }
70    }
71
72    pub fn grid(mut self, grid: [u32; 3]) -> Self {
73        self.grid = grid;
74        self
75    }
76
77    pub fn block(mut self, block: [u32; 3]) -> Self {
78        self.block = block;
79        self
80    }
81
82    pub fn shared_mem(mut self, bytes: u32) -> Self {
83        self.shared_mem = bytes;
84        self
85    }
86
87    /// 2026-09-25: Add a device buffer argument.
88    pub fn arg_ptr(mut self, p: DevicePtr) -> Self {
89        let slot = self.storage.len() as u32;
90        self.storage.push(p.0);
91        self.kinds.push(ArgKind {
92            is_buffer: true,
93            byte_len: 0,
94            slot,
95        });
96        self
97    }
98
99    /// 2026-09-25: Add a 128-byte `CUtensorMap` by value, for a kernel parameter
100    /// declared `__grid_constant__ const CUtensorMap`. The bytes fill 16
101    /// consecutive slots and form one argument.
102    pub fn arg_tensormap(mut self, map: &[u8; 128]) -> Self {
103        let slot = self.storage.len() as u32;
104        for c in map.chunks(8) {
105            let mut w = [0u8; 8];
106            w.copy_from_slice(c);
107            self.storage.push(u64::from_le_bytes(w));
108        }
109        self.kinds.push(ArgKind {
110            is_buffer: false,
111            byte_len: 128,
112            slot,
113        });
114        self
115    }
116
117    pub fn arg_u32(mut self, v: u32) -> Self {
118        let slot = self.storage.len() as u32;
119        self.storage.push(v as u64);
120        self.kinds.push(ArgKind {
121            is_buffer: false,
122            byte_len: 4,
123            slot,
124        });
125        self
126    }
127
128    pub fn arg_u64(mut self, v: u64) -> Self {
129        let slot = self.storage.len() as u32;
130        self.storage.push(v);
131        self.kinds.push(ArgKind {
132            is_buffer: false,
133            byte_len: 8,
134            slot,
135        });
136        self
137    }
138
139    pub fn arg_i32(mut self, v: i32) -> Self {
140        let slot = self.storage.len() as u32;
141        self.storage.push(v as u32 as u64);
142        self.kinds.push(ArgKind {
143            is_buffer: false,
144            byte_len: 4,
145            slot,
146        });
147        self
148    }
149
150    pub fn arg_f32(mut self, v: f32) -> Self {
151        let slot = self.storage.len() as u32;
152        self.storage.push(f32::to_bits(v) as u64);
153        self.kinds.push(ArgKind {
154            is_buffer: false,
155            byte_len: 4,
156            slot,
157        });
158        self
159    }
160
161    /// 2026-09-25: Launch on `stream` through `GpuBackend::launch_typed`.
162    ///
163    /// With `debug_sync_kernels` on, the stream is synchronised after a
164    /// successful launch and a fault is returned with the launch geometry and a
165    /// backtrace.
166    pub fn launch(self, stream: u64) -> Result<()> {
167        let mut args: Vec<KernelArg<'_>> = Vec::with_capacity(self.kinds.len());
168        for kind in self.kinds.iter() {
169            let slot = &self.storage[kind.slot as usize];
170            if kind.is_buffer {
171                args.push(KernelArg::Buffer(DevicePtr(*slot)));
172            } else {
173                // 2026-09-25: SAFETY: the arg's bytes start at `slot` and run
174                // `byte_len` bytes through consecutive slots of
175                // `self.storage`, which is not modified while `args` lives.
176                let bytes = unsafe {
177                    std::slice::from_raw_parts(
178                        slot as *const u64 as *const u8,
179                        kind.byte_len as usize,
180                    )
181                };
182                args.push(KernelArg::Bytes(bytes));
183            }
184        }
185        let r = self.gpu.launch_typed(
186            self.kernel,
187            self.grid,
188            self.block,
189            self.shared_mem,
190            stream,
191            &args,
192        );
193        // 2026-09-25: `METRALE_DEBUG_SYNC_KERNELS=1` on the CUDA backend: an
194        // asynchronous fault is reported at the launch that caused it.
195        if r.is_ok() && self.gpu.debug_sync_kernels() {
196            self.gpu.synchronize(stream).map_err(|e| {
197                let bt = std::backtrace::Backtrace::force_capture();
198                anyhow::anyhow!(
199                    "METRALE_DEBUG_SYNC_KERNELS: async GPU fault immediately after kernel launch \
200                     grid={:?} block={:?} shared_mem={}: {e}\nLAUNCH BACKTRACE:\n{bt}",
201                    self.grid,
202                    self.block,
203                    self.shared_mem
204                )
205            })?;
206        }
207        r
208    }
209}
210
211/// 2026-09-25: `a / b`, rounded up.
212pub fn div_ceil(a: u32, b: u32) -> u32 {
213    a.div_ceil(b)
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::gpu::mock::MockGpuBackend;
220
221    #[test]
222    fn test_kernel_launch_builder() {
223        let gpu = MockGpuBackend::new();
224        let kernel = gpu.kernel("test", "test_kernel").unwrap();
225
226        let result = KernelLaunch::new(&gpu, kernel)
227            .grid([4, 1, 1])
228            .block([256, 1, 1])
229            .arg_ptr(DevicePtr(0x1000))
230            .arg_u32(42)
231            .arg_f32(1.5)
232            .launch(0);
233
234        assert!(result.is_ok());
235        assert_eq!(gpu.launch_count(), 1);
236    }
237
238    /// 2026-09-25: A 128-byte by-value argument (a `CUtensorMap`) occupies 16
239    /// consecutive slots with one param entry, and its bytes round-trip.
240    #[test]
241    fn a_128_byte_arg_is_not_truncated() {
242        use crate::gpu::{KernelArg, pack_kernel_args};
243        let map: Vec<u8> = (0..128u16).map(|i| (i * 7 % 251) as u8).collect();
244        let args = [
245            KernelArg::Buffer(DevicePtr(0xDEAD_BEEF)),
246            KernelArg::Bytes(&map),
247            KernelArg::Bytes(&42u32.to_le_bytes()),
248        ];
249        let (storage, starts) = pack_kernel_args(&args);
250
251        assert_eq!(
252            starts.len(),
253            3,
254            "one param entry per argument, not per slot"
255        );
256        assert_eq!(starts, vec![0, 1, 17], "the map occupies slots 1..=16");
257        assert_eq!(storage.len(), 18);
258        assert_eq!(storage[0], 0xDEAD_BEEF);
259
260        let round_trip: Vec<u8> = storage[starts[1]..starts[1] + 16]
261            .iter()
262            .flat_map(|w| w.to_le_bytes())
263            .collect();
264        assert_eq!(
265            round_trip, map,
266            "128-byte struct arg was corrupted or truncated"
267        );
268        assert_eq!(
269            storage[starts[2]] as u32, 42,
270            "the arg after it is still intact"
271        );
272    }
273
274    /// 2026-09-25: Through the builder, a tensormap between two ordinary args
275    /// leaves both intact and counts as one argument.
276    #[test]
277    fn a_tensormap_arg_does_not_shift_the_args_around_it() {
278        let gpu = MockGpuBackend::new();
279        let kernel = gpu.kernel("test", "tma_kernel").unwrap();
280        let map = [0xABu8; 128];
281
282        let b = KernelLaunch::new(&gpu, kernel)
283            .arg_ptr(DevicePtr(0x1000))
284            .arg_tensormap(&map)
285            .arg_u32(7);
286
287        assert_eq!(b.kinds.len(), 3, "one param entry per arg, not per slot");
288        assert_eq!(b.storage.len(), 1 + 16 + 1);
289        assert_eq!(b.kinds[0].slot, 0);
290        assert_eq!(b.kinds[1].slot, 1);
291        assert_eq!(b.kinds[1].byte_len, 128);
292        assert_eq!(
293            b.kinds[2].slot, 17,
294            "the arg after the map must not be shifted"
295        );
296        assert_eq!(b.storage[b.kinds[2].slot as usize] as u32, 7);
297        assert!(b.launch(0).is_ok());
298    }
299
300    /// 2026-09-25: A zero-length byte arg still gets a slot to point at.
301    #[test]
302    fn an_empty_byte_arg_still_gets_one_slot() {
303        use crate::gpu::{KernelArg, pack_kernel_args};
304        let (storage, starts) =
305            pack_kernel_args(&[KernelArg::Bytes(&[]), KernelArg::Bytes(&[9u8])]);
306        assert_eq!(starts, vec![0, 1]);
307        assert_eq!(storage.len(), 2);
308        assert_eq!(storage[1], 9);
309    }
310
311    #[test]
312    fn test_div_ceil() {
313        assert_eq!(div_ceil(10, 3), 4);
314        assert_eq!(div_ceil(9, 3), 3);
315        assert_eq!(div_ceil(1, 256), 1);
316        assert_eq!(div_ceil(256, 256), 1);
317        assert_eq!(div_ceil(257, 256), 2);
318    }
319}
metrale_gpu_runtime/kernel_args.rs

metrale_gpu_runtime/
kernel_args.rs
