1use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5use std::sync::{Mutex, OnceLock};
6
7use vyre_foundation::ir::{MemoryKind, Node, Program};
8
9use crate::binding::Binding;
10use crate::program_walks::{
11 dispatch_element_count_for_program, infer_dispatch_grid_for_count,
12 program_uses_launch_geometry_ids, try_dispatch_param_words_into,
13};
14use crate::tuner::{
15 identity_fisher_q16, Mode, NaturalGradientPolicy, Tuner, TunerCache, TuningMeasurement,
16 WORKGROUP_CANDIDATES,
17};
18use crate::validation::{validate_launch_geometry, LaunchGeometryLimits};
19use crate::{BackendError, DispatchConfig};
20
21const COLD_START_GRID_STEP_NS: u64 = 1_024;
22const COLD_START_IDLE_LANE_NS: u64 = 8;
23const COLD_START_TEMPERATURE_NS: u64 = 4_096;
24const MAX_NATURAL_LAUNCH_CACHE_ENTRIES: usize = 4_096;
25
26static NATURAL_LAUNCH_CACHE: OnceLock<Mutex<BTreeMap<NaturalLaunchCacheKey, NaturalLaunchEntry>>> =
27 OnceLock::new();
28
29#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct LaunchPlan {
32 pub element_count: u32,
34 pub workgroup: [u32; 3],
36 pub grid: [u32; 3],
38 pub param_words: Vec<u32>,
40 pub max_binding_alignment: usize,
45}
46
47impl LaunchPlan {
48 #[must_use]
50 pub fn new() -> Self {
51 Self {
52 element_count: 1,
53 workgroup: [1, 1, 1],
54 grid: [1, 1, 1],
55 param_words: Vec::new(),
56 max_binding_alignment: 1,
57 }
58 }
59
60 pub fn from_bindings(
67 program: &Program,
68 bindings: &[Binding],
69 config: &DispatchConfig,
70 limits: LaunchGeometryLimits,
71 ) -> Result<Self, BackendError> {
72 let mut plan = Self::new();
73 plan.prepare_into(program, bindings, config, limits)?;
74 Ok(plan)
75 }
76
77 pub fn prepare_into(
84 &mut self,
85 program: &Program,
86 bindings: &[Binding],
87 config: &DispatchConfig,
88 limits: LaunchGeometryLimits,
89 ) -> Result<(), BackendError> {
90 self.prepare_into_for_mode(program, bindings, config, limits, Mode::from_env())
91 }
92
93 fn prepare_into_for_mode(
94 &mut self,
95 program: &Program,
96 bindings: &[Binding],
97 config: &DispatchConfig,
98 limits: LaunchGeometryLimits,
99 mode: Mode,
100 ) -> Result<(), BackendError> {
101 let workgroup =
102 effective_launch_workgroup_for_mode(program, bindings, config, limits, mode);
103 validate_launch_geometry(workgroup, [1, 1, 1], limits)?;
104 let element_count = launch_element_count(program, bindings, workgroup, config, limits)?;
105 let grid = match config.grid_override {
106 Some(grid) => grid,
107 None => {
108 if workgroup[1] != 1 || workgroup[2] != 1 {
114 return Err(BackendError::InvalidProgram {
115 fix: format!(
116 "Fix: backend `{}` requires DispatchConfig::grid_override for non-1D workgroups. \
117 workgroup={:?} has no unambiguous default grid; set grid_override to the logical [x, y, z] you want.",
118 limits.backend, workgroup,
119 ),
120 });
121 }
122 infer_dispatch_grid_for_count(element_count, workgroup)?
123 }
124 };
125 validate_launch_geometry(workgroup, grid, limits)?;
126 self.element_count = element_count;
127 self.workgroup = workgroup;
128 self.grid = grid;
129 self.max_binding_alignment = bindings
130 .iter()
131 .map(|binding| binding.preferred_alignment)
132 .max()
133 .unwrap_or(1);
134 try_dispatch_param_words_into(bindings, element_count, &mut self.param_words).map_err(
135 |error| BackendError::InvalidProgram {
136 fix: format!(
137 "Fix: {}: dispatch ABI parameter staging failed: {error}",
138 limits.backend
139 ),
140 },
141 )?;
142 Ok(())
143 }
144}
145
146impl Default for LaunchPlan {
147 fn default() -> Self {
148 Self::new()
149 }
150}
151
152fn launch_element_count(
153 program: &Program,
154 bindings: &[Binding],
155 workgroup: [u32; 3],
156 config: &DispatchConfig,
157 limits: LaunchGeometryLimits,
158) -> Result<u32, BackendError> {
159 let inferred = dispatch_element_count_for_program(program, bindings);
160 let Some(grid) = config.grid_override else {
161 return Ok(inferred);
162 };
163 if workgroup.contains(&0) || grid.contains(&0) {
164 return Err(BackendError::InvalidProgram {
165 fix: format!(
166 "Fix: {} grid_override and workgroup dimensions must all be non-zero.",
167 limits.backend
168 ),
169 });
170 }
171 grid[0]
172 .checked_mul(workgroup[0])
173 .filter(|count| *count != 0)
174 .ok_or_else(|| BackendError::InvalidProgram {
175 fix: format!(
176 "Fix: {} grid_override.x * workgroup_size.x must fit in u32.",
177 limits.backend
178 ),
179 })
180}
181
182fn effective_launch_workgroup_for_mode(
183 program: &Program,
184 bindings: &[Binding],
185 config: &DispatchConfig,
186 limits: LaunchGeometryLimits,
187 mode: Mode,
188) -> [u32; 3] {
189 let element_count = dispatch_element_count_for_program(program, bindings);
190 resolve_launch_workgroup_for_mode(program, config, limits, element_count, mode)
191}
192
193#[must_use]
200pub fn resolve_launch_workgroup(
201 program: &Program,
202 config: &DispatchConfig,
203 limits: LaunchGeometryLimits,
204 element_count: u32,
205) -> [u32; 3] {
206 resolve_launch_workgroup_for_mode(program, config, limits, element_count, Mode::from_env())
207}
208
209#[must_use]
214pub fn resolve_launch_workgroup_for_mode(
215 program: &Program,
216 config: &DispatchConfig,
217 limits: LaunchGeometryLimits,
218 element_count: u32,
219 mode: Mode,
220) -> [u32; 3] {
221 if let Some(workgroup) = config.workgroup_override {
222 return workgroup;
223 }
224 let declared = program.workgroup_size();
225 if mode != Mode::NaturalGradient || config.grid_override.is_some() {
226 return declared;
227 }
228 natural_gradient_cold_start_workgroup(program, declared, element_count, limits)
229 .unwrap_or(declared)
230}
231
232#[must_use]
241pub fn record_launch_measurement(
242 program: &Program,
243 config: &DispatchConfig,
244 limits: LaunchGeometryLimits,
245 element_count: u32,
246 observed_workgroup: [u32; 3],
247 elapsed_ns: u64,
248) -> bool {
249 record_launch_measurement_for_mode(
250 program,
251 config,
252 limits,
253 element_count,
254 observed_workgroup,
255 elapsed_ns,
256 Mode::from_env(),
257 )
258}
259
260fn record_launch_measurement_for_mode(
261 program: &Program,
262 config: &DispatchConfig,
263 limits: LaunchGeometryLimits,
264 element_count: u32,
265 observed_workgroup: [u32; 3],
266 elapsed_ns: u64,
267 mode: Mode,
268) -> bool {
269 record_launch_measurement_for_mode_with_store(
270 program,
271 config,
272 limits,
273 element_count,
274 observed_workgroup,
275 elapsed_ns,
276 mode,
277 None,
278 )
279}
280
281fn record_launch_measurement_for_mode_with_store(
282 program: &Program,
283 config: &DispatchConfig,
284 limits: LaunchGeometryLimits,
285 element_count: u32,
286 observed_workgroup: [u32; 3],
287 elapsed_ns: u64,
288 mode: Mode,
289 persistent_path: Option<&Path>,
290) -> bool {
291 if mode != Mode::NaturalGradient
292 || elapsed_ns == 0
293 || config.workgroup_override.is_some()
294 || config.grid_override.is_some()
295 || observed_workgroup[1] != 1
296 || observed_workgroup[2] != 1
297 || !candidate_x_fits_limits(observed_workgroup[0], limits)
298 {
299 return false;
300 }
301 let declared = program.workgroup_size();
302 if !is_natural_gradient_launch_tunable(program, declared, element_count) {
303 return false;
304 }
305 let cache_key = NaturalLaunchCacheKey::new(program, declared, element_count, limits);
306 let mut measurements = natural_launch_cache_measurements(cache_key).unwrap_or_default();
307 measurements
308 .entry(observed_workgroup)
309 .and_modify(|best_ns| *best_ns = (*best_ns).min(elapsed_ns))
310 .or_insert(elapsed_ns);
311 let Some(selected) =
312 select_natural_launch_workgroup(declared, element_count, limits, Some(&measurements))
313 else {
314 return false;
315 };
316 natural_launch_cache_set(
317 cache_key,
318 NaturalLaunchEntry {
319 selected,
320 measurements,
321 },
322 );
323 if let Err(error) =
324 persist_natural_launch_selection(cache_key, limits, selected, persistent_path)
325 {
326 tracing::debug!(
327 error,
328 "natural-gradient launch feedback accepted in memory but could not persist"
329 );
330 }
331 true
332}
333
334fn natural_gradient_cold_start_workgroup(
335 program: &Program,
336 declared: [u32; 3],
337 element_count: u32,
338 limits: LaunchGeometryLimits,
339) -> Option<[u32; 3]> {
340 natural_gradient_cold_start_workgroup_with_store(program, declared, element_count, limits, None)
341}
342
343fn natural_gradient_cold_start_workgroup_with_store(
344 program: &Program,
345 declared: [u32; 3],
346 element_count: u32,
347 limits: LaunchGeometryLimits,
348 persistent_path: Option<&Path>,
349) -> Option<[u32; 3]> {
350 if !is_natural_gradient_launch_tunable(program, declared, element_count) {
351 return None;
352 }
353 let cache_key = NaturalLaunchCacheKey::new(program, declared, element_count, limits);
354 if let Some(cached) = natural_launch_cache_get(cache_key) {
355 return Some(cached);
356 }
357 if let Some(persisted) = natural_launch_cache_get_persisted(cache_key, limits, persistent_path)
358 {
359 natural_launch_cache_set(
360 cache_key,
361 NaturalLaunchEntry {
362 selected: persisted,
363 measurements: BTreeMap::new(),
364 },
365 );
366 return Some(persisted);
367 }
368
369 let selected = select_natural_launch_workgroup(declared, element_count, limits, None)?;
370 natural_launch_cache_set(
371 cache_key,
372 NaturalLaunchEntry {
373 selected,
374 measurements: BTreeMap::new(),
375 },
376 );
377 Some(selected)
378}
379
380fn select_natural_launch_workgroup(
381 declared: [u32; 3],
382 element_count: u32,
383 limits: LaunchGeometryLimits,
384 measurements: Option<&BTreeMap<[u32; 3], u64>>,
385) -> Option<[u32; 3]> {
386 let peak_resident = peak_resident_threads_per_compute_unit(declared[0], limits);
387 let mut samples = Vec::with_capacity(WORKGROUP_CANDIDATES.len() + 1);
388 for candidate_x in WORKGROUP_CANDIDATES
389 .iter()
390 .copied()
391 .chain(std::iter::once(declared[0]))
392 {
393 if !candidate_x_fits_limits(candidate_x, limits)
394 || samples
395 .iter()
396 .any(|sample: &TuningMeasurement| sample.workgroup_size[0] == candidate_x)
397 {
398 continue;
399 }
400 let workgroup_size = [candidate_x, 1, 1];
401 let elapsed_ns = match measurements.and_then(|measured| measured.get(&workgroup_size)) {
402 Some(&measured_ns) => measured_ns,
403 None if cold_start_admits_width(candidate_x, limits, peak_resident) => {
404 estimate_cold_start_latency_ns(element_count, candidate_x)
405 }
406 None => continue,
407 };
408 samples.push(TuningMeasurement {
409 workgroup_size,
410 elapsed_ns,
411 });
412 }
413 if let Some(measured) = measurements {
414 for (&workgroup_size, &elapsed_ns) in measured {
415 if workgroup_size[1] != 1
416 || workgroup_size[2] != 1
417 || elapsed_ns == 0
418 || !candidate_x_fits_limits(workgroup_size[0], limits)
419 || samples
420 .iter()
421 .any(|sample| sample.workgroup_size == workgroup_size)
422 {
423 continue;
424 }
425 samples.push(TuningMeasurement {
426 workgroup_size,
427 elapsed_ns,
428 });
429 }
430 }
431
432 if samples.len() < 2 {
433 return None;
434 }
435 NaturalGradientPolicy {
436 temperature_ns: COLD_START_TEMPERATURE_NS,
437 }
438 .suggest(&samples, &identity_fisher_q16(samples.len()))
439 .ok()
440 .map(|step| step.selected_workgroup_size)
441}
442
443#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
444struct NaturalLaunchCacheKey {
445 fingerprint: [u8; 32],
446 declared: [u32; 3],
447 element_count: u32,
448 max_threads_per_block: u32,
449 max_block_dim: [u32; 3],
450 max_grid_dim: [u32; 3],
451 max_threads_per_sm: u32,
452}
453
454impl NaturalLaunchCacheKey {
455 fn new(
456 program: &Program,
457 declared: [u32; 3],
458 element_count: u32,
459 limits: LaunchGeometryLimits,
460 ) -> Self {
461 Self {
462 fingerprint: program.fingerprint(),
463 declared,
464 element_count,
465 max_threads_per_block: limits.max_threads_per_block,
466 max_block_dim: limits.max_block_dim,
467 max_grid_dim: limits.max_grid_dim,
468 max_threads_per_sm: limits.max_threads_per_sm,
469 }
470 }
471
472 fn persistent_key(self) -> String {
473 let mut hasher = blake3::Hasher::new();
474 hasher.update(b"vyre-natural-launch-feedback-v2\0");
478 hasher.update(&self.fingerprint);
479 for axis in self.declared {
480 hasher.update(&axis.to_le_bytes());
481 }
482 hasher.update(&self.element_count.to_le_bytes());
483 hasher.update(&self.max_threads_per_block.to_le_bytes());
484 for axis in self.max_block_dim {
485 hasher.update(&axis.to_le_bytes());
486 }
487 for axis in self.max_grid_dim {
488 hasher.update(&axis.to_le_bytes());
489 }
490 hasher.update(&self.max_threads_per_sm.to_le_bytes());
491 let digest = hasher.finalize();
492 let mut key = String::with_capacity(74);
493 key.push_str("launch-v2-");
494 crate::pipeline::hashing::push_lower_hex(digest.as_bytes(), &mut key);
495 key
496 }
497}
498
499#[derive(Clone, Debug, Eq, PartialEq)]
500
501struct NaturalLaunchEntry {
502 selected: [u32; 3],
503 measurements: BTreeMap<[u32; 3], u64>,
504}
505
506fn natural_launch_cache_get(key: NaturalLaunchCacheKey) -> Option<[u32; 3]> {
507 let cache = NATURAL_LAUNCH_CACHE.get_or_init(|| Mutex::new(BTreeMap::new()));
508 let guard = cache.lock().unwrap_or_else(|poison| poison.into_inner());
509 guard.get(&key).map(|entry| entry.selected)
510}
511
512fn natural_launch_cache_measurements(
513 key: NaturalLaunchCacheKey,
514) -> Option<BTreeMap<[u32; 3], u64>> {
515 let cache = NATURAL_LAUNCH_CACHE.get_or_init(|| Mutex::new(BTreeMap::new()));
516 let guard = cache.lock().unwrap_or_else(|poison| poison.into_inner());
517 guard.get(&key).map(|entry| entry.measurements.clone())
518}
519
520fn natural_launch_cache_set(key: NaturalLaunchCacheKey, value: NaturalLaunchEntry) {
521 let cache = NATURAL_LAUNCH_CACHE.get_or_init(|| Mutex::new(BTreeMap::new()));
522 let mut guard = cache.lock().unwrap_or_else(|poison| poison.into_inner());
523 if guard.len() >= MAX_NATURAL_LAUNCH_CACHE_ENTRIES && !guard.contains_key(&key) {
524 if let Some(oldest) = guard.keys().next().copied() {
525 guard.remove(&oldest);
526 }
527 }
528 guard.insert(key, value);
529}
530
531#[cfg(test)]
532fn natural_launch_cache_remove(key: NaturalLaunchCacheKey) {
533 if let Some(cache) = NATURAL_LAUNCH_CACHE.get() {
534 if let Ok(mut guard) = cache.lock() {
535 guard.remove(&key);
536 }
537 }
538}
539
540fn natural_launch_cache_get_persisted(
541 key: NaturalLaunchCacheKey,
542 limits: LaunchGeometryLimits,
543 persistent_path: Option<&Path>,
544) -> Option<[u32; 3]> {
545 let path = persistent_path
546 .map(Path::to_path_buf)
547 .unwrap_or_else(|| natural_launch_persistent_cache_path(limits));
548 let selected = TunerCache::load(&path).ok()?.get(&key.persistent_key())?;
549 valid_persisted_launch_selection(selected, limits).then_some(selected)
550}
551
552fn persist_natural_launch_selection(
553 key: NaturalLaunchCacheKey,
554 limits: LaunchGeometryLimits,
555 selected: [u32; 3],
556 persistent_path: Option<&Path>,
557) -> Result<(), String> {
558 let path = persistent_path
559 .map(Path::to_path_buf)
560 .unwrap_or_else(|| natural_launch_persistent_cache_path(limits));
561 persist_natural_launch_selection_to_path(key, selected, &path)
562}
563
564fn persist_natural_launch_selection_to_path(
565 key: NaturalLaunchCacheKey,
566 selected: [u32; 3],
567 path: &Path,
568) -> Result<(), String> {
569 let mut cache = TunerCache::load(path)?;
570 while cache.entries.len() >= MAX_NATURAL_LAUNCH_CACHE_ENTRIES {
571 let Some(oldest) = cache.entries.keys().next().cloned() else {
572 break;
573 };
574 cache.entries.remove(&oldest);
575 }
576 cache.set(key.persistent_key(), selected);
577 cache.save(path)
578}
579
580fn natural_launch_persistent_cache_path(limits: LaunchGeometryLimits) -> PathBuf {
581 Tuner::cache_path_for_adapter(&natural_launch_persistent_adapter_key(limits))
582}
583
584fn natural_launch_persistent_adapter_key(limits: LaunchGeometryLimits) -> String {
585 let mut hasher = blake3::Hasher::new();
586 hasher.update(b"vyre-natural-launch-adapter-v2\0");
587 hasher.update(limits.backend.as_bytes());
588 hasher.update(&limits.max_threads_per_block.to_le_bytes());
589 for axis in limits.max_block_dim {
590 hasher.update(&axis.to_le_bytes());
591 }
592 for axis in limits.max_grid_dim {
593 hasher.update(&axis.to_le_bytes());
594 }
595 hasher.update(&limits.max_threads_per_sm.to_le_bytes());
596 let digest = hasher.finalize();
597 let mut key = String::with_capacity(92);
598 key.push_str("natural-launch-feedback-v2-");
599 crate::pipeline::hashing::push_lower_hex(digest.as_bytes(), &mut key);
600 key
601}
602
603fn valid_persisted_launch_selection(selected: [u32; 3], limits: LaunchGeometryLimits) -> bool {
604 selected[1] == 1 && selected[2] == 1 && candidate_x_fits_limits(selected[0], limits)
605}
606
607fn is_natural_gradient_launch_tunable(
608 program: &Program,
609 declared: [u32; 3],
610 element_count: u32,
611) -> bool {
612 declared[0] != 0
613 && declared[1] == 1
614 && declared[2] == 1
615 && element_count != 0
616 && program
617 .entry
618 .iter()
619 .any(|node| !matches!(node, Node::Return))
620 && !program.non_composable_with_self
621 && !program_uses_launch_geometry_ids(program)
622 && program
623 .buffers
624 .iter()
625 .all(|buffer| buffer.kind() != MemoryKind::Shared)
626}
627
628fn candidate_x_fits_limits(candidate_x: u32, limits: LaunchGeometryLimits) -> bool {
629 candidate_x != 0
630 && candidate_x <= limits.max_threads_per_block
631 && candidate_x <= limits.max_block_dim[0]
632}
633
634fn peak_resident_threads_per_compute_unit(
642 declared_x: u32,
643 limits: LaunchGeometryLimits,
644) -> Option<u32> {
645 WORKGROUP_CANDIDATES
646 .iter()
647 .copied()
648 .chain(std::iter::once(declared_x))
649 .filter(|&candidate_x| candidate_x_fits_limits(candidate_x, limits))
650 .filter_map(|candidate_x| limits.resident_threads_per_compute_unit(candidate_x))
651 .max()
652}
653
654fn cold_start_admits_width(
679 candidate_x: u32,
680 limits: LaunchGeometryLimits,
681 peak_resident: Option<u32>,
682) -> bool {
683 let Some(peak_resident) = peak_resident else {
684 return true;
685 };
686 limits
687 .resident_threads_per_compute_unit(candidate_x)
688 .is_none_or(|resident| resident >= peak_resident)
689}
690
691fn estimate_cold_start_latency_ns(element_count: u32, candidate_x: u32) -> u64 {
692 let groups = u64::from(element_count.div_ceil(candidate_x));
693 let scheduled_lanes = groups.saturating_mul(u64::from(candidate_x));
694 let idle_lanes = scheduled_lanes.saturating_sub(u64::from(element_count));
695 groups
696 .saturating_mul(COLD_START_GRID_STEP_NS)
697 .saturating_add(idle_lanes.saturating_mul(COLD_START_IDLE_LANE_NS))
698}
699
700#[must_use]
702pub fn program_vsa_fingerprint(program: &Program) -> Vec<u32> {
703 program_vsa_fingerprint_words(program).to_vec()
704}
705
706#[must_use]
708pub fn program_vsa_fingerprint_words(program: &Program) -> [u32; 8] {
709 let fingerprint = program.fingerprint();
710 let mut words = [0u32; 8];
711 for (word, chunk) in words.iter_mut().zip(fingerprint.chunks_exact(4)) {
712 *word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
713 }
714 words
715}
716
717#[cfg(test)]
718mod tests {
719 use super::*;
720 use crate::binding::BindingRole;
721 use vyre_foundation::ir::{BufferDecl, DataType, Expr, Node, Program};
722
723 #[test]
724 fn program_vsa_fingerprint_words_match_wire_decoder() {
725 let program = Program::wrapped(vec![], [64, 1, 1], vec![]);
726 let words = program_vsa_fingerprint_words(&program);
727 let fingerprint = program.fingerprint();
728
729 for (index, chunk) in fingerprint.chunks_exact(4).enumerate() {
730 assert_eq!(
731 words[index],
732 u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
733 );
734 }
735 assert_eq!(program_vsa_fingerprint(&program), words.to_vec());
736 }
737
738 #[test]
739 fn launch_plan_prepare_into_reuses_param_words() {
740 let program = Program::wrapped(vec![], [64, 1, 1], vec![]);
741 let bindings = vec![Binding {
742 name: std::sync::Arc::from("input"),
743 binding: 0,
744 buffer_index: 0,
745 role: BindingRole::Input,
746 element_size: 4,
747 preferred_alignment: 64,
748 element_count: 7,
749 static_byte_len: Some(28),
750 input_index: Some(0),
751 output_index: None,
752 }];
753 let limits = LaunchGeometryLimits {
754 backend: "test",
755 max_threads_per_block: 1024,
756 max_block_dim: [1024, 1024, 64],
757 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
758 max_threads_per_sm: 1536,
759 };
760 let mut plan = LaunchPlan {
761 param_words: Vec::with_capacity(8),
762 ..LaunchPlan::new()
763 };
764 let ptr = plan.param_words.as_ptr();
765 plan.prepare_into(&program, &bindings, &DispatchConfig::default(), limits)
766 .unwrap();
767 assert_eq!(plan.element_count, 7);
768 assert_eq!(plan.grid, [1, 1, 1]);
769 assert_eq!(plan.param_words, vec![7, 7]);
770 assert_eq!(plan.max_binding_alignment, 64);
771 assert_eq!(plan.param_words.as_ptr(), ptr);
772 }
773
774 #[test]
775 fn natural_gradient_launch_tunes_safe_1d_storage_program() {
776 let program = Program::wrapped(
777 vec![BufferDecl::output("out", 0, DataType::U32).with_count(4096)],
778 [32, 1, 1],
779 vec![],
780 );
781 let bindings = vec![Binding {
782 name: std::sync::Arc::from("out"),
783 binding: 0,
784 buffer_index: 0,
785 role: BindingRole::Output,
786 element_size: 4,
787 preferred_alignment: 128,
788 element_count: 4096,
789 static_byte_len: Some(16_384),
790 input_index: None,
791 output_index: Some(0),
792 }];
793 let limits = LaunchGeometryLimits {
794 backend: "test",
795 max_threads_per_block: 1024,
796 max_block_dim: [1024, 1024, 64],
797 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
798 max_threads_per_sm: 1536,
799 };
800 let mut plan = LaunchPlan::new();
801
802 plan.prepare_into_for_mode(
803 &program,
804 &bindings,
805 &DispatchConfig::default(),
806 limits,
807 Mode::NaturalGradient,
808 )
809 .expect("Fix: safe 1D storage launch should accept natural-gradient cold start");
810
811 assert_eq!(
812 plan.workgroup,
813 [512, 1, 1],
814 "Fix: cold start must pick the widest width that keeps every resident thread slot usable. Was [1024,1,1], which is 1536/1024 = 1 block per SM and 512 stranded slots on every SM."
815 );
816 assert_eq!(
817 limits.resident_threads_per_compute_unit(plan.workgroup[0]),
818 Some(1536),
819 "Fix: the chosen width must strand no per-SM thread slot when a candidate dividing 1536 evenly exists."
820 );
821 assert_eq!(plan.grid, [8, 1, 1]);
822 assert_eq!(plan.element_count, 4096);
823 }
824
825 #[test]
826 fn natural_gradient_launch_preserves_declared_shape_for_local_workgroup_ids() {
827 let program = Program::wrapped(
828 vec![BufferDecl::output("out_local_ids", 0, DataType::U32).with_count(4096)],
829 [1024, 1, 1],
830 vec![
831 Node::let_bind("lane", Expr::LocalId { axis: 0 }),
832 Node::let_bind("block", Expr::WorkgroupId { axis: 0 }),
833 Node::let_bind(
834 "global",
835 Expr::add(
836 Expr::mul(Expr::var("block"), Expr::u32(1024)),
837 Expr::var("lane"),
838 ),
839 ),
840 Node::store("out_local_ids", Expr::var("global"), Expr::var("lane")),
841 ],
842 );
843 let bindings = vec![Binding {
844 name: std::sync::Arc::from("out_local_ids"),
845 binding: 0,
846 buffer_index: 0,
847 role: BindingRole::Output,
848 element_size: 4,
849 preferred_alignment: 128,
850 element_count: 4096,
851 static_byte_len: Some(16_384),
852 input_index: None,
853 output_index: Some(0),
854 }];
855 let limits = LaunchGeometryLimits {
856 backend: "test",
857 max_threads_per_block: 1024,
858 max_block_dim: [1024, 1024, 64],
859 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
860 max_threads_per_sm: 1536,
861 };
862
863 assert_eq!(
864 effective_launch_workgroup_for_mode(
865 &program,
866 &bindings,
867 &DispatchConfig::default(),
868 limits,
869 Mode::NaturalGradient,
870 ),
871 [1024, 1, 1],
872 "Fix: automatic launch tuning must not change kernels whose LocalId/WorkgroupId arithmetic makes workgroup shape semantic."
873 );
874 }
875
876 #[test]
877 fn measured_launch_feedback_overrides_heuristic_cold_start() {
878 let dir = tempfile::tempdir()
879 .expect("Fix: measured launch feedback test needs an isolated tuner cache");
880 let path = dir.path().join("launch-feedback.toml");
881 let program = Program::wrapped(
882 vec![BufferDecl::output("out_feedback_isolated", 0, DataType::U32).with_count(8192)],
883 [32, 1, 1],
884 vec![],
885 );
886 let config = DispatchConfig::default();
887 let limits = LaunchGeometryLimits {
888 backend: "test",
889 max_threads_per_block: 1024,
890 max_block_dim: [1024, 1024, 64],
891 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
892 max_threads_per_sm: 1536,
893 };
894 let key = NaturalLaunchCacheKey::new(&program, [32, 1, 1], 8192, limits);
895 natural_launch_cache_remove(key);
896
897 assert_eq!(
898 natural_gradient_cold_start_workgroup_with_store(
899 &program,
900 [32, 1, 1],
901 8192,
902 limits,
903 Some(&path),
904 ),
905 Some([512, 1, 1]),
906 "Fix: this pins the cold-start selector's output, not a required constant. It was [1024,1,1] under a heuristic with no occupancy term at all, so the old message's claim of an occupancy-efficient shape described the opposite of what it selected."
907 );
908 assert!(
909 record_launch_measurement_for_mode_with_store(
910 &program,
911 &config,
912 limits,
913 8192,
914 [64, 1, 1],
915 1,
916 Mode::NaturalGradient,
917 Some(&path),
918 ),
919 "Fix: natural-gradient resolver must accept measured backend timing for safe 1D launches."
920 );
921 assert_eq!(
922 natural_gradient_cold_start_workgroup_with_store(
923 &program,
924 [32, 1, 1],
925 8192,
926 limits,
927 Some(&path),
928 ),
929 Some([64, 1, 1]),
930 "Fix: measured launch feedback must steer future automatic launch choices."
931 );
932 }
933
934 #[test]
935 fn persisted_launch_feedback_rehydrates_measured_selection() {
936 let dir = tempfile::tempdir()
937 .expect("Fix: launch feedback persistence test needs a temporary cache directory");
938 let path = dir.path().join("launch-feedback.toml");
939 let program = Program::wrapped(
940 vec![BufferDecl::output("out_persisted", 0, DataType::U32).with_count(16_384)],
941 [32, 1, 1],
942 vec![],
943 );
944 let limits = LaunchGeometryLimits {
945 backend: "test",
946 max_threads_per_block: 1024,
947 max_block_dim: [1024, 1024, 64],
948 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
949 max_threads_per_sm: 1536,
950 };
951 let key = NaturalLaunchCacheKey::new(&program, [32, 1, 1], 16_384, limits);
952 natural_launch_cache_remove(key);
953
954 persist_natural_launch_selection_to_path(key, [64, 1, 1], &path)
955 .expect("Fix: measured launch feedback should persist through the tuner cache format");
956
957 assert_eq!(
958 natural_gradient_cold_start_workgroup_with_store(
959 &program,
960 [32, 1, 1],
961 16_384,
962 limits,
963 Some(&path),
964 ),
965 Some([64, 1, 1]),
966 "Fix: automatic launch resolution must rehydrate measured feedback from the bounded tuner cache before falling back to heuristics."
967 );
968 }
969
970 #[test]
971 fn natural_gradient_launch_preserves_explicit_and_shared_memory_shapes() {
972 let program = Program::wrapped(
973 vec![
974 BufferDecl::output("out", 0, DataType::U32).with_count(4096),
975 BufferDecl::workgroup("scratch", 64, DataType::U32),
976 ],
977 [64, 1, 1],
978 vec![],
979 );
980 let bindings = vec![Binding {
981 name: std::sync::Arc::from("out"),
982 binding: 0,
983 buffer_index: 0,
984 role: BindingRole::Output,
985 element_size: 4,
986 preferred_alignment: 128,
987 element_count: 4096,
988 static_byte_len: Some(16_384),
989 input_index: None,
990 output_index: Some(0),
991 }];
992 let limits = LaunchGeometryLimits {
993 backend: "test",
994 max_threads_per_block: 1024,
995 max_block_dim: [1024, 1024, 64],
996 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
997 max_threads_per_sm: 1536,
998 };
999 let mut config = DispatchConfig::default();
1000 config.workgroup_override = Some([256, 1, 1]);
1001
1002 assert_eq!(
1003 effective_launch_workgroup_for_mode(
1004 &program,
1005 &bindings,
1006 &config,
1007 limits,
1008 Mode::NaturalGradient,
1009 ),
1010 [256, 1, 1],
1011 "Fix: explicit dispatch workgroup overrides must remain authoritative."
1012 );
1013
1014 let default_config = DispatchConfig::default();
1015 assert_eq!(
1016 effective_launch_workgroup_for_mode(
1017 &program,
1018 &bindings,
1019 &default_config,
1020 limits,
1021 Mode::NaturalGradient,
1022 ),
1023 [64, 1, 1],
1024 "Fix: workgroup-local scratch kernels must keep their declared shape."
1025 );
1026 }
1027
1028 #[test]
1041 fn record_launch_measurement_starts_fresh_only_when_no_prior_history_exists() {
1042 let dir = tempfile::tempdir()
1043 .expect("Fix: measurement history test needs a temporary cache directory");
1044 let path = dir.path().join("measurements-test.toml");
1045 let program = Program::wrapped(
1046 vec![BufferDecl::output("out_meas_history", 0, DataType::U32).with_count(4096)],
1047 [32, 1, 1],
1048 vec![],
1049 );
1050 let config = DispatchConfig::default();
1051 let limits = LaunchGeometryLimits {
1052 backend: "test-measurements",
1053 max_threads_per_block: 1024,
1054 max_block_dim: [1024, 1024, 64],
1055 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
1056 max_threads_per_sm: 0,
1057 };
1058 let key = NaturalLaunchCacheKey::new(&program, [32, 1, 1], 4096, limits);
1059 natural_launch_cache_remove(key);
1060
1061 assert!(
1063 record_launch_measurement_for_mode_with_store(
1064 &program,
1065 &config,
1066 limits,
1067 4096,
1068 [256, 1, 1],
1069 100,
1070 Mode::NaturalGradient,
1071 Some(&path),
1072 ),
1073 "Fix: first measurement must be accepted into the cache"
1074 );
1075
1076 let after_first = natural_launch_cache_get(key);
1078 assert!(
1079 after_first.is_some(),
1080 "Fix: cache must hold a selection after the first measurement"
1081 );
1082
1083 assert!(
1086 record_launch_measurement_for_mode_with_store(
1087 &program,
1088 &config,
1089 limits,
1090 4096,
1091 [128, 1, 1],
1092 50,
1093 Mode::NaturalGradient,
1094 Some(&path),
1095 ),
1096 "Fix: second measurement must be accepted into the cache"
1097 );
1098
1099 let measurements = natural_launch_cache_measurements(key)
1100 .expect("Fix: cache must hold measurements after two records");
1101 assert!(
1102 measurements.len() >= 2,
1103 "Fix: measurement history must accumulate across calls, got {} entries, expected >= 2",
1104 measurements.len()
1105 );
1106 assert_eq!(
1107 measurements.get(&[256, 1, 1]),
1108 Some(&100),
1109 "Fix: first measurement (workgroup=[256,1,1], 100ns) must be retained in history"
1110 );
1111 assert_eq!(
1112 measurements.get(&[128, 1, 1]),
1113 Some(&50),
1114 "Fix: second measurement (workgroup=[128,1,1], 50ns) must be present in history"
1115 );
1116 }
1117
1118 const RTX_5090_SM_COUNT: u32 = 170;
1120
1121 fn blackwell_5090_limits() -> LaunchGeometryLimits {
1124 LaunchGeometryLimits {
1125 backend: "blackwell-5090-test",
1126 max_threads_per_block: 1024,
1127 max_block_dim: [1024, 1024, 64],
1128 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
1129 max_threads_per_sm: 1536,
1130 }
1131 }
1132
1133 fn tunable_1d_program(output: &'static str, element_count: u32, declared: [u32; 3]) -> Program {
1140 Program::wrapped(
1141 vec![BufferDecl::output(output, 0, DataType::U32).with_count(element_count)],
1142 declared,
1143 vec![],
1144 )
1145 }
1146
1147 #[test]
1158 fn cold_start_never_strands_resident_thread_slots_when_an_even_divisor_exists() {
1159 let limits = blackwell_5090_limits();
1160 for (output, element_count) in [
1161 ("out_no_strand_1k", 1024u32),
1162 ("out_no_strand_4k", 4096),
1163 ("out_no_strand_64k", 65_536),
1164 ("out_no_strand_tail", 4097),
1165 ("out_no_strand_100k", 100_000),
1166 ] {
1167 let program = tunable_1d_program(output, element_count, [32, 1, 1]);
1168 let resolved = resolve_launch_workgroup_for_mode(
1169 &program,
1170 &DispatchConfig::default(),
1171 limits,
1172 element_count,
1173 Mode::NaturalGradient,
1174 );
1175 let resident = limits.resident_threads_per_compute_unit(resolved[0]);
1176 assert_eq!(
1177 resident,
1178 Some(1536),
1179 "Fix: cold start chose {resolved:?} for {element_count} elements, leaving {} of every SM's 1536 thread slots unusable. Prefer a width that divides the per-SM budget evenly.",
1180 1536 - resident.unwrap_or(1536)
1181 );
1182 }
1183 }
1184
1185 #[test]
1192 fn cold_start_still_admits_1024_where_the_per_sm_budget_divides_evenly() {
1193 let limits = LaunchGeometryLimits {
1194 backend: "even-divisor-test",
1195 max_threads_per_block: 1024,
1196 max_block_dim: [1024, 1024, 64],
1197 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
1198 max_threads_per_sm: 2048,
1199 };
1200 let program = tunable_1d_program("out_even_divisor", 65_536, [32, 1, 1]);
1201
1202 assert_eq!(
1203 limits.resident_threads_per_compute_unit(1024),
1204 Some(2048),
1205 "Fix: 1024 must reach every thread slot on a 2048-thread SM, otherwise this test's premise is wrong."
1206 );
1207 assert_eq!(
1208 resolve_launch_workgroup_for_mode(
1209 &program,
1210 &DispatchConfig::default(),
1211 limits,
1212 65_536,
1213 Mode::NaturalGradient,
1214 ),
1215 [1024, 1, 1],
1216 "Fix: residency-aware cold start must stay a residency rule. A width that strands nothing has to remain selectable on every device."
1217 );
1218 }
1219
1220 #[test]
1231 fn unreported_per_sm_budget_leaves_cold_start_byte_identical() {
1232 let limits = LaunchGeometryLimits {
1233 backend: "unreported-residency-test",
1234 max_threads_per_block: 1024,
1235 max_block_dim: [1024, 1024, 64],
1236 max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
1237 max_threads_per_sm: 0,
1238 };
1239 assert_eq!(
1240 limits.resident_threads_per_compute_unit(1024),
1241 None,
1242 "Fix: an unreported per-SM budget must answer `unknown`, never a guessed number."
1243 );
1244
1245 for (output, element_count) in [
1246 ("out_inert_1k", 1024u32),
1247 ("out_inert_4k", 4096),
1248 ("out_inert_64k", 65_536),
1249 ("out_inert_1000", 1_000),
1250 ("out_inert_4097", 4097),
1251 ("out_inert_100k", 100_000),
1252 ] {
1253 let program = tunable_1d_program(output, element_count, [32, 1, 1]);
1254 assert_eq!(
1255 resolve_launch_workgroup_for_mode(
1256 &program,
1257 &DispatchConfig::default(),
1258 limits,
1259 element_count,
1260 Mode::NaturalGradient,
1261 ),
1262 [1024, 1, 1],
1263 "Fix: residency-aware cold start must be inert for a backend that reports no per-SM budget. {element_count} elements resolved differently than they did before residency entered this decision."
1264 );
1265 }
1266 }
1267
1268 #[test]
1275 fn explicit_geometry_pins_outrank_residency_aware_cold_start() {
1276 let limits = blackwell_5090_limits();
1277 let declared = [256, 1, 1];
1278 let program = tunable_1d_program("out_pinned_geometry", 262_144, declared);
1279
1280 let mut pinned_workgroup = DispatchConfig::default();
1281 pinned_workgroup.workgroup_override = Some([64, 1, 1]);
1282 assert_eq!(
1283 resolve_launch_workgroup_for_mode(
1284 &program,
1285 &pinned_workgroup,
1286 limits,
1287 262_144,
1288 Mode::NaturalGradient,
1289 ),
1290 [64, 1, 1],
1291 "Fix: an explicit workgroup override stays authoritative even when residency prefers another width."
1292 );
1293
1294 let mut pinned_grid = DispatchConfig::default();
1295 pinned_grid.grid_override = Some([1024, 1, 1]);
1296 assert_eq!(
1297 resolve_launch_workgroup_for_mode(
1298 &program,
1299 &pinned_grid,
1300 limits,
1301 262_144,
1302 Mode::NaturalGradient,
1303 ),
1304 declared,
1305 "Fix: an explicit grid override must keep the declared workgroup, since the caller sized the grid against it."
1306 );
1307 }
1308
1309 #[test]
1323 fn cooperative_lane_ceiling_follows_the_resolved_width_not_the_declared_one() {
1324 let limits = blackwell_5090_limits();
1325 let declared = [256, 1, 1];
1326 let program = tunable_1d_program("out_resolved_ceiling", 262_144, declared);
1327 let lane_ceiling = |width: u32| -> u64 {
1328 u64::from(
1329 limits
1330 .resident_threads_per_compute_unit(width)
1331 .expect("Fix: this device model reports a per-SM thread budget"),
1332 ) * u64::from(RTX_5090_SM_COUNT)
1333 };
1334
1335 let resolved = resolve_launch_workgroup_for_mode(
1336 &program,
1337 &DispatchConfig::default(),
1338 limits,
1339 262_144,
1340 Mode::NaturalGradient,
1341 );
1342 assert_ne!(
1343 resolved, declared,
1344 "Fix: this program is tunable, so a bound taken from the declared width would bound a width nothing launches."
1345 );
1346 assert_eq!(
1347 lane_ceiling(1024),
1348 174_080,
1349 "Fix: 1024 wide is 1 block/SM x 170 SMs x 1024 lanes. This is the ceiling the defect produced."
1350 );
1351 assert_eq!(
1352 lane_ceiling(resolved[0]),
1353 261_120,
1354 "Fix: the resolved width must reach the device's full cooperative capacity, 1536 resident threads x 170 SMs. Seeing 174,080 here means the tuner resolved 1024 again."
1355 );
1356
1357 let mut pinned = DispatchConfig::default();
1358 pinned.workgroup_override = Some(declared);
1359 assert_eq!(
1360 resolve_launch_workgroup_for_mode(
1361 &program,
1362 &pinned,
1363 limits,
1364 262_144,
1365 Mode::NaturalGradient,
1366 ),
1367 declared,
1368 "Fix: a pinned width must resolve to itself so the declared and resolved bounds coincide."
1369 );
1370 assert_eq!(lane_ceiling(declared[0]), 261_120);
1371 }
1372
1373 #[test]
1379 fn measured_feedback_can_still_select_a_width_cold_start_would_reject() {
1380 let dir =
1381 tempfile::tempdir().expect("Fix: measured feedback test needs an isolated tuner cache");
1382 let path = dir.path().join("residency-feedback.toml");
1383 let limits = blackwell_5090_limits();
1384 let declared = [32, 1, 1];
1385 let program = tunable_1d_program("out_measured_beats_residency", 65_536, declared);
1386 let key = NaturalLaunchCacheKey::new(&program, declared, 65_536, limits);
1387 natural_launch_cache_remove(key);
1388
1389 assert_eq!(
1390 natural_gradient_cold_start_workgroup_with_store(
1391 &program,
1392 declared,
1393 65_536,
1394 limits,
1395 Some(&path),
1396 ),
1397 Some([512, 1, 1]),
1398 "Fix: with no measurements the residency preference decides."
1399 );
1400 natural_launch_cache_remove(key);
1401 assert!(
1402 record_launch_measurement_for_mode_with_store(
1403 &program,
1404 &DispatchConfig::default(),
1405 limits,
1406 65_536,
1407 [1024, 1, 1],
1408 1,
1409 Mode::NaturalGradient,
1410 Some(&path),
1411 ),
1412 "Fix: a real timing for a residency-poor width must still be accepted."
1413 );
1414 assert_eq!(
1415 natural_gradient_cold_start_workgroup_with_store(
1416 &program,
1417 declared,
1418 65_536,
1419 limits,
1420 Some(&path),
1421 ),
1422 Some([1024, 1, 1]),
1423 "Fix: residency governs the cold start only. Measured feedback must remain able to choose a width cold start would never propose."
1424 );
1425 }
1426}