1use super::MegakernelWorkItem;
7
8mod prologue;
9pub use prologue::shared_prologue_length;
10
11pub(super) const MAX_DENSE_FUSION_ITEMS: usize = 4096;
15
16#[derive(Debug, Default)]
21pub struct FusionSelectionScratch {
22 order: Vec<usize>,
23 result: Vec<u32>,
24 conflict_degrees: Vec<u32>,
25 selected: Vec<usize>,
26}
27
28impl FusionSelectionScratch {
29 #[must_use]
31 pub fn result(&self) -> &[u32] {
32 &self.result
33 }
34
35 #[must_use]
37 pub fn take_result(&mut self) -> Vec<u32> {
38 std::mem::take(&mut self.result)
39 }
40
41 fn prepare(&mut self, n: usize) {
42 self.order.clear();
43 self.order.extend(0..n);
44 self.result.clear();
45 self.result.resize(n, 0);
46 self.conflict_degrees.clear();
47 self.conflict_degrees.resize(n, 0);
48 self.selected.clear();
49 }
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
54pub enum FusionSelectionError {
55 ExchangeSizeOverflow {
57 n: usize,
59 },
60 CostLen {
62 expected: usize,
64 actual: usize,
66 },
67 ExchangeAdjLen {
69 expected: usize,
71 actual: usize,
73 },
74}
75
76impl std::fmt::Display for FusionSelectionError {
77 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78 match self {
79 Self::ExchangeSizeOverflow { n } => write!(
80 f,
81 "megakernel fusion selector n*n overflow for n={n}. Fix: shard the work batch before fusion selection."
82 ),
83 Self::CostLen { expected, actual } => write!(
84 f,
85 "megakernel fusion selector cost length {actual} does not match n={expected}. Fix: pass one cost per work item."
86 ),
87 Self::ExchangeAdjLen { expected, actual } => write!(
88 f,
89 "megakernel fusion selector exchange_adj length {actual} does not match n*n={expected}. Fix: pass a dense row-major n*n exchange graph."
90 ),
91 }
92 }
93}
94
95impl std::error::Error for FusionSelectionError {}
96
97fn validate_selector_shape(
98 cost_len: usize,
99 n: u32,
100 exchange_adj_len: usize,
101) -> Result<(usize, usize), FusionSelectionError> {
102 let n_usize = usize::try_from(n)
103 .map_err(|_| FusionSelectionError::ExchangeSizeOverflow { n: usize::MAX })?;
104 let cells = n_usize
105 .checked_mul(n_usize)
106 .ok_or(FusionSelectionError::ExchangeSizeOverflow { n: n_usize })?;
107 if cost_len != n_usize {
108 return Err(FusionSelectionError::CostLen {
109 expected: n_usize,
110 actual: cost_len,
111 });
112 }
113 if exchange_adj_len != cells {
114 return Err(FusionSelectionError::ExchangeAdjLen {
115 expected: cells,
116 actual: exchange_adj_len,
117 });
118 }
119 Ok((n_usize, cells))
120}
121
122#[derive(Debug, Default)]
127pub struct CompactFusionPlanningScratch {
128 costs_q16: Vec<u16>,
129 stalks: Vec<f32>,
130 diffused_stalks: Vec<f32>,
131 effective_divergence: Vec<u32>,
132 deltas: Vec<f32>,
133 sorted_deltas: Vec<f32>,
134 exchange_adj: Vec<u32>,
135 order: Vec<usize>,
136 selection: FusionSelectionScratch,
137}
138
139impl CompactFusionPlanningScratch {
140 #[must_use]
142 pub fn exchange_adj(&self) -> &[u32] {
143 &self.exchange_adj
144 }
145
146 #[must_use]
148 pub fn selected(&self) -> &[u32] {
149 self.selection.result()
150 }
151}
152
153pub fn plan_compact_fusion_into<'a>(
158 work_items: &[MegakernelWorkItem],
159 scratch: &'a mut CompactFusionPlanningScratch,
160) -> &'a [u32] {
161 let n = work_items.len();
162 if n > MAX_DENSE_FUSION_ITEMS {
163 scratch.selection.prepare(n);
164 scratch.selection.result.fill(1);
165 scratch.exchange_adj.clear();
166 return scratch.selection.result();
167 }
168
169 if n == 0 {
170 scratch.costs_q16.clear();
171 scratch.stalks.clear();
172 scratch.diffused_stalks.clear();
173 scratch.effective_divergence.clear();
174 scratch.deltas.clear();
175 scratch.sorted_deltas.clear();
176 scratch.exchange_adj.clear();
177 scratch.selection.prepare(0);
178 return scratch.selection.result();
179 }
180
181 scratch.costs_q16.clear();
182 scratch.costs_q16.resize(n, u16::MAX);
183
184 scratch.stalks.clear();
185 scratch.stalks.extend(
186 work_items
187 .iter()
188 .enumerate()
189 .map(|(item_idx, _item)| (item_idx as f32) * 0.001),
190 );
191 scratch.diffused_stalks.clear();
192 scratch.diffused_stalks.extend_from_slice(&scratch.stalks);
193 for _ in 0..8 {
194 for value in &mut scratch.diffused_stalks {
195 *value -= 0.5_f32 * 0.7_f32 * *value;
196 }
197 }
198
199 let divergence_threshold = 0.05_f32;
200 let mut delta_sum = 0.0_f32;
201 let mut delta_max = 0.0_f32;
202 scratch.effective_divergence.clear();
203 for (&initial, &diffused) in scratch.stalks.iter().zip(scratch.diffused_stalks.iter()) {
204 let delta = (initial - diffused).abs();
205 delta_sum += delta;
206 delta_max = delta_max.max(delta);
207 scratch
208 .effective_divergence
209 .push(u32::from(delta > divergence_threshold));
210 }
211
212 let n_f32 = n as f32;
213 let gap_signal = if delta_max > 0.0_f32 && n_f32 > 0.0_f32 {
214 delta_sum / (n_f32 * delta_max)
215 } else {
216 1.0_f32
217 };
218 if gap_signal < 0.3 {
219 scratch.deltas.clear();
220 scratch.deltas.extend(
221 scratch
222 .stalks
223 .iter()
224 .zip(scratch.diffused_stalks.iter())
225 .map(|(s, d)| (s - d).abs()),
226 );
227 scratch.sorted_deltas.clear();
228 scratch.sorted_deltas.extend_from_slice(&scratch.deltas);
229 scratch
230 .sorted_deltas
231 .sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
232 let median = scratch
233 .sorted_deltas
234 .get(scratch.sorted_deltas.len() / 2)
235 .copied()
236 .unwrap_or(0.0);
237 for (flag, delta) in scratch
238 .effective_divergence
239 .iter_mut()
240 .zip(scratch.deltas.iter())
241 {
242 if *delta < median {
243 *flag = 0;
244 }
245 }
246 }
247
248 scratch.exchange_adj.clear();
249 let dense_cells = n * n;
250 scratch.exchange_adj.resize(dense_cells, 0);
251 let mut has_exchange_conflict = false;
252
253 let mut has_op_conflict = false;
254 scratch.order.clear();
255 scratch.order.extend(0..n);
256 if scratch.order.len() > 1 {
257 scratch
258 .order
259 .sort_unstable_by_key(|&item_idx| work_items[item_idx].op_handle);
260 if scratch
261 .order
262 .windows(2)
263 .any(|window| work_items[window[0]].op_handle == work_items[window[1]].op_handle)
264 {
265 has_op_conflict = true;
266 }
267 }
268 let has_output_input_chain = (0..n.checked_sub(1).unwrap_or(0)).any(|i| {
269 work_items.get(i).map(|w| w.output_handle) == work_items.get(i + 1).map(|w| w.input_handle)
270 });
271 let has_divergence_conflict = scratch.effective_divergence.iter().any(|&v| v != 0);
272 scratch.selection.prepare(n);
273
274 if !has_op_conflict && !has_divergence_conflict {
275 if has_output_input_chain {
276 for cost in scratch.costs_q16.iter_mut() {
277 *cost = discount_q16(*cost, 3_276);
278 }
279 }
280
281 scratch.selection.result.fill(1);
282 return scratch.selection.result();
283 }
284
285 {
286 let conflict_degrees = &mut scratch.selection.conflict_degrees;
287 for i in 0..n {
288 let row_start = i * n;
289 for j in 0..n {
290 if i == j {
291 continue;
292 }
293 let same_op = work_items[i].op_handle == work_items[j].op_handle;
294 if n <= 32 && same_op {
295 scratch.costs_q16[i] = discount_q16(scratch.costs_q16[i], 3_276);
296 }
297 let divergent =
298 scratch.effective_divergence[i] != 0 && scratch.effective_divergence[j] != 0;
299 if same_op || divergent {
300 scratch.exchange_adj[row_start + j] = 1;
301 if i < j {
302 conflict_degrees[i] = increment_degree(conflict_degrees[i]);
303 conflict_degrees[j] = increment_degree(conflict_degrees[j]);
304 }
305 has_exchange_conflict = true;
306 }
307 }
308 }
309 }
310 if has_output_input_chain {
311 for cost in scratch.costs_q16.iter_mut() {
312 *cost = discount_q16(*cost, 3_276);
313 }
314 }
315 if !has_exchange_conflict {
316 scratch.selection.result.fill(1);
317 return scratch.selection.result();
318 }
319
320 let conflict_degrees = &scratch.selection.conflict_degrees;
321 scratch.selection.order.sort_unstable_by(|&a, &b| {
322 scratch.costs_q16[a]
323 .cmp(&scratch.costs_q16[b])
324 .then_with(|| conflict_degrees[a].cmp(&conflict_degrees[b]))
325 .then_with(|| a.cmp(&b))
326 });
327 select_ordered_maximal(
328 &scratch.exchange_adj,
329 n,
330 &scratch.selection.order,
331 &mut scratch.selection.selected,
332 &mut scratch.selection.result,
333 );
334 scratch.selection.result()
335}
336
337#[must_use]
345pub fn select_fused_subset(costs: &[f64], n: u32, exchange_adj: &[u32]) -> Vec<u32> {
346 let mut scratch = FusionSelectionScratch::default();
347 select_fused_subset_into(costs, n, exchange_adj, &mut scratch);
348 scratch.take_result()
349}
350
351pub fn select_fused_subset_into(
353 costs: &[f64],
354 n: u32,
355 exchange_adj: &[u32],
356 scratch: &mut FusionSelectionScratch,
357) {
358 if let Ok((n_usize, _cells)) = validate_selector_shape(costs.len(), n, exchange_adj.len()) {
359 if n_usize <= MAX_DENSE_FUSION_ITEMS && exchange_adj.iter().all(|&edge| edge == 0) {
360 scratch.prepare(n_usize);
361 scratch.result.fill(1);
362 return;
363 }
364 }
365 if let Err(error) = select_fused_subset_checked_into(costs, n, exchange_adj, scratch) {
370 panic!("vyre-runtime fusion subset selection failed on malformed planner input: {error}");
371 }
372}
373
374pub fn select_fused_subset_checked_into(
376 costs: &[f64],
377 n: u32,
378 exchange_adj: &[u32],
379 scratch: &mut FusionSelectionScratch,
380) -> Result<(), FusionSelectionError> {
381 let (n_usize, _cells) = validate_selector_shape(costs.len(), n, exchange_adj.len())?;
382 if n_usize > MAX_DENSE_FUSION_ITEMS {
383 scratch.prepare(n_usize);
384 scratch.result.fill(1);
385 return Ok(());
386 }
387 scratch.prepare(n_usize);
388 if exchange_adj.iter().all(|&edge| edge == 0) {
389 scratch.result.fill(1);
390 return Ok(());
391 }
392 if !compute_conflict_degrees_with_conflict(exchange_adj, n_usize, &mut scratch.conflict_degrees)
393 {
394 scratch.result.fill(1);
395 return Ok(());
396 }
397 scratch.order.sort_unstable_by(|&a, &b| {
398 costs[a]
399 .total_cmp(&costs[b])
400 .then_with(|| scratch.conflict_degrees[a].cmp(&scratch.conflict_degrees[b]))
401 .then_with(|| a.cmp(&b))
402 });
403 select_ordered_maximal(
404 exchange_adj,
405 n_usize,
406 &scratch.order,
407 &mut scratch.selected,
408 &mut scratch.result,
409 );
410 Ok(())
411}
412
413#[must_use]
419pub fn select_fused_subset_compact(costs_q16: &[u16], n: u32, exchange_adj: &[u32]) -> Vec<u32> {
420 let mut scratch = FusionSelectionScratch::default();
421 select_fused_subset_compact_into(costs_q16, n, exchange_adj, &mut scratch);
422 scratch.take_result()
423}
424
425pub fn select_fused_subset_compact_into(
427 costs_q16: &[u16],
428 n: u32,
429 exchange_adj: &[u32],
430 scratch: &mut FusionSelectionScratch,
431) {
432 if let Ok((n_usize, _cells)) = validate_selector_shape(costs_q16.len(), n, exchange_adj.len()) {
433 if n_usize <= MAX_DENSE_FUSION_ITEMS && exchange_adj.iter().all(|&edge| edge == 0) {
434 scratch.prepare(n_usize);
435 scratch.result.fill(1);
436 return;
437 }
438 }
439 if let Err(error) = select_fused_subset_compact_checked_into(costs_q16, n, exchange_adj, scratch)
442 {
443 panic!(
444 "vyre-runtime compact fusion subset selection failed on malformed planner input: {error}"
445 );
446 }
447}
448
449pub fn select_fused_subset_compact_checked_into(
451 costs_q16: &[u16],
452 n: u32,
453 exchange_adj: &[u32],
454 scratch: &mut FusionSelectionScratch,
455) -> Result<(), FusionSelectionError> {
456 let (n_usize, _cells) = validate_selector_shape(costs_q16.len(), n, exchange_adj.len())?;
457 if n_usize > MAX_DENSE_FUSION_ITEMS {
458 scratch.prepare(n_usize);
459 scratch.result.fill(1);
460 return Ok(());
461 }
462 scratch.prepare(n_usize);
463 if exchange_adj.iter().all(|&edge| edge == 0) {
464 scratch.result.fill(1);
465 return Ok(());
466 }
467 if !compute_conflict_degrees_with_conflict(exchange_adj, n_usize, &mut scratch.conflict_degrees)
468 {
469 scratch.result.fill(1);
470 return Ok(());
471 }
472 scratch.order.sort_unstable_by(|&a, &b| {
473 costs_q16[a]
474 .cmp(&costs_q16[b])
475 .then_with(|| scratch.conflict_degrees[a].cmp(&scratch.conflict_degrees[b]))
476 .then_with(|| a.cmp(&b))
477 });
478 select_ordered_maximal(
479 exchange_adj,
480 n_usize,
481 &scratch.order,
482 &mut scratch.selected,
483 &mut scratch.result,
484 );
485 Ok(())
486}
487
488#[must_use]
491
492pub fn select_optimal_fused_subset(costs: &[f64], n: u32, exchange_adj: &[u32]) -> Vec<u32> {
493 select_fused_subset(costs, n, exchange_adj)
494}
495
496#[must_use]
498pub fn select_fused_subset_with_rate(costs: &[f64], n: u32, exchange_adj: &[u32]) -> Vec<u32> {
499 select_fused_subset(costs, n, exchange_adj)
500}
501
502#[must_use]
509pub fn select_fused_subset_pruned(
510 costs: &[f64],
511 n: u32,
512 exchange_adj: &[u32],
513 dead_mask: &[bool],
514) -> Vec<u32> {
515 let mut selection = select_fused_subset(costs, n, exchange_adj);
516 prune_dead_arms_inplace(&mut selection, dead_mask);
517 selection
518}
519
520pub fn select_fused_subset_pruned_into(
522 costs: &[f64],
523 n: u32,
524 exchange_adj: &[u32],
525 dead_mask: &[bool],
526 scratch: &mut FusionSelectionScratch,
527) {
528 select_fused_subset_into(costs, n, exchange_adj, scratch);
529 prune_dead_arms_inplace(&mut scratch.result, dead_mask);
530}
531
532pub fn prune_dead_arms_inplace(selection: &mut [u32], dead_mask: &[bool]) -> u32 {
554 if selection.len() != dead_mask.len() {
555 return 0;
556 }
557 let mut eliminated = 0_u32;
558 for (slot, &dead) in selection.iter_mut().zip(dead_mask.iter()) {
559 if dead && *slot != 0 {
560 *slot = 0;
561 eliminated = eliminated.saturating_add(1);
562 }
563 }
564 eliminated
565}
566
567fn compute_conflict_degrees_with_conflict(exchange_adj: &[u32], n: usize, out: &mut [u32]) -> bool {
568 debug_assert_eq!(out.len(), n);
569 out.fill(0);
570 let mut has_conflict = false;
571 for i in 0..n {
572 let row = i * n;
573 for j in (i + 1)..n {
574 if exchange_adj[row + j] != 0 || exchange_adj[j * n + i] != 0 {
575 out[i] = increment_degree(out[i]);
576 out[j] = increment_degree(out[j]);
577 has_conflict = true;
578 }
579 }
580 }
581 has_conflict
582}
583
584fn discount_q16(value: u16, amount: u16) -> u16 {
585 value.saturating_sub(amount)
586}
587
588fn increment_degree(value: u32) -> u32 {
589 value.saturating_add(1)
590}
591
592fn select_ordered_maximal(
593 exchange_adj: &[u32],
594 n: usize,
595 order: &[usize],
596 selected: &mut Vec<usize>,
597 result: &mut [u32],
598) {
599 result.fill(0);
600 selected.clear();
601
602 if n == 0 {
603 return;
604 }
605
606 if n <= 64 {
607 let mut conflict_masks = [0_u64; 64];
608 for i in 0..n {
609 let row = i * n;
610 let mut mask = 0_u64;
611 for j in 0..n {
612 if i == j {
613 continue;
614 }
615 if exchange_adj[row + j] != 0 || exchange_adj[j * n + i] != 0 {
616 mask |= 1_u64 << j;
617 }
618 }
619 conflict_masks[i] = mask;
620 }
621
622 let mut selected_mask = 0_u64;
623 for &item in order {
624 if item >= n {
625 continue;
626 }
627 if conflict_masks[item] & selected_mask == 0 {
628 result[item] = 1;
629 selected_mask |= 1_u64 << item;
630 selected.push(item);
631 }
632 }
633 return;
634 }
635
636 if n <= 128 {
637 let mut conflict_masks_lo = [0_u64; 128];
638 let mut conflict_masks_hi = [0_u64; 128];
639 for i in 0..n {
640 let row = i * n;
641 let mut mask_lo = 0_u64;
642 let mut mask_hi = 0_u64;
643 for j in 0..n {
644 if i == j {
645 continue;
646 }
647 if exchange_adj[row + j] != 0 || exchange_adj[j * n + i] != 0 {
648 if j < 64 {
649 mask_lo |= 1_u64 << j;
650 } else {
651 mask_hi |= 1_u64 << (j - 64);
652 }
653 }
654 }
655 conflict_masks_lo[i] = mask_lo;
656 conflict_masks_hi[i] = mask_hi;
657 }
658
659 let mut selected_lo = 0_u64;
660 let mut selected_hi = 0_u64;
661 for &item in order {
662 if item >= n {
663 continue;
664 }
665 let conflict = (conflict_masks_lo[item] & selected_lo) != 0
666 || (conflict_masks_hi[item] & selected_hi) != 0;
667 if !conflict {
668 result[item] = 1;
669 if item < 64 {
670 selected_lo |= 1_u64 << item;
671 } else {
672 selected_hi |= 1_u64 << (item - 64);
673 }
674 selected.push(item);
675 }
676 }
677 return;
678 }
679
680 if n <= 192 {
681 let mut conflict_masks_0 = [0_u64; 192];
682 let mut conflict_masks_1 = [0_u64; 192];
683 let mut conflict_masks_2 = [0_u64; 192];
684 for i in 0..n {
685 let row = i * n;
686 let mut mask_0 = 0_u64;
687 let mut mask_1 = 0_u64;
688 let mut mask_2 = 0_u64;
689 for j in 0..n {
690 if i == j {
691 continue;
692 }
693 if exchange_adj[row + j] != 0 || exchange_adj[j * n + i] != 0 {
694 match j / 64 {
695 0 => mask_0 |= 1_u64 << (j % 64),
696 1 => mask_1 |= 1_u64 << (j % 64),
697 2 => mask_2 |= 1_u64 << (j % 64),
698 _ => {}
699 }
700 }
701 }
702 conflict_masks_0[i] = mask_0;
703 conflict_masks_1[i] = mask_1;
704 conflict_masks_2[i] = mask_2;
705 }
706
707 let mut selected_0 = 0_u64;
708 let mut selected_1 = 0_u64;
709 let mut selected_2 = 0_u64;
710 for &item in order {
711 if item >= n {
712 continue;
713 }
714 let conflict = (conflict_masks_0[item] & selected_0 != 0)
715 || (conflict_masks_1[item] & selected_1 != 0)
716 || (conflict_masks_2[item] & selected_2 != 0);
717 if !conflict {
718 result[item] = 1;
719 let bit = 1_u64 << (item % 64);
720 match item / 64 {
721 0 => selected_0 |= bit,
722 1 => selected_1 |= bit,
723 2 => selected_2 |= bit,
724 _ => {}
725 }
726 selected.push(item);
727 }
728 }
729 return;
730 }
731
732 if n <= 256 {
733 let mut conflict_masks_0 = [0_u64; 256];
734 let mut conflict_masks_1 = [0_u64; 256];
735 let mut conflict_masks_2 = [0_u64; 256];
736 let mut conflict_masks_3 = [0_u64; 256];
737 for i in 0..n {
738 let row = i * n;
739 let mut mask_0 = 0_u64;
740 let mut mask_1 = 0_u64;
741 let mut mask_2 = 0_u64;
742 let mut mask_3 = 0_u64;
743 for j in 0..n {
744 if i == j {
745 continue;
746 }
747 if exchange_adj[row + j] != 0 || exchange_adj[j * n + i] != 0 {
748 match j / 64 {
749 0 => mask_0 |= 1_u64 << (j % 64),
750 1 => mask_1 |= 1_u64 << (j % 64),
751 2 => mask_2 |= 1_u64 << (j % 64),
752 _ => mask_3 |= 1_u64 << (j % 64),
753 }
754 }
755 }
756 conflict_masks_0[i] = mask_0;
757 conflict_masks_1[i] = mask_1;
758 conflict_masks_2[i] = mask_2;
759 conflict_masks_3[i] = mask_3;
760 }
761
762 let mut selected_0 = 0_u64;
763 let mut selected_1 = 0_u64;
764 let mut selected_2 = 0_u64;
765 let mut selected_3 = 0_u64;
766 for &item in order {
767 if item >= n {
768 continue;
769 }
770 let conflict = (conflict_masks_0[item] & selected_0 != 0)
771 || (conflict_masks_1[item] & selected_1 != 0)
772 || (conflict_masks_2[item] & selected_2 != 0)
773 || (conflict_masks_3[item] & selected_3 != 0);
774 if !conflict {
775 result[item] = 1;
776 let bit = 1_u64 << (item % 64);
777 match item / 64 {
778 0 => selected_0 |= bit,
779 1 => selected_1 |= bit,
780 2 => selected_2 |= bit,
781 _ => selected_3 |= bit,
782 }
783 selected.push(item);
784 }
785 }
786 return;
787 }
788
789 let chunks = n.div_ceil(64);
790 let mut conflict_masks = vec![0_u64; n * chunks];
791 for i in 0..n {
792 for j in (i + 1)..n {
793 if exchange_adj[i * n + j] != 0 || exchange_adj[j * n + i] != 0 {
794 let i_word = i / 64;
795 let i_bit = 1_u64 << (i % 64);
796 let j_word = j / 64;
797 let j_bit = 1_u64 << (j % 64);
798
799 let i_base = i * chunks;
800 let j_base = j * chunks;
801 conflict_masks[i_base + j_word] |= j_bit;
802 conflict_masks[j_base + i_word] |= i_bit;
803 }
804 }
805 }
806
807 let mut selected_mask = vec![0_u64; chunks];
808 for &item in order {
809 if item >= n {
810 continue;
811 }
812 let base = item * chunks;
813 let mut conflict = false;
814 for chunk in 0..chunks {
815 if conflict_masks[base + chunk] & selected_mask[chunk] != 0 {
816 conflict = true;
817 break;
818 }
819 }
820 if !conflict {
821 result[item] = 1;
822 selected.push(item);
823 selected_mask[item / 64] |= 1_u64 << (item % 64);
824 }
825 }
826}
827
828#[cfg(test)]
829mod tests;