trueno_gpu/
launch_budget.rs

1//! Launch budget validation — the enforcement `ptx-codegen-safety-v1` has always declared.
2//!
3//! `contracts/trueno/ptx-codegen-safety-v1.yaml` equation `register_budget` states
4//!
5//! ```text
6//! forall kernel K:
7//!   reg_count(K)  <= max_regs_per_thread(sm)
8//!   shared_mem(K) <= max_shared_per_block(sm)
9//! postcondition: cuOccupancyMaxActiveBlocksPerMultiprocessor > 0
10//! ```
11//!
12//! and nothing enforced it: the generated `contract_register_budget!` macro is invoked
13//! nowhere in the tree, and its postcondition names an unbound identifier that would not
14//! compile if it ever were. This module is that missing enforcement.
15//!
16//! **Arch-agnostic by construction.** The contract's own `domain` stops at sm_90, and so
17//! does the hand-written arch table in `driver/sys/mod.rs` (`CU_TARGET_COMPUTE_90` is the
18//! last constant). A per-SM limit table is a thing that goes stale every GPU generation —
19//! this one has, twice. So the limits here are **queried from the device**
20//! (`cuDeviceGetAttribute`) rather than looked up, and adding a new architecture requires
21//! no change to this file.
22//!
23//! The policy — [`validate_launch`] — is pure and needs no GPU, so its case table runs in
24//! the required check. Only [`DeviceLimits::query`] and [`KernelAttributes::query`] need
25//! CUDA.
26
27/// Per-kernel resource usage, as reported by the JIT for a *compiled* kernel.
28///
29/// These are properties of the cubin the driver actually produced — not of the PTX we
30/// emitted — which is why they can only be read back after a module load.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct KernelAttributes {
33    /// Largest block size this kernel can be launched with (`CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK`).
34    pub max_threads_per_block: u32,
35    /// Registers used per thread (`CU_FUNC_ATTRIBUTE_NUM_REGS`).
36    pub num_regs: u32,
37    /// Statically declared shared memory, bytes (`CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES`).
38    pub static_shared_bytes: u32,
39    /// Per-thread local memory, bytes (`CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES`). Non-zero means
40    /// the kernel spilled registers to local memory — legal, but a performance cliff.
41    pub local_bytes: u32,
42}
43
44/// Device-side limits, queried rather than tabulated.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct DeviceLimits {
47    /// `CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK`.
48    pub max_threads_per_block: u32,
49    /// `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK`.
50    pub max_shared_per_block: u32,
51    /// `CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK`.
52    pub max_regs_per_block: u32,
53    /// `CU_DEVICE_ATTRIBUTE_WARP_SIZE`.
54    pub warp_size: u32,
55}
56
57/// A way a launch violates the kernel's or the device's budget.
58///
59/// Every variant names both sides of the comparison: a violation you cannot act on is a
60/// log line, not a gate.
61#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
62pub enum LaunchBudgetViolation {
63    /// Block size exceeds what the compiled kernel supports (usually register pressure).
64    #[error("block size {requested} exceeds the kernel's max_threads_per_block {kernel_max} (register pressure); launch would fail with CUDA_ERROR_INVALID_VALUE")]
65    BlockExceedsKernelMax {
66        /// Threads per block the launch site asked for.
67        requested: u32,
68        /// What the compiled kernel supports.
69        kernel_max: u32,
70    },
71    /// Block size exceeds the device maximum.
72    #[error("block size {requested} exceeds the device's max_threads_per_block {device_max}")]
73    BlockExceedsDeviceMax {
74        /// Threads per block the launch site asked for.
75        requested: u32,
76        /// What the device supports.
77        device_max: u32,
78    },
79    /// Static + dynamic shared memory exceeds the per-block limit.
80    #[error("shared memory {static_bytes}+{dynamic_bytes}={total} bytes exceeds the device's max_shared_per_block {device_max}")]
81    SharedMemoryExceeded {
82        /// Statically declared bytes.
83        static_bytes: u32,
84        /// Dynamically requested bytes.
85        dynamic_bytes: u32,
86        /// Their sum.
87        total: u32,
88        /// Device per-block limit.
89        device_max: u32,
90    },
91    /// Registers for the whole block exceed the per-block register file.
92    #[error("register budget {num_regs}regs x {block_size}threads = {total} exceeds the device's max_regs_per_block {device_max}")]
93    RegisterBudgetExceeded {
94        /// Registers per thread.
95        num_regs: u32,
96        /// Threads per block.
97        block_size: u32,
98        /// Their product.
99        total: u32,
100        /// Device per-block limit.
101        device_max: u32,
102    },
103    /// A zero-sized launch. Always a bug at the call site, never a legal request.
104    #[error("block size is zero")]
105    ZeroBlockSize,
106}
107
108/// Validate one launch configuration against a compiled kernel and its device.
109///
110/// This is the `register_budget` equation, executable. It is **pure** — no CUDA, no global
111/// state — so the case table below runs anywhere, including the CPU-only required check.
112///
113/// # Errors
114///
115/// Returns the first [`LaunchBudgetViolation`] found. Checks are ordered cheapest-first and
116/// most-specific-first, so `BlockExceedsKernelMax` (the actionable one, naming the kernel's
117/// own ceiling) is reported ahead of the device-wide limit.
118pub fn validate_launch(
119    attrs: &KernelAttributes,
120    limits: &DeviceLimits,
121    block_size: u32,
122    dynamic_shared_bytes: u32,
123) -> Result<(), LaunchBudgetViolation> {
124    if block_size == 0 {
125        return Err(LaunchBudgetViolation::ZeroBlockSize);
126    }
127    if block_size > attrs.max_threads_per_block {
128        return Err(LaunchBudgetViolation::BlockExceedsKernelMax {
129            requested: block_size,
130            kernel_max: attrs.max_threads_per_block,
131        });
132    }
133    if block_size > limits.max_threads_per_block {
134        return Err(LaunchBudgetViolation::BlockExceedsDeviceMax {
135            requested: block_size,
136            device_max: limits.max_threads_per_block,
137        });
138    }
139    let total_shared = attrs
140        .static_shared_bytes
141        .saturating_add(dynamic_shared_bytes);
142    if total_shared > limits.max_shared_per_block {
143        return Err(LaunchBudgetViolation::SharedMemoryExceeded {
144            static_bytes: attrs.static_shared_bytes,
145            dynamic_bytes: dynamic_shared_bytes,
146            total: total_shared,
147            device_max: limits.max_shared_per_block,
148        });
149    }
150    let total_regs = attrs.num_regs.saturating_mul(block_size);
151    if total_regs > limits.max_regs_per_block {
152        return Err(LaunchBudgetViolation::RegisterBudgetExceeded {
153            num_regs: attrs.num_regs,
154            block_size,
155            total: total_regs,
156            device_max: limits.max_regs_per_block,
157        });
158    }
159    Ok(())
160}
161
162/// Whether a kernel spilled registers to local memory.
163///
164/// Not a violation — a spilled kernel is correct, just slow — so this is reported
165/// separately from [`validate_launch`] rather than failing it. FALSIFY-PTX-003's
166/// "register spilling causes performance cliff" half.
167#[must_use]
168pub fn spills_to_local(attrs: &KernelAttributes) -> bool {
169    attrs.local_bytes > 0
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    /// A kernel that comfortably fits: the shape `probe.ptx` reported on GB10.
177    fn ok_kernel() -> KernelAttributes {
178        KernelAttributes {
179            max_threads_per_block: 1024,
180            num_regs: 10,
181            static_shared_bytes: 0,
182            local_bytes: 0,
183        }
184    }
185
186    /// Limits as an sm_89/sm_121-class device reports them.
187    fn limits() -> DeviceLimits {
188        DeviceLimits {
189            max_threads_per_block: 1024,
190            max_shared_per_block: 49152,
191            max_regs_per_block: 65536,
192            warp_size: 32,
193        }
194    }
195
196    #[test]
197    fn accepts_a_legal_launch() {
198        assert!(validate_launch(&ok_kernel(), &limits(), 256, 0).is_ok());
199    }
200
201    #[test]
202    fn accepts_the_kernel_max_exactly() {
203        // Boundary: block_size == max_threads_per_block is legal, not off-by-one.
204        assert!(validate_launch(&ok_kernel(), &limits(), 1024, 0).is_ok());
205    }
206
207    #[test]
208    fn rejects_zero_block_size() {
209        assert_eq!(
210            validate_launch(&ok_kernel(), &limits(), 0, 0),
211            Err(LaunchBudgetViolation::ZeroBlockSize)
212        );
213    }
214
215    #[test]
216    fn rejects_block_above_kernel_max() {
217        // The real GH-613 shape: a register-heavy kernel caps below the device max, and a
218        // hardcoded block=256 launch site walks straight into CUDA_ERROR_INVALID_VALUE.
219        let heavy = KernelAttributes {
220            max_threads_per_block: 128,
221            num_regs: 200,
222            static_shared_bytes: 0,
223            local_bytes: 0,
224        };
225        assert_eq!(
226            validate_launch(&heavy, &limits(), 256, 0),
227            Err(LaunchBudgetViolation::BlockExceedsKernelMax {
228                requested: 256,
229                kernel_max: 128,
230            })
231        );
232    }
233
234    #[test]
235    fn kernel_max_is_reported_before_device_max() {
236        // Both are violated; the kernel's own ceiling is the actionable one.
237        let heavy = KernelAttributes {
238            max_threads_per_block: 128,
239            ..ok_kernel()
240        };
241        let tight = DeviceLimits {
242            max_threads_per_block: 512,
243            ..limits()
244        };
245        assert!(matches!(
246            validate_launch(&heavy, &tight, 2048, 0),
247            Err(LaunchBudgetViolation::BlockExceedsKernelMax { .. })
248        ));
249    }
250
251    #[test]
252    fn rejects_block_above_device_max() {
253        let permissive = KernelAttributes {
254            max_threads_per_block: 4096,
255            ..ok_kernel()
256        };
257        assert_eq!(
258            validate_launch(&permissive, &limits(), 2048, 0),
259            Err(LaunchBudgetViolation::BlockExceedsDeviceMax {
260                requested: 2048,
261                device_max: 1024,
262            })
263        );
264    }
265
266    #[test]
267    fn rejects_shared_memory_overflow_counting_both_halves() {
268        // static alone fits and dynamic alone fits; only the SUM violates.
269        let k = KernelAttributes {
270            static_shared_bytes: 32768,
271            ..ok_kernel()
272        };
273        assert_eq!(
274            validate_launch(&k, &limits(), 256, 32768),
275            Err(LaunchBudgetViolation::SharedMemoryExceeded {
276                static_bytes: 32768,
277                dynamic_bytes: 32768,
278                total: 65536,
279                device_max: 49152,
280            })
281        );
282    }
283
284    #[test]
285    fn accepts_shared_memory_exactly_at_the_limit() {
286        let k = KernelAttributes {
287            static_shared_bytes: 49152,
288            ..ok_kernel()
289        };
290        assert!(validate_launch(&k, &limits(), 256, 0).is_ok());
291    }
292
293    #[test]
294    fn rejects_register_budget_overflow() {
295        // 255 regs x 1024 threads = 261_120 > 65_536.
296        let k = KernelAttributes {
297            max_threads_per_block: 1024,
298            num_regs: 255,
299            ..ok_kernel()
300        };
301        assert_eq!(
302            validate_launch(&k, &limits(), 1024, 0),
303            Err(LaunchBudgetViolation::RegisterBudgetExceeded {
304                num_regs: 255,
305                block_size: 1024,
306                total: 261_120,
307                device_max: 65536,
308            })
309        );
310    }
311
312    #[test]
313    fn register_product_saturates_instead_of_overflowing() {
314        // u32 multiply of two large values must not wrap into a FALSE PASS.
315        let k = KernelAttributes {
316            max_threads_per_block: u32::MAX,
317            num_regs: u32::MAX,
318            ..ok_kernel()
319        };
320        let wide = DeviceLimits {
321            max_threads_per_block: u32::MAX,
322            ..limits()
323        };
324        assert!(matches!(
325            validate_launch(&k, &wide, u32::MAX, 0),
326            Err(LaunchBudgetViolation::RegisterBudgetExceeded { .. })
327        ));
328    }
329
330    #[test]
331    fn shared_memory_sum_saturates_instead_of_overflowing() {
332        let k = KernelAttributes {
333            static_shared_bytes: u32::MAX,
334            ..ok_kernel()
335        };
336        assert!(matches!(
337            validate_launch(&k, &limits(), 256, u32::MAX),
338            Err(LaunchBudgetViolation::SharedMemoryExceeded { .. })
339        ));
340    }
341
342    #[test]
343    fn spill_detection_is_separate_from_validity() {
344        let spilled = KernelAttributes {
345            local_bytes: 128,
346            ..ok_kernel()
347        };
348        // Spilling is a performance cliff, not an illegal launch.
349        assert!(validate_launch(&spilled, &limits(), 256, 0).is_ok());
350        assert!(spills_to_local(&spilled));
351        assert!(!spills_to_local(&ok_kernel()));
352    }
353
354    #[test]
355    fn violations_name_both_sides_of_the_comparison() {
356        // A violation you cannot act on is a log line, not a gate.
357        let heavy = KernelAttributes {
358            max_threads_per_block: 128,
359            ..ok_kernel()
360        };
361        let Err(v) = validate_launch(&heavy, &limits(), 256, 0) else {
362            panic!("expected a violation");
363        };
364        let msg = v.to_string();
365        assert!(msg.contains("256"), "message must name the request: {msg}");
366        assert!(msg.contains("128"), "message must name the limit: {msg}");
367    }
368}
trueno_gpu/launch_budget.rs

trueno_gpu/
launch_budget.rs
