1use std::collections::{BTreeMap, BTreeSet};
15
16use serde::{Deserialize, Serialize};
17
18use super::certificate::ChoreographyCertificate;
19use super::target::{
20 precision_bytes, precision_int_bits, DataflowPolicy, HandshakeKind, Interconnect, MeshShape,
21 TileShape,
22};
23use crate::codegen_sva::fabric_rtl::{LoopCell, Src};
24use crate::codegen_sva::fabric_target::Precision;
25
26macro_rules! id_newtype {
27 ($(#[$m:meta])* $name:ident) => {
28 $(#[$m])*
29 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
30 pub struct $name(pub u32);
31 impl $name {
32 pub fn idx(self) -> usize {
33 self.0 as usize
34 }
35 pub fn new(i: usize) -> Self {
36 $name(i as u32)
37 }
38 }
39 };
40}
41
42id_newtype!(
43 #[doc = "A global tensor."]
44 TensorId
45);
46id_newtype!(
47 #[doc = "A stage."]
48 StageId
49);
50id_newtype!(
51 #[doc = "A buffer."]
52 BufferId
53);
54id_newtype!(
55 #[doc = "An edge."]
56 EdgeId
57);
58id_newtype!(
59 #[doc = "A multicast group."]
60 GroupId
61);
62
63#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
67pub struct ChipId(pub u16);
68
69fn chip_is_zero(c: &u16) -> bool {
70 *c == 0
71}
72
73fn one_bank() -> usize {
74 1
75}
76
77fn is_one_bank(b: &usize) -> bool {
78 *b == 1
79}
80id_newtype!(
81 #[doc = "A semaphore."]
82 SemId
83);
84id_newtype!(
85 #[doc = "A work item (one output block)."]
86 WorkItemId
87);
88id_newtype!(
89 #[doc = "A loop variable of the choreography."]
90 LoopVar
91);
92id_newtype!(
93 #[doc = "A loop count the program reads from its data ([`Trip`])."]
94 TripId
95);
96
97#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
99pub struct Site {
100 pub x: u16,
101 pub y: u16,
102 #[serde(default, skip_serializing_if = "chip_is_zero")]
105 pub chip: u16,
106}
107
108impl Site {
109 pub fn new(x: usize, y: usize) -> Site {
110 Site {
111 x: x as u16,
112 y: y as u16,
113 chip: 0,
114 }
115 }
116 pub fn on_chip(chip: ChipId, x: usize, y: usize) -> Site {
118 Site {
119 x: x as u16,
120 y: y as u16,
121 chip: chip.0,
122 }
123 }
124 pub fn chip_id(self) -> ChipId {
125 ChipId(self.chip)
126 }
127}
128
129#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
131pub struct Rect {
132 pub x0: u16,
133 pub y0: u16,
134 pub x1: u16,
135 pub y1: u16,
136}
137
138impl Rect {
139 pub fn contains(&self, s: Site) -> bool {
140 s.x >= self.x0 && s.x <= self.x1 && s.y >= self.y0 && s.y <= self.y1
141 }
142 pub fn sites(&self) -> Vec<Site> {
144 let mut v = Vec::new();
145 for y in self.y0..=self.y1 {
146 for x in self.x0..=self.x1 {
147 v.push(Site { x, y, chip: 0 });
148 }
149 }
150 v
151 }
152 pub fn len(&self) -> usize {
153 (self.x1 - self.x0 + 1) as usize * (self.y1 - self.y0 + 1) as usize
154 }
155}
156
157#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
159pub enum Engine {
160 Movement(u8),
161 Compute,
162}
163
164#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
166pub struct NocId(pub u8);
167
168#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
169pub enum TensorRole {
170 Input,
171 Output,
172 Param,
173 Activation,
178 State,
182}
183
184impl TensorRole {
185 pub fn starts_filled(self) -> bool {
188 matches!(
189 self,
190 TensorRole::Input | TensorRole::Param | TensorRole::State
191 )
192 }
193 pub fn is_written(self) -> bool {
195 matches!(
196 self,
197 TensorRole::Output | TensorRole::Activation | TensorRole::State
198 )
199 }
200 pub fn is_read(self) -> bool {
202 !matches!(self, TensorRole::Output)
203 }
204 pub fn is_fully_covered(self) -> bool {
208 matches!(self, TensorRole::Output | TensorRole::Activation)
209 }
210}
211
212#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
213pub enum TensorSpace {
214 Dram,
215 Hbm,
216 Sram,
219}
220
221#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
222pub enum Layout {
223 RowMajorTiles,
224}
225
226#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
227pub struct ElemShape {
228 pub rows: usize,
229 pub cols: usize,
230}
231
232#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
233pub struct TileGrid {
234 pub rows_t: usize,
235 pub cols_t: usize,
236}
237
238impl TileGrid {
239 pub fn count(&self) -> usize {
240 self.rows_t * self.cols_t
241 }
242}
243
244#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
247pub enum TensorInit {
248 Zero,
249 Pattern {
251 modulus: i64,
252 shift: i64,
253 },
254 Values(Vec<i64>),
255 FloatValues(Vec<u64>),
257 ConstF(u64),
264 RandomBf16 {
266 seed: u32,
267 shift_milli: i32,
268 },
269 External,
271}
272
273impl TensorInit {
274 pub fn abs_bound(&self) -> Option<i64> {
276 match self {
277 TensorInit::Zero => Some(0),
278 TensorInit::Pattern { modulus, shift } => {
279 let lo = *shift;
280 let hi = (modulus - 1) + shift;
281 Some(lo.abs().max(hi.abs()))
282 }
283 TensorInit::Values(v) => Some(v.iter().map(|x| x.abs()).max().unwrap_or(0)),
284 TensorInit::FloatValues(_)
285 | TensorInit::ConstF(_)
286 | TensorInit::RandomBf16 { .. }
287 | TensorInit::External => None,
288 }
289 }
290 pub fn range(&self) -> Option<(i64, i64)> {
292 match self {
293 TensorInit::Zero => Some((0, 0)),
294 TensorInit::Pattern { modulus, shift } => Some((*shift, (modulus - 1) + shift)),
295 TensorInit::Values(v) => Some((v.iter().copied().min().unwrap_or(0), v.iter().copied().max().unwrap_or(0))),
296 TensorInit::FloatValues(_)
297 | TensorInit::ConstF(_)
298 | TensorInit::RandomBf16 { .. }
299 | TensorInit::External => None,
300 }
301 }
302 pub fn element(&self, p: usize) -> Option<i64> {
304 match self {
305 TensorInit::Zero => Some(0),
306 TensorInit::Pattern { modulus, shift } => Some((p as i64 % modulus) + shift),
307 TensorInit::Values(v) => v.get(p).copied(),
308 TensorInit::FloatValues(_)
309 | TensorInit::ConstF(_)
310 | TensorInit::RandomBf16 { .. }
311 | TensorInit::External => None,
312 }
313 }
314 pub fn element_f64(&self, p: usize) -> Option<f64> {
316 match self {
317 TensorInit::FloatValues(v) => v.get(p).map(|b| f64::from_bits(*b)),
318 TensorInit::ConstF(b) => Some(f64::from_bits(*b)),
319 other => other.element(p).map(|x| x as f64),
320 }
321 }
322 pub fn is_float(&self) -> bool {
323 matches!(self, TensorInit::FloatValues(_) | TensorInit::ConstF(_))
324 }
325}
326
327#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
328pub struct TensorDecl {
329 pub id: TensorId,
330 pub name: String,
331 pub role: TensorRole,
332 pub elems: ElemShape,
333 pub tiles: TileGrid,
334 pub dtype: Precision,
335 pub layout: Layout,
336 pub space: TensorSpace,
337 pub init: TensorInit,
338 pub range: Option<(i64, i64)>,
340 #[serde(default, skip_serializing_if = "Option::is_none")]
343 pub chip: Option<ChipId>,
344}
345
346impl TensorDecl {
347 pub fn tile_index(&self, row_t: usize, col_t: usize) -> usize {
348 row_t * self.tiles.cols_t + col_t
349 }
350 pub fn abs_bound(&self) -> Option<i64> {
352 self.init
353 .abs_bound()
354 .or_else(|| self.range.map(|(lo, hi)| lo.abs().max(hi.abs())))
355 }
356}
357
358#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
368pub struct IndirectTile {
369 pub index: BufferId,
371 pub tile_stride: usize,
373 pub elems_per_tile: usize,
375 pub bound: (usize, usize),
377}
378
379#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
382pub struct TileAddr {
383 pub base: usize,
384 pub steps: Vec<(LoopVar, usize)>,
385 #[serde(default)]
386 pub indirect: Option<IndirectTile>,
387}
388
389impl TileAddr {
390 pub fn constant(base: usize) -> TileAddr {
391 TileAddr {
392 base,
393 steps: Vec::new(),
394 indirect: None,
395 }
396 }
397 pub fn resolve(&self, env: &BTreeMap<LoopVar, usize>) -> usize {
400 self.base
401 + self
402 .steps
403 .iter()
404 .map(|(v, s)| env.get(v).copied().unwrap_or(0) * s)
405 .sum::<usize>()
406 }
407}
408
409#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
412pub struct AccessPattern {
413 pub origin: TileAddr,
414 pub block_h: usize,
415 pub block_w: usize,
416 pub stride_w: usize,
417 pub stride_h: usize,
418}
419
420impl AccessPattern {
421 pub fn block_tiles(&self) -> usize {
422 self.block_h * self.block_w
423 }
424
425 pub fn tiles_at(
429 &self,
430 coord: (usize, usize),
431 sub_blocks: (usize, usize),
432 ) -> Option<Vec<usize>> {
433 use super::passes::{LV_ITEM_X, LV_ITEM_Y, LV_SUB_H, LV_SUB_W};
434 if self.origin.indirect.is_some()
435 || self
436 .origin
437 .steps
438 .iter()
439 .any(|(v, _)| ![LV_ITEM_Y, LV_ITEM_X, LV_SUB_H, LV_SUB_W].contains(v))
440 {
441 return None;
442 }
443 let mut tiles = Vec::with_capacity(sub_blocks.0 * sub_blocks.1 * self.block_tiles());
444 for sh in 0..sub_blocks.0 {
445 for sw in 0..sub_blocks.1 {
446 let at = self
447 .origin
448 .steps
449 .iter()
450 .fold(self.origin.base, |a, (v, stride)| {
451 let value = match *v {
452 LV_ITEM_Y => coord.0,
453 LV_ITEM_X => coord.1,
454 LV_SUB_H => sh,
455 _ => sw,
456 };
457 a + value * stride
458 });
459 for h in 0..self.block_h {
460 for w in 0..self.block_w {
461 tiles.push(at + h * self.stride_h + w * self.stride_w);
462 }
463 }
464 }
465 }
466 Some(tiles)
467 }
468}
469
470#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
472pub struct BlockShape {
473 pub osh: usize,
474 pub osw: usize,
475 pub ibw: usize,
476}
477
478impl BlockShape {
479 pub fn acc_tiles(&self) -> usize {
480 self.osh * self.osw
481 }
482}
483
484#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
486pub enum AccPolicy {
487 I32,
488 I64,
489 F32,
490}
491
492impl AccPolicy {
493 pub fn for_result(p: Precision) -> AccPolicy {
497 match p {
498 Precision::Int32 => AccPolicy::I32,
499 Precision::Int8 | Precision::Int16 | Precision::Int64 => AccPolicy::I64,
500 Precision::Fp8 | Precision::Bf16 | Precision::Fp16 | Precision::Fp32 => AccPolicy::F32,
501 }
502 }
503 pub fn interm_precision(self) -> Precision {
505 match self {
506 AccPolicy::I32 => Precision::Int32,
507 AccPolicy::I64 => Precision::Int64,
508 AccPolicy::F32 => Precision::Fp32,
509 }
510 }
511 pub fn bits(self) -> u32 {
512 match self {
513 AccPolicy::I32 => 32,
514 AccPolicy::I64 => 64,
515 AccPolicy::F32 => 32,
516 }
517 }
518}
519
520#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
522pub enum NumericsPolicy {
523 Exact,
524 AllowSpillRounding,
525}
526
527#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
529pub enum NumericsClaim {
530 Exact,
531 Approx { reason: String },
532}
533
534#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
536pub enum Granularity {
537 PerTile,
538 PerSubBlock,
539}
540
541#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
542pub enum ComputeOp {
543 MatmulTiles {
544 block: BlockShape,
545 sub_blocks: (usize, usize),
547 num_blocks: usize,
548 operand: Precision,
549 acc: AccPolicy,
550 dataflow: DataflowPolicy,
551 interm: Option<BufferId>,
552 granularity: Granularity,
553 },
554 Requant {
555 m: i64,
556 s: u32,
557 },
558 Bias {
559 bias: BufferId,
560 },
561 Relu,
562 Generic {
580 iters: Vec<IterSpec>,
581 ins: Vec<ElemAccess>,
583 out: ElemAccess,
584 body: TensorBody,
585 },
586 Scalar(ScalarBody),
589}
590
591#[derive(Clone, Debug)]
596pub enum ComputeShape<'a> {
597 Tiles {
598 block: &'a BlockShape,
599 sub_blocks: (usize, usize),
600 operand: Precision,
601 },
602 AccMap {
605 body: TensorBody,
606 reads: Vec<AccRead>,
607 },
608 Generic {
609 iters: &'a [IterSpec],
610 ins: &'a [ElemAccess],
611 out: &'a ElemAccess,
612 body: &'a TensorBody,
613 },
614 Scalar(&'a ScalarBody),
615}
616
617#[derive(Clone, Copy, Debug, PartialEq, Eq)]
619pub enum AccRead {
620 Tile,
622 RowOfTile,
624}
625
626impl ComputeOp {
627 pub fn shape(&self) -> Result<ComputeShape<'_>, &'static str> {
630 let konst = |v: i64| Box::new(TensorBody::Const((v as f64).to_bits()));
631 let x = || Box::new(TensorBody::Operand(0));
632 Ok(match self {
633 ComputeOp::MatmulTiles { block, sub_blocks, operand, .. } => ComputeShape::Tiles {
634 block,
635 sub_blocks: *sub_blocks,
636 operand: *operand,
637 },
638 ComputeOp::Requant { m, s } => {
639 if (*m as f64) as i64 != *m {
640 return Err("a requant multiplier past ±2⁵³");
641 }
642 ComputeShape::AccMap {
643 body: TensorBody::Shr { arg: Box::new(TensorBody::Mul(konst(*m), x())), by: *s },
644 reads: vec![AccRead::Tile],
645 }
646 }
647 ComputeOp::Bias { .. } => ComputeShape::AccMap {
648 body: TensorBody::Add(x(), Box::new(TensorBody::Operand(1))),
649 reads: vec![AccRead::Tile, AccRead::RowOfTile],
650 },
651 ComputeOp::Relu => ComputeShape::AccMap {
652 body: TensorBody::Max(x(), konst(0)),
653 reads: vec![AccRead::Tile],
654 },
655 ComputeOp::Generic { iters, ins, out, body } => ComputeShape::Generic { iters, ins, out, body },
656 ComputeOp::Scalar(b) => ComputeShape::Scalar(b),
657 })
658 }
659}
660
661#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
663pub struct IterSpec {
664 pub var: LoopVar,
665 pub extent: usize,
666 pub kind: IterKind,
667}
668
669#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
671pub enum IterKind {
672 Parallel,
673 Reduce {
675 pool: ReducePool,
676 },
677}
678
679#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
685pub struct IndirectTerm {
686 pub operand: usize,
687 pub stride: i64,
688 pub modulus: usize,
689}
690
691#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
709pub struct ElemAccess {
710 pub base: i64,
711 pub terms: Vec<(LoopVar, i64)>,
712 #[serde(default)]
713 pub indirect: Option<IndirectTerm>,
714}
715
716impl ElemAccess {
717 pub fn offset(&self, env: &BTreeMap<LoopVar, usize>) -> i64 {
719 self.base
720 + self
721 .terms
722 .iter()
723 .map(|(v, s)| env.get(v).copied().unwrap_or(0) as i64 * s)
724 .sum::<i64>()
725 }
726 pub fn offset_at(&self, row: LoopVar, r: usize, col: LoopVar, c: usize) -> i64 {
728 self.base
729 + self
730 .terms
731 .iter()
732 .map(|(v, s)| {
733 if *v == row {
734 r as i64 * s
735 } else if *v == col {
736 c as i64 * s
737 } else {
738 0
739 }
740 })
741 .sum::<i64>()
742 }
743 pub fn offset_at_k(
745 &self,
746 row: LoopVar,
747 r: usize,
748 col: LoopVar,
749 c: usize,
750 k: LoopVar,
751 x: usize,
752 ) -> i64 {
753 self.offset_at(row, r, col, c)
754 + self
755 .terms
756 .iter()
757 .filter(|(v, _)| *v == k)
758 .map(|(_, s)| x as i64 * s)
759 .sum::<i64>()
760 }
761 pub fn varies_along(&self, var: LoopVar) -> bool {
763 self.terms.iter().any(|(v, s)| *v == var && *s != 0)
764 }
765 pub fn broadcast_class(
774 &self,
775 row: LoopVar,
776 col: LoopVar,
777 width: usize,
778 ) -> Option<(bool, bool)> {
779 if self.base != 0 || self.indirect.is_some() {
780 return None;
781 }
782 let (mut row_seen, mut col_seen) = (false, false);
783 for (v, s) in &self.terms {
784 if *s == 0 {
785 continue;
786 }
787 if *v == row && !row_seen && *s == width as i64 {
788 row_seen = true;
789 } else if *v == col && !col_seen && *s == 1 {
790 col_seen = true;
791 } else {
792 return None;
793 }
794 }
795 Some((row_seen, col_seen))
796 }
797}
798
799#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
805pub enum TensorBody {
806 Operand(usize),
808 Const(u64),
810 Add(Box<TensorBody>, Box<TensorBody>),
811 Sub(Box<TensorBody>, Box<TensorBody>),
812 Mul(Box<TensorBody>, Box<TensorBody>),
813 Min(Box<TensorBody>, Box<TensorBody>),
815 Max(Box<TensorBody>, Box<TensorBody>),
816 FloorDiv(Box<TensorBody>, Box<TensorBody>),
818 Shr {
820 arg: Box<TensorBody>,
821 by: u32,
822 },
823 Unary {
824 kind: SfpuKind,
825 arg: Box<TensorBody>,
826 },
827 Select {
832 cond: Pred,
833 then: Box<TensorBody>,
834 otherwise: Box<TensorBody>,
835 },
836}
837
838impl TensorBody {
839 pub fn post_order(&self) -> Vec<&TensorBody> {
842 fn walk<'a>(b: &'a TensorBody, out: &mut Vec<&'a TensorBody>) {
843 match b {
844 TensorBody::Operand(_) | TensorBody::Const(_) => {}
845 TensorBody::Add(x, y)
846 | TensorBody::Sub(x, y)
847 | TensorBody::Mul(x, y)
848 | TensorBody::Min(x, y)
849 | TensorBody::Max(x, y)
850 | TensorBody::FloorDiv(x, y) => {
851 walk(x, out);
852 walk(y, out);
853 }
854 TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => walk(arg, out),
855 TensorBody::Select { then, otherwise, .. } => {
856 walk(then, out);
857 walk(otherwise, out);
858 }
859 }
860 out.push(b);
861 }
862 let mut out = Vec::new();
863 walk(self, &mut out);
864 out
865 }
866 pub fn max_operand(&self) -> Option<usize> {
869 match self {
870 TensorBody::Operand(i) => Some(*i),
871 TensorBody::Const(_) => None,
872 TensorBody::Add(x, y)
873 | TensorBody::Sub(x, y)
874 | TensorBody::Mul(x, y)
875 | TensorBody::Min(x, y)
876 | TensorBody::Max(x, y)
877 | TensorBody::FloorDiv(x, y) => x.max_operand().max(y.max_operand()),
878 TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => arg.max_operand(),
879 TensorBody::Select {
880 cond,
881 then,
882 otherwise,
883 } => cond
884 .operand()
885 .max(then.max_operand())
886 .max(otherwise.max_operand()),
887 }
888 }
889 pub fn pred_names_operand(&self, i: usize) -> bool {
891 match self {
892 TensorBody::Operand(_) | TensorBody::Const(_) => false,
893 TensorBody::Add(x, y)
894 | TensorBody::Sub(x, y)
895 | TensorBody::Mul(x, y)
896 | TensorBody::Min(x, y)
897 | TensorBody::Max(x, y)
898 | TensorBody::FloorDiv(x, y) => x.pred_names_operand(i) || y.pred_names_operand(i),
899 TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => arg.pred_names_operand(i),
900 TensorBody::Select {
901 cond,
902 then,
903 otherwise,
904 } => {
905 cond.operand() == Some(i)
906 || then.pred_names_operand(i)
907 || otherwise.pred_names_operand(i)
908 }
909 }
910 }
911 pub fn integral(&self) -> bool {
914 match self {
915 TensorBody::Operand(_) => true,
916 TensorBody::Const(bits) => {
917 let v = f64::from_bits(*bits);
918 v.fract() == 0.0 && v.abs() < 9_007_199_254_740_992.0
919 }
920 TensorBody::Add(x, y)
921 | TensorBody::Sub(x, y)
922 | TensorBody::Mul(x, y)
923 | TensorBody::Min(x, y)
924 | TensorBody::Max(x, y)
925 | TensorBody::FloorDiv(x, y) => x.integral() && y.integral(),
926 TensorBody::Shr { arg, by } => *by <= 63 && arg.integral(),
927 TensorBody::Unary { .. } => false,
928 TensorBody::Select {
929 then, otherwise, ..
930 } => then.integral() && otherwise.integral(),
931 }
932 }
933}
934
935#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
937pub enum Pred {
938 AxisLt(LoopVar, PredRhs),
940 AxisGe(LoopVar, PredRhs),
942}
943
944impl Pred {
945 pub fn var(&self) -> LoopVar {
946 match self {
947 Pred::AxisLt(v, _) | Pred::AxisGe(v, _) => *v,
948 }
949 }
950 pub fn rhs(&self) -> PredRhs {
951 match self {
952 Pred::AxisLt(_, r) | Pred::AxisGe(_, r) => *r,
953 }
954 }
955 pub fn operand(&self) -> Option<usize> {
957 match self.rhs() {
958 PredRhs::Operand(i) => Some(i),
959 PredRhs::Const(_) => None,
960 }
961 }
962 pub fn map_operand(&self, shift: &dyn Fn(usize) -> usize) -> Pred {
964 let rhs = match self.rhs() {
965 PredRhs::Operand(i) => PredRhs::Operand(shift(i)),
966 c => c,
967 };
968 match self {
969 Pred::AxisLt(v, _) => Pred::AxisLt(*v, rhs),
970 Pred::AxisGe(v, _) => Pred::AxisGe(*v, rhs),
971 }
972 }
973}
974
975#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
977pub enum PredRhs {
978 Const(usize),
980 Operand(usize),
983}
984
985#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
987pub enum SfpuKind {
988 Exp,
989 Recip,
990 Rsqrt {
993 eps_nano: u64,
994 over: usize,
995 },
996 Silu,
998 Relu,
999 Ln,
1002}
1003
1004#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1006pub enum ReducePool {
1007 Sum,
1008 Max,
1009 ArgMax,
1014}
1015
1016#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1018pub enum RowXformKind {
1019 Rsqrt { eps_nano: u64, over: usize },
1033}
1034
1035#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1037pub struct ScalarBody {
1038 pub n_in: usize,
1039 pub cells: Vec<LoopCell>,
1040 pub result: Src,
1041}
1042
1043#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1044pub enum StageKind {
1045 Load {
1046 tensor: TensorId,
1047 pattern: AccessPattern,
1048 num_blocks: usize,
1049 multicast: Option<GroupId>,
1051 },
1052 Store {
1053 tensor: TensorId,
1054 pattern: AccessPattern,
1055 sub_blocks: (usize, usize),
1056 },
1057 Recv {
1059 group: GroupId,
1060 },
1061 AllReduce {
1077 pool: ReducePool,
1078 tensor: TensorId,
1079 combine: Box<ComputeOp>,
1080 },
1081 Compute(ComputeOp),
1082 Scalar(ScalarBody),
1083}
1084
1085impl StageKind {
1086 pub fn remap_buffers(&mut self, f: &dyn Fn(BufferId) -> BufferId) {
1090 match self {
1091 StageKind::Load { pattern, .. } | StageKind::Store { pattern, .. } => {
1092 if let Some(ind) = &mut pattern.origin.indirect {
1093 ind.index = f(ind.index);
1094 }
1095 }
1096 StageKind::Compute(ComputeOp::MatmulTiles { interm, .. }) => *interm = interm.map(f),
1097 StageKind::Compute(ComputeOp::Bias { bias }) => *bias = f(*bias),
1098 _ => {}
1099 }
1100 }
1101
1102 pub fn role(&self) -> SiteRole {
1103 match self {
1104 StageKind::Load { .. } | StageKind::Store { .. } | StageKind::Recv { .. } => {
1105 SiteRole::Movement
1106 }
1107 StageKind::AllReduce { .. } => SiteRole::Compute,
1109 StageKind::Compute(_) | StageKind::Scalar(_) => SiteRole::Compute,
1110 }
1111 }
1112}
1113
1114#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1115pub enum SiteRole {
1116 Compute,
1117 Movement,
1118}
1119
1120#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1121pub struct Stage {
1122 pub id: StageId,
1123 pub name: String,
1124 pub kind: StageKind,
1125 pub role: SiteRole,
1126 pub ins: Vec<BufferId>,
1127 pub outs: Vec<BufferId>,
1128 pub site: Option<Site>,
1129 pub item: Option<WorkItemId>,
1130 pub origin: Option<StageId>,
1132}
1133
1134#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
1135pub enum MemSpace {
1136 Local,
1137 Acc,
1138 Pmu,
1139}
1140
1141#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
1142pub enum BufferKind {
1143 Fifo {
1144 depth_pages: usize,
1145 page_tiles: usize,
1146 },
1147 Interm {
1149 depth_pages: usize,
1150 page_tiles: usize,
1151 },
1152 Accumulator {
1153 tiles: usize,
1154 },
1155}
1156
1157impl BufferKind {
1158 pub fn tiles(&self) -> usize {
1159 match self {
1160 BufferKind::Fifo {
1161 depth_pages,
1162 page_tiles,
1163 }
1164 | BufferKind::Interm {
1165 depth_pages,
1166 page_tiles,
1167 } => depth_pages * page_tiles,
1168 BufferKind::Accumulator { tiles } => *tiles,
1169 }
1170 }
1171 pub fn depth_pages(&self) -> usize {
1173 match self {
1174 BufferKind::Fifo { depth_pages, .. } | BufferKind::Interm { depth_pages, .. } => {
1175 *depth_pages
1176 }
1177 BufferKind::Accumulator { .. } => 1,
1178 }
1179 }
1180 pub fn page_tiles(&self) -> usize {
1181 match self {
1182 BufferKind::Fifo { page_tiles, .. } | BufferKind::Interm { page_tiles, .. } => {
1183 *page_tiles
1184 }
1185 BufferKind::Accumulator { tiles } => *tiles,
1186 }
1187 }
1188}
1189
1190#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1191pub struct Buffer {
1192 pub id: BufferId,
1193 pub name: String,
1194 pub kind: BufferKind,
1195 pub dtype: Precision,
1196 pub space: MemSpace,
1197 pub site: Option<Site>,
1198 pub origin: Option<BufferId>,
1199 pub alias_of: Option<BufferId>,
1202}
1203
1204#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1205pub enum Handshake {
1206 Fifo,
1207 Semaphore {
1208 ready: SemId,
1209 valid: SemId,
1210 ready_count: usize,
1211 },
1212 Barrier {
1213 group: GroupId,
1214 },
1215 ValidReady,
1216 Token,
1217}
1218
1219impl Handshake {
1220 pub fn kind(&self) -> HandshakeKind {
1221 match self {
1222 Handshake::Fifo => HandshakeKind::Fifo,
1223 Handshake::Semaphore { .. } => HandshakeKind::Semaphore,
1224 Handshake::Barrier { .. } => HandshakeKind::Barrier,
1225 Handshake::ValidReady => HandshakeKind::ValidReady,
1226 Handshake::Token => HandshakeKind::Token,
1227 }
1228 }
1229 pub fn is_local(&self) -> bool {
1231 matches!(self, Handshake::Fifo | Handshake::ValidReady)
1232 }
1233}
1234
1235#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1236pub struct Edge {
1237 pub id: EdgeId,
1238 pub from: StageId,
1239 pub buffer: BufferId,
1240 pub to: Vec<StageId>,
1241 pub handshake: Handshake,
1242}
1243
1244#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1245pub struct TileRect {
1246 pub row0_t: usize,
1247 pub col0_t: usize,
1248 pub h_t: usize,
1249 pub w_t: usize,
1250}
1251
1252impl TileRect {
1253 pub fn tiles(&self, cols_t: usize) -> Vec<usize> {
1254 let mut v = Vec::with_capacity(self.h_t * self.w_t);
1255 for r in self.row0_t..self.row0_t + self.h_t {
1256 for c in self.col0_t..self.col0_t + self.w_t {
1257 v.push(r * cols_t + c);
1258 }
1259 }
1260 v
1261 }
1262}
1263
1264#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
1266pub enum FoldLevel {
1267 Prologue,
1269 Body,
1271 Epilogue,
1273}
1274
1275#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1276pub struct WorkItem {
1277 pub id: WorkItemId,
1278 pub coord: (usize, usize),
1280 pub out: TileRect,
1281}
1282
1283#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1284pub struct InnerLoop {
1285 pub num_blocks: usize,
1286 pub block_k_tiles: usize,
1287}
1288
1289#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1291pub struct BlockPlan {
1292 pub per_site_m: usize,
1293 pub per_site_n: usize,
1294 pub grid: (usize, usize),
1296 pub inner: InnerLoop,
1297 pub block: BlockShape,
1298 pub acc: AccPolicy,
1299 pub pack_accumulate: bool,
1302 pub engines: EnginePacking,
1304 pub granularity: Granularity,
1306}
1307
1308impl BlockPlan {
1309 pub fn site_block(
1323 &self,
1324 mt: usize,
1325 nt: usize,
1326 i: usize,
1327 j: usize,
1328 ) -> (usize, usize, usize, usize) {
1329 let h = self.per_site_m.min(mt.saturating_sub(i * self.per_site_m));
1330 let w = self.per_site_n.min(nt.saturating_sub(j * self.per_site_n));
1331 let fit = |nominal: usize, extent: usize| -> usize {
1332 (1..=nominal.min(extent))
1333 .rev()
1334 .find(|d| extent % d == 0)
1335 .unwrap_or(1)
1336 };
1337 (h, w, fit(self.block.osh, h), fit(self.block.osw, w))
1338 }
1339
1340 pub fn site_blocks(
1342 &self,
1343 mt: usize,
1344 nt: usize,
1345 ) -> Vec<((usize, usize), (usize, usize, usize, usize))> {
1346 let (ny, nx) = self.grid;
1347 let mut v = Vec::with_capacity(ny * nx);
1348 for i in 0..ny {
1349 for j in 0..nx {
1350 let b = self.site_block(mt, nt, i, j);
1351 if b.0 > 0 && b.1 > 0 {
1352 v.push(((i, j), b));
1353 }
1354 }
1355 }
1356 v
1357 }
1358}
1359
1360#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1365pub enum EnginePacking {
1366 LoadsTogether,
1368 OperandWithStore,
1370}
1371
1372#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1373pub struct WorkList {
1374 pub items: Vec<WorkItem>,
1375 pub k_tiles: usize,
1376 pub plan: Option<BlockPlan>,
1377 #[serde(default, skip_serializing_if = "Option::is_none")]
1380 pub k_trip: Option<TripId>,
1381}
1382
1383#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1384pub enum Direction {
1385 Row,
1386 Col,
1387}
1388
1389#[derive(PartialOrd, Ord, Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1390pub enum Operand {
1391 A,
1392 B,
1393}
1394
1395#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1396pub struct Group {
1397 pub id: GroupId,
1398 pub rect: Rect,
1399 pub sender: Site,
1400 pub direction: Direction,
1401 pub operand: Operand,
1402 pub noc: NocId,
1403}
1404
1405#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1406pub struct EngineSlot {
1407 pub site: Site,
1408 pub engine: Engine,
1409 pub noc: Option<NocId>,
1410}
1411
1412#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1413pub struct Placement {
1414 pub mesh: MeshShape,
1416 pub site_of: Vec<Site>,
1418 pub engine_of: Vec<EngineSlot>,
1420 pub groups: Vec<Group>,
1421 pub items_of_site: Vec<(Site, Vec<WorkItemId>)>,
1422}
1423
1424impl Placement {
1425 pub fn items_at(&self, site: Site) -> &[WorkItemId] {
1426 self.items_of_site
1427 .iter()
1428 .find(|(s, _)| *s == site)
1429 .map(|(_, v)| v.as_slice())
1430 .unwrap_or(&[])
1431 }
1432 pub fn sites(&self) -> BTreeSet<Site> {
1433 self.site_of.iter().copied().collect()
1434 }
1435}
1436
1437#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
1438pub struct BufferAt {
1439 pub site: Site,
1440 pub space: MemSpace,
1441 pub offset_bytes: u32,
1442 pub bytes: u32,
1443}
1444
1445#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1446pub struct MemoryMap {
1447 pub buffer_at: Vec<BufferAt>,
1449}
1450
1451#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1452pub enum SemRole {
1453 ReadyCounter {
1454 expected: usize,
1455 },
1456 ValidFlag,
1457 Barrier {
1459 n: usize,
1460 phase: BarrierPhase,
1461 },
1462}
1463
1464#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1470pub enum BarrierPhase {
1471 Gather,
1474 Release,
1477}
1478
1479#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1481pub struct Rendezvous {
1482 pub gather: SemId,
1483 pub release: SemId,
1484}
1485
1486#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1487pub struct Semaphore {
1488 pub id: SemId,
1489 pub init: u32,
1490 pub role: SemRole,
1491 pub operand: Option<Operand>,
1495 pub sites: Vec<Site>,
1497}
1498
1499#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1500pub enum SemValue {
1501 Invalid,
1502 Valid,
1503 Count(usize),
1504}
1505
1506impl SemValue {
1507 pub fn as_u32(self) -> u32 {
1508 match self {
1509 SemValue::Invalid => 0,
1510 SemValue::Valid => 1,
1511 SemValue::Count(n) => n as u32,
1512 }
1513 }
1514}
1515
1516#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1517pub enum SemCond {
1518 Equals(usize),
1519 AtLeast(usize),
1520}
1521
1522#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1524pub enum LoopCount {
1525 Const(usize),
1527 Trip { trip: TripId, max: usize },
1530}
1531
1532impl LoopCount {
1533 pub fn bound(self) -> usize {
1535 match self {
1536 LoopCount::Const(n) => n,
1537 LoopCount::Trip { max, .. } => max,
1538 }
1539 }
1540
1541 pub fn at(self, counts: &[usize]) -> usize {
1543 match self {
1544 LoopCount::Const(n) => n,
1545 LoopCount::Trip { trip, .. } => counts[trip.idx()],
1546 }
1547 }
1548
1549 pub fn trip(self) -> Option<TripId> {
1551 match self {
1552 LoopCount::Const(_) => None,
1553 LoopCount::Trip { trip, .. } => Some(trip),
1554 }
1555 }
1556}
1557
1558#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1570pub struct Trip {
1571 pub tensor: TensorId,
1572 pub elem: usize,
1573 pub admit: (i64, i64),
1574 pub bias: i64,
1575 pub per: usize,
1576}
1577
1578impl Trip {
1579 fn count(&self, v: i64) -> i128 {
1581 if self.per == 0 {
1582 return 0;
1583 }
1584 let (x, per) = (v as i128 + self.bias as i128, self.per as i128);
1585 if x <= 0 {
1586 -((-x) / per)
1587 } else {
1588 (x + per - 1) / per
1589 }
1590 }
1591
1592 pub fn lo(&self) -> usize {
1594 self.count(self.admit.0).max(0) as usize
1595 }
1596
1597 pub fn hi(&self) -> usize {
1599 self.count(self.admit.1).max(0) as usize
1600 }
1601
1602 pub fn count_of(&self, v: f64) -> Option<usize> {
1604 if v.fract() != 0.0 || !v.is_finite() || v < self.admit.0 as f64 || v > self.admit.1 as f64
1605 {
1606 return None;
1607 }
1608 Some(self.count(v as i64) as usize)
1609 }
1610}
1611
1612#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1613pub enum ArchCond {
1614 Always,
1615 Blackhole,
1616}
1617
1618#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1620pub enum ChoreoOp {
1621 Reserve {
1622 buf: BufferId,
1623 pages: usize,
1624 },
1625 Push {
1626 buf: BufferId,
1627 pages: usize,
1628 },
1629 Wait {
1630 buf: BufferId,
1631 pages: usize,
1632 },
1633 Pop {
1634 buf: BufferId,
1635 pages: usize,
1636 },
1637 ReadTiles {
1638 tensor: TensorId,
1639 at: TileAddr,
1640 block_h: usize,
1641 block_w: usize,
1642 stride_w: usize,
1643 stride_h: usize,
1644 into: BufferId,
1645 },
1646 ReadBarrier,
1647 WriteTiles {
1648 from: BufferId,
1649 tensor: TensorId,
1650 at: TileAddr,
1651 block_h: usize,
1652 block_w: usize,
1653 stride_w: usize,
1654 stride_h: usize,
1655 },
1656 WriteBarrier,
1657 McastSend {
1662 buf: BufferId,
1663 pages: usize,
1664 group: GroupId,
1665 hold_path: bool,
1666 },
1667 McastSetValid {
1668 sem: SemId,
1669 group: GroupId,
1670 },
1671 Flush {
1672 cond: ArchCond,
1673 },
1674 SemSet {
1675 sem: SemId,
1676 value: SemValue,
1677 },
1678 SemInc {
1679 sem: SemId,
1680 at: Site,
1681 by: usize,
1682 },
1683 SemWait {
1684 sem: SemId,
1685 until: SemCond,
1686 },
1687 AcquireAcc {
1688 tiles: usize,
1689 },
1690 ReleaseAcc,
1691 Compute {
1692 op: ComputeOp,
1693 ins: Vec<BufferId>,
1694 acc_reload: Option<BufferId>,
1695 },
1696 PackAcc {
1700 to: BufferId,
1701 tiles: usize,
1702 convert: Precision,
1703 accumulate: bool,
1704 },
1705 Barrier {
1706 group: GroupId,
1707 },
1708 Loop {
1710 var: LoopVar,
1711 count: LoopCount,
1712 body: Vec<TaggedOp>,
1713 },
1714}
1715
1716#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
1717pub struct TaggedOp {
1718 pub stage: StageId,
1719 pub op: ChoreoOp,
1720}
1721
1722#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1723pub struct EngineProgram {
1724 pub site: Site,
1725 pub engine: Engine,
1726 pub stages: Vec<StageId>,
1727 pub ops: Vec<TaggedOp>,
1728}
1729
1730#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1731pub struct Schedule {
1732 pub programs: Vec<EngineProgram>,
1733}
1734
1735impl Schedule {
1736 pub fn program_at(&self, site: Site, engine: Engine) -> Option<&EngineProgram> {
1737 self.programs
1738 .iter()
1739 .find(|p| p.site == site && p.engine == engine)
1740 }
1741 pub fn per_stage(&self, stage: StageId) -> Vec<&ChoreoOp> {
1744 fn collect<'a>(ops: &'a [TaggedOp], stage: StageId, out: &mut Vec<&'a ChoreoOp>) {
1745 for t in ops {
1746 if t.stage == stage {
1747 out.push(&t.op);
1748 }
1749 if let ChoreoOp::Loop { body, .. } = &t.op {
1750 collect(body, stage, out);
1751 }
1752 }
1753 }
1754 let mut out = Vec::new();
1755 for p in &self.programs {
1756 collect(&p.ops, stage, &mut out);
1757 }
1758 out
1759 }
1760}
1761
1762#[derive(Clone, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
1765pub struct Requirements {
1766 pub handshakes: BTreeSet<HandshakeKind>,
1767 pub interconnect: BTreeSet<Interconnect>,
1768}
1769
1770#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
1771pub enum Phase {
1772 Template,
1773 Tiled,
1774 Placed,
1775 Sized,
1776 Scheduled,
1777 Verified,
1778}
1779
1780#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
1782pub enum IrViolation {
1783 IdNotDense {
1784 what: &'static str,
1785 index: usize,
1786 },
1787 DanglingTensor {
1788 stage: StageId,
1789 tensor: TensorId,
1790 },
1791 DanglingBuffer {
1792 stage: Option<StageId>,
1793 buffer: BufferId,
1794 },
1795 DanglingStage {
1796 edge: EdgeId,
1797 stage: StageId,
1798 },
1799 DanglingGroup {
1800 stage: StageId,
1801 group: GroupId,
1802 },
1803 DanglingSemaphore {
1804 sem: SemId,
1805 },
1806 DanglingItem {
1807 stage: StageId,
1808 item: WorkItemId,
1809 },
1810 EdgeBufferNotInStagePorts {
1811 edge: EdgeId,
1812 stage: StageId,
1813 buffer: BufferId,
1814 },
1815 FifoProducerCount {
1816 buffer: BufferId,
1817 count: usize,
1818 },
1819 FifoConsumerCount {
1820 buffer: BufferId,
1821 count: usize,
1822 },
1823 CarryWithoutInit {
1825 buffer: BufferId,
1826 },
1827 CarryUpdatedTwice {
1829 buffer: BufferId,
1830 },
1831 CarryInitMovesWithTheBlocks {
1833 buffer: BufferId,
1834 },
1835 CarryNotARing {
1837 buffer: BufferId,
1838 },
1839 LoadReadsOutput {
1840 stage: StageId,
1841 tensor: TensorId,
1842 },
1843 StoreWritesInput {
1844 stage: StageId,
1845 tensor: TensorId,
1846 },
1847 TileGridMismatch {
1848 tensor: TensorId,
1849 },
1850 OutputUncovered {
1851 tensor: TensorId,
1852 tile: usize,
1853 },
1854 OutputOverlap {
1855 tensor: TensorId,
1856 tile: usize,
1857 },
1858 OperandDtypeMismatch {
1859 stage: StageId,
1860 buffer: BufferId,
1861 expected: Precision,
1862 found: Precision,
1863 },
1864 IntermNarrowerThanAcc {
1865 stage: StageId,
1866 interm: BufferId,
1867 acc: AccPolicy,
1868 found: Precision,
1869 },
1870 AccTilesExceeded {
1871 stage: StageId,
1872 needed: usize,
1873 available: usize,
1874 },
1875 BlockDivisibility {
1876 what: &'static str,
1877 value: usize,
1878 divisor: usize,
1879 },
1880 GridMismatch {
1881 expected: (usize, usize),
1882 found: (usize, usize),
1883 },
1884 PhaseOrder {
1885 have: &'static str,
1886 missing: &'static str,
1887 },
1888 StageUnplaced {
1889 stage: StageId,
1890 },
1891 EngineRoleMismatch {
1892 stage: StageId,
1893 },
1894 EngineSiteMismatch {
1895 stage: StageId,
1896 },
1897 GroupMalformed {
1898 group: GroupId,
1899 reason: &'static str,
1900 },
1901 GroupMemberMismatch {
1902 group: GroupId,
1903 site: Site,
1904 reason: &'static str,
1905 },
1906 CrossSiteLocalEdge {
1907 edge: EdgeId,
1908 },
1909 LocalCrossSiteEdge {
1910 edge: EdgeId,
1911 },
1912 MemoryOverlap {
1913 site: Site,
1914 a: BufferId,
1915 b: BufferId,
1916 },
1917 MemoryUnmapped {
1918 buffer: BufferId,
1919 },
1920 DuplicateEngineProgram {
1921 site: Site,
1922 engine: Engine,
1923 },
1924 CountingImbalance {
1925 buffer: BufferId,
1926 reserved: u64,
1927 pushed: u64,
1928 popped: u64,
1929 },
1930 AccumulatorNeverComputed {
1933 site: Site,
1934 engine: Engine,
1935 },
1936 SemaphoreImbalance {
1937 sem: SemId,
1938 incs: u64,
1939 waits: u64,
1940 },
1941 SemaphoreNotReset {
1944 sem: SemId,
1945 waits: u64,
1946 resets: u64,
1947 },
1948 AccImbalance {
1949 site: Site,
1950 acquires: u64,
1951 releases: u64,
1952 },
1953 HandshakeNotRequired {
1954 kind: HandshakeKind,
1955 },
1956 ItemCountMismatch {
1957 group: GroupId,
1958 },
1959 StageNotScheduled {
1962 stage: StageId,
1963 op: &'static str,
1964 },
1965 GenericOperandCount {
1967 stage: StageId,
1968 operands: usize,
1969 maps: usize,
1970 },
1971 GenericOperandOutOfRange {
1974 stage: StageId,
1975 operand: usize,
1976 operands: usize,
1977 },
1978 ContractionMalformed {
1981 stage: StageId,
1982 reason: &'static str,
1983 },
1984 IndirectOperandOutOfRange {
1986 stage: StageId,
1987 operand: usize,
1988 operands: usize,
1989 },
1990 IndirectMalformed {
1993 stage: StageId,
1994 reason: &'static str,
1995 },
1996 TripMalformed {
1999 trip: TripId,
2000 reason: &'static str,
2001 },
2002 TripLoopMalformed {
2006 reason: &'static str,
2007 },
2008 IndirectOutWithoutScatter {
2011 stage: StageId,
2012 },
2013 IndirectIndexNotAnOperand {
2015 stage: StageId,
2016 buffer: BufferId,
2017 },
2018 ArgMaxOutputNotInteger {
2020 stage: StageId,
2021 found: Precision,
2022 },
2023 TensorRoleDeferred {
2026 tensor: TensorId,
2027 role: TensorRole,
2028 until: &'static str,
2029 },
2030 TensorSpaceDeferred {
2032 tensor: TensorId,
2033 space: TensorSpace,
2034 until: &'static str,
2035 },
2036 SectionSpanMalformed {
2039 section: usize,
2040 what: &'static str,
2041 },
2042 SectionBarrierUnbalanced {
2045 sem: SemId,
2046 participants: usize,
2047 sites: usize,
2048 incs: u64,
2049 wait_ops: u64,
2050 },
2051 SectionBarrierMisplaced {
2054 section: usize,
2055 site: Site,
2056 engine: Engine,
2057 what: &'static str,
2058 },
2059 SectionBufferEscape {
2062 stage: StageId,
2063 buffer: BufferId,
2064 stage_section: usize,
2065 buffer_section: usize,
2066 },
2067}
2068
2069#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
2071pub struct PlanSummary {
2072 pub phase: Phase,
2073 pub sites_used: usize,
2074 pub grid: (usize, usize),
2075 pub per_site_block: (usize, usize),
2076 pub block: Option<BlockShape>,
2077 pub num_blocks: usize,
2078 pub streaming: bool,
2079 pub multicast_groups: usize,
2080 pub row_groups: usize,
2081 pub col_groups: usize,
2082 pub acc: Option<AccPolicy>,
2083 pub dram_tiles_read: u64,
2084 pub semaphores: usize,
2085 pub fifos: usize,
2086 pub engine_program_classes: usize,
2087 pub engines: EnginePacking,
2089}
2090
2091#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
2093pub enum DimAxis {
2094 Rows,
2095 Cols,
2096}
2097
2098#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
2108pub struct SectionSpan {
2109 pub name: String,
2110 pub stages: (usize, usize),
2112 pub buffers: (usize, usize),
2113 pub edges: (usize, usize),
2114 pub semaphores: (usize, usize),
2115 pub items: (usize, usize),
2116 pub plan: Option<BlockPlan>,
2117 pub k_tiles: usize,
2118 pub output: Option<TensorId>,
2120 pub barrier_after: Option<Rendezvous>,
2123}
2124
2125pub struct SiteItems<'a> {
2127 sections: usize,
2128 whole: BTreeMap<Site, &'a [WorkItemId]>,
2130 by_section: BTreeMap<(usize, Site), Vec<WorkItemId>>,
2131}
2132
2133impl SiteItems<'_> {
2134 pub fn at(&self, section: Option<usize>, site: Site) -> &[WorkItemId] {
2136 match section.filter(|k| *k < self.sections) {
2137 Some(k) => self
2138 .by_section
2139 .get(&(k, site))
2140 .map(Vec::as_slice)
2141 .unwrap_or(&[]),
2142 None => self.whole.get(&site).copied().unwrap_or(&[]),
2143 }
2144 }
2145}
2146
2147impl SectionSpan {
2148 pub fn holds_stage(&self, s: StageId) -> bool {
2149 s.idx() >= self.stages.0 && s.idx() < self.stages.1
2150 }
2151 pub fn holds_buffer(&self, b: BufferId) -> bool {
2152 b.idx() >= self.buffers.0 && b.idx() < self.buffers.1
2153 }
2154 pub fn holds_item(&self, i: WorkItemId) -> bool {
2155 i.idx() >= self.items.0 && i.idx() < self.items.1
2156 }
2157}
2158
2159#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
2160pub struct SpatialProgram {
2161 pub name: String,
2162 pub tile: TileShape,
2163 pub tensors: Vec<TensorDecl>,
2164 pub stages: Vec<Stage>,
2165 pub buffers: Vec<Buffer>,
2166 pub edges: Vec<Edge>,
2167 pub semaphores: Vec<Semaphore>,
2168 #[serde(default, skip_serializing_if = "Vec::is_empty")]
2172 pub sections: Vec<SectionSpan>,
2173 #[serde(default, skip_serializing_if = "Vec::is_empty")]
2176 pub trips: Vec<Trip>,
2177 #[serde(default, skip_serializing_if = "Vec::is_empty")]
2180 pub links: Vec<LinkPoint>,
2181 pub work: WorkList,
2182 pub numerics: NumericsPolicy,
2183 pub claim: NumericsClaim,
2184 pub acc_tiles_used: usize,
2186 #[serde(default = "one_bank", skip_serializing_if = "is_one_bank")]
2190 pub mem_banks: usize,
2191 pub requires: Requirements,
2192 pub placement: Option<Placement>,
2193 pub memory: Option<MemoryMap>,
2194 pub schedule: Option<Schedule>,
2195 #[serde(skip)]
2196 pub certificate: Option<ChoreographyCertificate>,
2197}
2198
2199#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
2209pub struct LinkPoint {
2210 pub after: Option<usize>,
2211 pub sends: Vec<TensorId>,
2212 pub receives: Vec<TensorId>,
2213}
2214
2215pub fn walk_ops(ops: &[ChoreoOp], mult: u64, f: &mut dyn FnMut(&ChoreoOp, u64)) {
2218 for op in ops {
2219 f(op, mult);
2220 if let ChoreoOp::Loop { count, body, .. } = op {
2221 walk_tagged_ops(body, mult * count.bound() as u64, f);
2222 }
2223 }
2224}
2225
2226pub fn walk_tagged_ops(ops: &[TaggedOp], mult: u64, f: &mut dyn FnMut(&ChoreoOp, u64)) {
2228 for t in ops {
2229 f(&t.op, mult);
2230 if let ChoreoOp::Loop { count, body, .. } = &t.op {
2231 walk_tagged_ops(body, mult * count.bound() as u64, f);
2232 }
2233 }
2234}
2235
2236pub fn walk_tagged(ops: &[TaggedOp], mult: u64, f: &mut dyn FnMut(&TaggedOp, u64)) {
2239 walk_tagged_counted(ops, mult, &|c| c.bound(), f);
2240}
2241
2242pub fn walk_tagged_counted(
2245 ops: &[TaggedOp],
2246 mult: u64,
2247 count: &dyn Fn(LoopCount) -> usize,
2248 f: &mut dyn FnMut(&TaggedOp, u64),
2249) {
2250 for t in ops {
2251 f(t, mult);
2252 if let ChoreoOp::Loop { count: c, body, .. } = &t.op {
2253 walk_tagged_counted(body, mult * count(*c) as u64, count, f);
2254 }
2255 }
2256}
2257
2258pub fn body_ops(body: &[TaggedOp]) -> Vec<ChoreoOp> {
2260 body.iter().map(|t| t.op.clone()).collect()
2261}
2262
2263impl SpatialProgram {
2264 pub fn tensor(&self, id: TensorId) -> &TensorDecl {
2265 &self.tensors[id.idx()]
2266 }
2267 pub fn stage(&self, id: StageId) -> &Stage {
2268 &self.stages[id.idx()]
2269 }
2270 pub fn buffer(&self, id: BufferId) -> &Buffer {
2271 &self.buffers[id.idx()]
2272 }
2273 pub fn page_extents(&self) -> Vec<usize> {
2283 let mut extent: Vec<usize> = self.buffers.iter().map(|b| b.kind.page_tiles()).collect();
2284 let Some(schedule) = &self.schedule else {
2285 return extent;
2286 };
2287 for program in &schedule.programs {
2288 let mut acquired = 0;
2289 walk_tagged_ops(&program.ops, 1, &mut |op, _| {
2290 let (b, tiles) = match op {
2291 ChoreoOp::ReadTiles { into, block_h, block_w, .. } => (*into, block_h * block_w),
2292 ChoreoOp::WriteTiles { from, block_h, block_w, .. } => (*from, block_h * block_w),
2293 ChoreoOp::PackAcc { to, tiles, .. } => (*to, *tiles),
2294 ChoreoOp::Compute { acc_reload: Some(r), .. } => (*r, acquired),
2295 ChoreoOp::AcquireAcc { tiles } => {
2296 acquired = *tiles;
2297 return;
2298 }
2299 _ => return,
2300 };
2301 extent[b.idx()] = extent[b.idx()].max(tiles);
2302 });
2303 }
2304 extent
2305 }
2306 pub fn output_tensor(&self) -> Option<&TensorDecl> {
2307 self.tensors.iter().find(|t| t.role == TensorRole::Output).or_else(|| {
2310 self.tensors.iter().find(|t| t.role == TensorRole::State && self.stages.iter().any(|s| matches!(&s.kind, StageKind::Store { tensor, .. } if *tensor == t.id)))
2311 })
2312 }
2313
2314 pub fn section_result(&self, section: Option<usize>) -> Option<&TensorDecl> {
2318 match section.and_then(|k| self.sections.get(k)) {
2319 Some(sp) => sp.output.map(|t| self.tensor(t)).or_else(|| {
2320 self.stages[sp.stages.0..sp.stages.1]
2321 .iter()
2322 .find_map(|s| match &s.kind {
2323 StageKind::Store { tensor, .. }
2324 if self.tensor(*tensor).role == TensorRole::State =>
2325 {
2326 Some(self.tensor(*tensor))
2327 }
2328 _ => None,
2329 })
2330 }),
2331 None => self.output_tensor(),
2332 }
2333 }
2334
2335 pub fn carries(&self) -> BTreeMap<BufferId, StageId> {
2339 let mut out = BTreeMap::new();
2340 for s in &self.stages {
2341 for b in &s.outs {
2342 if s.ins.contains(b) {
2343 out.entry(*b).or_insert(s.id);
2344 }
2345 }
2346 }
2347 out
2348 }
2349
2350 pub fn fold_levels(&self) -> Vec<FoldLevel> {
2357 let carries = self.carries();
2358 let mut level = vec![FoldLevel::Prologue; self.stages.len()];
2359 let mut producers: BTreeMap<BufferId, Vec<StageId>> = BTreeMap::new();
2360 let mut readers: BTreeMap<BufferId, Vec<StageId>> = BTreeMap::new();
2361 for s in &self.stages {
2362 for b in &s.outs {
2363 producers.entry(*b).or_default().push(s.id);
2364 }
2365 for b in &s.ins {
2366 readers.entry(*b).or_default().push(s.id);
2367 }
2368 }
2369 loop {
2370 let mut changed = false;
2371 for s in &self.stages {
2372 let reads = |l: FoldLevel| {
2373 s.ins.iter().filter(|b| !carries.contains_key(b)).any(|b| {
2374 producers
2375 .get(b)
2376 .is_some_and(|ps| ps.iter().any(|p| level[p.idx()] == l))
2377 })
2378 };
2379 let reads_carry = s.ins.iter().any(|b| carries.contains_key(b));
2380 let feeds_body = s.outs.iter().filter(|b| !carries.contains_key(b)).any(|b| {
2381 readers.get(b).is_some_and(|rs| {
2382 rs.iter()
2383 .any(|r| *r != s.id && level[r.idx()] == FoldLevel::Body)
2384 })
2385 });
2386 let next = match &s.kind {
2387 StageKind::Load { pattern, .. }
2388 if pattern
2389 .origin
2390 .steps
2391 .iter()
2392 .any(|(v, _)| *v == super::passes::LV_KBLOCK) =>
2393 {
2394 FoldLevel::Body
2395 }
2396 StageKind::Recv { .. } => FoldLevel::Body,
2397 _ if carries.values().any(|c| *c == s.id)
2398 || reads(FoldLevel::Body)
2399 || (reads_carry && feeds_body) =>
2400 {
2401 FoldLevel::Body
2402 }
2403 _ if reads_carry || reads(FoldLevel::Epilogue) => FoldLevel::Epilogue,
2404 _ => FoldLevel::Prologue,
2405 };
2406 if next != level[s.id.idx()] {
2407 level[s.id.idx()] = next;
2408 changed = true;
2409 }
2410 }
2411 if !changed {
2412 return level;
2413 }
2414 }
2415 }
2416
2417 pub fn section_of_stage(&self, s: StageId) -> Option<usize> {
2420 self.section_holding(s.idx(), |sp| sp.stages)
2421 }
2422
2423 pub fn section_of_buffer(&self, b: BufferId) -> Option<usize> {
2425 self.section_holding(b.idx(), |sp| sp.buffers)
2426 }
2427
2428 pub fn section_of_item(&self, i: WorkItemId) -> Option<usize> {
2430 self.section_holding(i.idx(), |sp| sp.items)
2431 }
2432
2433 fn section_holding(
2442 &self,
2443 id: usize,
2444 span: impl Fn(&SectionSpan) -> (usize, usize),
2445 ) -> Option<usize> {
2446 let holds = |sp: &SectionSpan| {
2447 let (a, b) = span(sp);
2448 a <= id && id < b
2449 };
2450 let k = self.sections.partition_point(|sp| span(sp).1 <= id);
2451 match self.sections.get(k) {
2452 Some(sp) if holds(sp) => Some(k),
2453 _ => self.sections.iter().position(holds),
2454 }
2455 }
2456
2457 pub fn items_at_in(&self, section: Option<usize>, site: Site) -> Vec<WorkItemId> {
2460 let all = self
2461 .placement
2462 .as_ref()
2463 .map(|pl| pl.items_at(site).to_vec())
2464 .unwrap_or_default();
2465 match section.and_then(|k| self.sections.get(k)) {
2466 Some(sp) => all.into_iter().filter(|i| sp.holds_item(*i)).collect(),
2467 None => all,
2468 }
2469 }
2470
2471 pub fn site_items(&self) -> SiteItems<'_> {
2474 let mut whole: BTreeMap<Site, &[WorkItemId]> = BTreeMap::new();
2475 let mut by_section: BTreeMap<(usize, Site), Vec<WorkItemId>> = BTreeMap::new();
2476 for (site, items) in self.placement.iter().flat_map(|pl| pl.items_of_site.iter()) {
2477 if whole.contains_key(site) {
2478 continue;
2479 }
2480 whole.insert(*site, items);
2481 for i in items {
2482 if let Some(k) = self.section_of_item(*i) {
2483 by_section.entry((k, *site)).or_default().push(*i);
2484 }
2485 }
2486 }
2487 SiteItems {
2488 sections: self.sections.len(),
2489 whole,
2490 by_section,
2491 }
2492 }
2493
2494 pub fn plan_of_stage(&self, s: StageId) -> Option<BlockPlan> {
2496 self.section_of_stage(s)
2497 .and_then(|k| self.sections[k].plan)
2498 .or(self.work.plan)
2499 }
2500
2501 pub fn state_tensors(&self) -> Vec<&TensorDecl> {
2503 self.tensors
2504 .iter()
2505 .filter(|t| t.role == TensorRole::State)
2506 .collect()
2507 }
2508
2509 fn buffer_tensor(&self, b: BufferId) -> Option<TensorId> {
2512 let template = self.buffers[b.idx()].origin.unwrap_or(b);
2513 self.stages.iter().find_map(|s| match &s.kind {
2514 StageKind::Load { tensor, .. }
2515 if s.outs.iter().any(|o| {
2516 *o == b || *o == template || self.buffers[o.idx()].origin == Some(template)
2517 }) =>
2518 {
2519 Some(*tensor)
2520 }
2521 StageKind::Store { tensor, .. }
2522 if s.ins.iter().any(|i| {
2523 *i == b || *i == template || self.buffers[i.idx()].origin == Some(template)
2524 }) =>
2525 {
2526 Some(*tensor)
2527 }
2528 _ => None,
2529 })
2530 }
2531
2532 fn stage_output_tensor(&self, stage: &Stage, depth: usize) -> Option<TensorId> {
2538 if depth > self.stages.len() {
2539 return None;
2540 }
2541 for out in &stage.outs {
2542 if let Some(t) = self.buffer_tensor(*out) {
2543 return Some(t);
2544 }
2545 if let Some(next) = self.stages.iter().find(|s| {
2546 s.id != stage.id && s.ins.contains(out) && matches!(s.kind, StageKind::Compute(_))
2547 }) {
2548 if let Some(t) = self.stage_output_tensor(next, depth + 1) {
2549 return Some(t);
2550 }
2551 }
2552 }
2553 None
2554 }
2555
2556 pub fn dim_classes(&self) -> Vec<Vec<(TensorId, DimAxis)>> {
2561 let n = self.tensors.len() * 2;
2562 let key = |t: TensorId, a: DimAxis| {
2563 t.idx() * 2
2564 + match a {
2565 DimAxis::Rows => 0,
2566 DimAxis::Cols => 1,
2567 }
2568 };
2569 let mut parent: Vec<usize> = (0..n).collect();
2570 fn find(parent: &mut Vec<usize>, i: usize) -> usize {
2571 let mut r = i;
2572 while parent[r] != r {
2573 r = parent[r];
2574 }
2575 let mut c = i;
2576 while parent[c] != r {
2577 let next = parent[c];
2578 parent[c] = r;
2579 c = next;
2580 }
2581 r
2582 }
2583 let union = |parent: &mut Vec<usize>, a: usize, b: usize| {
2584 let (ra, rb) = (find(parent, a), find(parent, b));
2585 if ra != rb {
2586 parent[ra.max(rb)] = ra.min(rb);
2587 }
2588 };
2589 for s in &self.stages {
2590 let StageKind::Compute(op) = &s.kind else {
2591 continue;
2592 };
2593 let ins: Vec<Option<TensorId>> = s.ins.iter().map(|b| self.buffer_tensor(*b)).collect();
2594 let out = self.stage_output_tensor(s, 0);
2595 match op {
2596 ComputeOp::MatmulTiles { .. } => {
2597 if let (Some(Some(a)), Some(Some(b))) = (ins.first(), ins.get(1)) {
2598 union(&mut parent, key(*a, DimAxis::Cols), key(*b, DimAxis::Rows));
2599 if let Some(c) = out {
2600 union(&mut parent, key(*a, DimAxis::Rows), key(c, DimAxis::Rows));
2601 union(&mut parent, key(*b, DimAxis::Cols), key(c, DimAxis::Cols));
2602 }
2603 }
2604 }
2605 ComputeOp::Generic {
2611 iters,
2612 ins: maps,
2613 out: omap,
2614 ..
2615 } => {
2616 for (ax, dim) in [
2617 (iters.first(), DimAxis::Rows),
2618 (iters.get(1), DimAxis::Cols),
2619 ] {
2620 let Some(ax) = ax else { continue };
2621 let mut sharing: Vec<TensorId> = ins
2622 .iter()
2623 .enumerate()
2624 .filter(|(i, _)| {
2625 maps.get(*i).map(|m| m.varies_along(ax.var)).unwrap_or(true)
2626 })
2627 .filter_map(|(_, t)| *t)
2628 .collect();
2629 if omap.varies_along(ax.var) {
2630 sharing.extend(out);
2631 }
2632 for w in sharing.windows(2) {
2633 union(&mut parent, key(w[0], dim), key(w[1], dim));
2634 }
2635 }
2636 }
2637 _ => {
2638 if let (Some(Some(x)), Some(o)) = (ins.first(), out) {
2640 union(&mut parent, key(*x, DimAxis::Rows), key(o, DimAxis::Rows));
2641 union(&mut parent, key(*x, DimAxis::Cols), key(o, DimAxis::Cols));
2642 }
2643 }
2644 }
2645 }
2646 let mut classes: std::collections::BTreeMap<usize, Vec<(TensorId, DimAxis)>> =
2647 std::collections::BTreeMap::new();
2648 for t in &self.tensors {
2649 for a in [DimAxis::Rows, DimAxis::Cols] {
2650 let k = key(t.id, a);
2651 classes
2652 .entry(find(&mut parent, k))
2653 .or_default()
2654 .push((t.id, a));
2655 }
2656 }
2657 classes.into_values().collect()
2658 }
2659
2660 pub fn dim_representative(&self, tensor: TensorId, axis: DimAxis) -> (TensorId, DimAxis) {
2663 let classes = self.dim_classes();
2664 let class = classes
2665 .iter()
2666 .find(|c| c.contains(&(tensor, axis)))
2667 .cloned()
2668 .unwrap_or_else(|| vec![(tensor, axis)]);
2669 class
2670 .iter()
2671 .copied()
2672 .find(|(t, _)| self.tensors[t.idx()].role == TensorRole::Output)
2673 .unwrap_or(class[0])
2674 }
2675 pub fn buffer_bytes(&self, b: &Buffer) -> usize {
2676 b.kind.tiles() * self.tile.elems() * precision_bytes(b.dtype)
2677 }
2678
2679 pub fn phase(&self) -> Phase {
2680 if self.certificate.is_some() {
2681 Phase::Verified
2682 } else if self.schedule.is_some() {
2683 Phase::Scheduled
2684 } else if self.memory.is_some() {
2685 Phase::Sized
2686 } else if self.placement.is_some() {
2687 Phase::Placed
2688 } else if self.work.plan.is_some() {
2689 Phase::Tiled
2690 } else {
2691 Phase::Template
2692 }
2693 }
2694
2695 pub fn neutral_dump(&self) -> String {
2697 serde_json::to_string_pretty(self).expect("the spatial IR serializes")
2698 }
2699
2700 pub fn parse_neutral(s: &str) -> Result<SpatialProgram, String> {
2701 serde_json::from_str(s).map_err(|e| e.to_string())
2702 }
2703
2704 pub fn sites(&self) -> BTreeSet<Site> {
2706 self.stages.iter().filter_map(|s| s.site).collect()
2707 }
2708
2709 pub fn stage_group(&self, s: &Stage) -> Option<GroupId> {
2711 match &s.kind {
2712 StageKind::Load { multicast, .. } => *multicast,
2713 StageKind::Recv { group } => Some(*group),
2714 _ => None,
2715 }
2716 }
2717
2718 pub fn dram_tiles_read(&self) -> u64 {
2720 let items_template = self.work.items.len().max(1) as u64;
2721 let mut total = 0u64;
2722 for s in &self.stages {
2723 if let StageKind::Load {
2724 pattern,
2725 num_blocks,
2726 ..
2727 } = &s.kind
2728 {
2729 let per_item = pattern.block_tiles() as u64 * (*num_blocks).max(1) as u64;
2730 let items = match (&self.placement, s.site) {
2731 (Some(p), Some(site)) => p.items_at(site).len().max(1) as u64,
2732 _ => items_template,
2733 };
2734 total += per_item * items;
2735 }
2736 }
2737 total
2738 }
2739
2740 pub fn engine_signature(&self, p: &EngineProgram) -> String {
2744 let buf_key = |b: BufferId| -> String {
2745 let bb = self.buffer(b);
2746 format!(
2747 "b{}p{}",
2748 bb.origin.map(|o| o.idx()).unwrap_or(b.idx()),
2749 bb.kind.page_tiles()
2750 )
2751 };
2752 let sem_key = |s: SemId| -> String {
2753 match self.semaphores.get(s.idx()) {
2754 Some(sem) => format!(
2755 "s{:?}{}",
2756 sem.role,
2757 sem.operand.map(|o| format!("{o:?}")).unwrap_or_default()
2758 ),
2759 None => format!("s?{}", s.idx()),
2760 }
2761 };
2762 let group_key = |g: GroupId| -> String {
2763 self.placement
2764 .as_ref()
2765 .and_then(|pl| pl.groups.get(g.idx()))
2766 .map(|gr| format!("g{:?}{:?}", gr.operand, gr.direction))
2767 .unwrap_or_else(|| "g".into())
2768 };
2769 fn sig(
2770 ops: &[TaggedOp],
2771 out: &mut String,
2772 buf_key: &dyn Fn(BufferId) -> String,
2773 sem_key: &dyn Fn(SemId) -> String,
2774 group_key: &dyn Fn(GroupId) -> String,
2775 ) {
2776 for t in ops {
2777 match &t.op {
2778 ChoreoOp::Reserve { buf, pages } => {
2779 out.push_str(&format!("R({},{});", buf_key(*buf), pages))
2780 }
2781 ChoreoOp::Push { buf, pages } => {
2782 out.push_str(&format!("P({},{});", buf_key(*buf), pages))
2783 }
2784 ChoreoOp::Wait { buf, pages } => {
2785 out.push_str(&format!("W({},{});", buf_key(*buf), pages))
2786 }
2787 ChoreoOp::Pop { buf, pages } => {
2788 out.push_str(&format!("O({},{});", buf_key(*buf), pages))
2789 }
2790 ChoreoOp::ReadTiles {
2791 tensor,
2792 at,
2793 block_h,
2794 block_w,
2795 stride_w,
2796 stride_h,
2797 into,
2798 } => out.push_str(&format!(
2799 "RT(t{},{:?},{block_h},{block_w},{stride_w},{stride_h},{});",
2800 tensor.idx(),
2801 at.steps,
2802 buf_key(*into)
2803 )),
2804 ChoreoOp::ReadBarrier => out.push_str("RB;"),
2805 ChoreoOp::WriteTiles {
2806 from,
2807 tensor,
2808 at,
2809 block_h,
2810 block_w,
2811 stride_w,
2812 stride_h,
2813 } => out.push_str(&format!(
2814 "WT({},t{},{:?},{block_h},{block_w},{stride_w},{stride_h});",
2815 buf_key(*from),
2816 tensor.idx(),
2817 at.steps
2818 )),
2819 ChoreoOp::WriteBarrier => out.push_str("WB;"),
2820 ChoreoOp::McastSend {
2821 buf,
2822 pages,
2823 group,
2824 hold_path,
2825 } => out.push_str(&format!(
2826 "MS({},{hold_path},{pages},{});",
2827 buf_key(*buf),
2828 group_key(*group)
2829 )),
2830 ChoreoOp::McastSetValid { sem, group } => {
2831 out.push_str(&format!("MV({},{});", sem_key(*sem), group_key(*group)))
2832 }
2833 ChoreoOp::Flush { cond } => out.push_str(&format!("F{cond:?};")),
2834 ChoreoOp::SemSet { sem, value } => {
2835 out.push_str(&format!("SS({},{value:?});", sem_key(*sem)))
2836 }
2837 ChoreoOp::SemInc { sem, by, .. } => {
2838 out.push_str(&format!("SI({},{by});", sem_key(*sem)))
2839 }
2840 ChoreoOp::SemWait { sem, until } => {
2841 out.push_str(&format!("SW({},{until:?});", sem_key(*sem)))
2842 }
2843 ChoreoOp::AcquireAcc { tiles } => out.push_str(&format!("AA({tiles});")),
2844 ChoreoOp::ReleaseAcc => out.push_str("RA;"),
2845 ChoreoOp::Compute {
2846 op,
2847 ins,
2848 acc_reload,
2849 } => {
2850 let op_key = match op {
2852 ComputeOp::MatmulTiles { block, sub_blocks, num_blocks, operand, acc, dataflow, interm, granularity } => format!(
2853 "MM({block:?},{sub_blocks:?},{num_blocks},{operand:?},{acc:?},{dataflow:?},{:?},{granularity:?})",
2854 interm.map(buf_key)
2855 ),
2856 ComputeOp::Bias { bias } => format!("Bias({})", buf_key(*bias)),
2857 other => format!("{other:?}"),
2858 };
2859 out.push_str(&format!(
2860 "C({op_key},[{}],{:?});",
2861 ins.iter()
2862 .map(|b| buf_key(*b))
2863 .collect::<Vec<_>>()
2864 .join(","),
2865 acc_reload.map(buf_key)
2866 ))
2867 }
2868 ChoreoOp::PackAcc {
2869 to,
2870 tiles,
2871 convert,
2872 accumulate,
2873 } => out.push_str(&format!(
2874 "PA({},{tiles},{convert:?},{accumulate});",
2875 buf_key(*to)
2876 )),
2877 ChoreoOp::Barrier { group } => {
2878 out.push_str(&format!("B({});", group_key(*group)))
2879 }
2880 ChoreoOp::Loop { body, .. } => {
2881 out.push_str("L{");
2882 sig(body, out, buf_key, sem_key, group_key);
2883 out.push_str("}");
2884 }
2885 }
2886 }
2887 }
2888 let mut out = format!("{:?}:", p.engine);
2889 sig(&p.ops, &mut out, &buf_key, &sem_key, &group_key);
2890 out
2891 }
2892
2893 pub fn summary(&self) -> PlanSummary {
2895 let plan = self.work.plan;
2896 let (grid, per_site_block, block, num_blocks, acc) = match plan {
2897 Some(p) => (
2898 p.grid,
2899 (p.per_site_m, p.per_site_n),
2900 Some(p.block),
2901 p.inner.num_blocks,
2902 Some(p.acc),
2903 ),
2904 None => ((0, 0), (0, 0), None, 0, None),
2905 };
2906 let streaming = plan
2907 .map(|p| {
2908 p.per_site_m == 1
2909 && p.per_site_n == 1
2910 && p.block
2911 == BlockShape {
2912 osh: 1,
2913 osw: 1,
2914 ibw: 1,
2915 }
2916 })
2917 .unwrap_or(false);
2918 let (multicast_groups, row_groups, col_groups) = match &self.placement {
2919 Some(pl) => (
2920 pl.groups.len(),
2921 pl.groups
2922 .iter()
2923 .filter(|g| g.direction == Direction::Row)
2924 .count(),
2925 pl.groups
2926 .iter()
2927 .filter(|g| g.direction == Direction::Col)
2928 .count(),
2929 ),
2930 None => (0, 0, 0),
2931 };
2932 let engine_program_classes = match &self.schedule {
2933 Some(s) => s
2934 .programs
2935 .iter()
2936 .map(|p| self.engine_signature(p))
2937 .collect::<BTreeSet<_>>()
2938 .len(),
2939 None => 0,
2940 };
2941 PlanSummary {
2942 phase: self.phase(),
2943 sites_used: self.sites().len(),
2944 grid,
2945 per_site_block,
2946 block,
2947 num_blocks,
2948 streaming,
2949 multicast_groups,
2950 row_groups,
2951 col_groups,
2952 acc,
2953 dram_tiles_read: self.dram_tiles_read(),
2954 semaphores: self.semaphores.len(),
2955 fifos: self
2956 .buffers
2957 .iter()
2958 .filter(|b| matches!(b.kind, BufferKind::Fifo { .. }))
2959 .count(),
2960 engine_program_classes,
2961 engines: self
2962 .work
2963 .plan
2964 .map(|pl| pl.engines)
2965 .unwrap_or(EnginePacking::LoadsTogether),
2966 }
2967 }
2968
2969 pub fn validate(&self) -> Result<(), IrViolation> {
2971 self.validate_ids()?;
2972 self.validate_sections()?;
2973 self.validate_references()?;
2974 self.validate_ports()?;
2975 self.validate_carries()?;
2976 self.validate_tensors()?;
2977 self.validate_compute()?;
2978 self.validate_indirect()?;
2979 self.validate_plan()?;
2980 self.validate_phase_order()?;
2981 self.validate_placement()?;
2982 self.validate_memory()?;
2983 self.validate_trips()?;
2984 self.validate_schedule()?;
2985 self.validate_rendezvous()?;
2986 self.validate_requirements()?;
2987 Ok(())
2988 }
2989
2990 pub fn bound_counts(&self) -> Vec<usize> {
2992 self.trips.iter().map(Trip::hi).collect()
2993 }
2994
2995 pub fn trip(&self, id: TripId) -> &Trip {
2996 &self.trips[id.idx()]
2997 }
2998
2999 fn validate_trips(&self) -> Result<(), IrViolation> {
3006 for (i, t) in self.trips.iter().enumerate() {
3007 let bad = |reason| {
3008 Err(IrViolation::TripMalformed {
3009 trip: TripId::new(i),
3010 reason,
3011 })
3012 };
3013 let Some(td) = self.tensors.get(t.tensor.idx()) else {
3014 return bad("a trip read from a tensor the program does not declare");
3015 };
3016 if td.role != TensorRole::Input {
3017 return bad("a trip read from a tensor that is not a program input");
3018 }
3019 if t.elem >= td.elems.rows * td.elems.cols {
3020 return bad("a trip read from an element past its tensor");
3021 }
3022 if t.per == 0 {
3023 return bad("a trip counted per zero elements");
3024 }
3025 if t.admit.0 > t.admit.1 {
3026 return bad("a trip that admits no value");
3027 }
3028 if t.lo() == 0 {
3029 return bad("a trip whose least count is zero");
3030 }
3031 }
3032 let Some(sch) = &self.schedule else {
3033 return Ok(());
3034 };
3035 fn walk(
3036 p: &SpatialProgram,
3037 ops: &[TaggedOp],
3038 inside: bool,
3039 used: &mut [bool],
3040 ) -> Result<(), IrViolation> {
3041 for t in ops {
3042 let ChoreoOp::Loop { count, body, .. } = &t.op else {
3043 continue;
3044 };
3045 let Some(trip) = count.trip() else {
3046 walk(p, body, inside, used)?;
3047 continue;
3048 };
3049 let bad = |reason| Err(IrViolation::TripLoopMalformed { reason });
3050 let Some(declared) = p.trips.get(trip.idx()) else {
3051 return bad("a loop run to a trip the program does not declare");
3052 };
3053 if inside {
3054 return bad("a loop run to a trip inside another");
3055 }
3056 if count.bound() != declared.hi() {
3057 return bad("a loop bounded apart from the trip it runs to");
3058 }
3059 let (mut foreign, mut stores) = (false, false);
3060 walk_tagged_ops(body, 1, &mut |op, _| {
3061 foreign |= matches!(
3062 op,
3063 ChoreoOp::SemSet { .. }
3064 | ChoreoOp::SemInc { .. }
3065 | ChoreoOp::SemWait { .. }
3066 | ChoreoOp::McastSend { .. }
3067 | ChoreoOp::McastSetValid { .. }
3068 | ChoreoOp::Barrier { .. }
3069 );
3070 stores |= matches!(op, ChoreoOp::WriteTiles { .. });
3071 });
3072 if foreign {
3073 return bad("a semaphore or multicast inside a loop run to a trip");
3074 }
3075 if stores {
3076 return bad(
3077 "a store inside a loop run to a trip, covering tiles the data decides",
3078 );
3079 }
3080 used[trip.idx()] = true;
3081 walk(p, body, true, used)?;
3082 }
3083 Ok(())
3084 }
3085 let mut used = vec![false; self.trips.len()];
3086 for prog in &sch.programs {
3087 walk(self, &prog.ops, false, &mut used)?;
3088 }
3089 if let Some(i) = used.iter().position(|u| !u) {
3090 return Err(IrViolation::TripMalformed {
3091 trip: TripId::new(i),
3092 reason: "a trip no loop runs to",
3093 });
3094 }
3095 Ok(())
3096 }
3097
3098 fn validate_ids(&self) -> Result<(), IrViolation> {
3099 for (i, t) in self.tensors.iter().enumerate() {
3100 if t.id.idx() != i {
3101 return Err(IrViolation::IdNotDense {
3102 what: "tensor",
3103 index: i,
3104 });
3105 }
3106 }
3107 for (i, s) in self.stages.iter().enumerate() {
3108 if s.id.idx() != i {
3109 return Err(IrViolation::IdNotDense {
3110 what: "stage",
3111 index: i,
3112 });
3113 }
3114 }
3115 for (i, b) in self.buffers.iter().enumerate() {
3116 if b.id.idx() != i {
3117 return Err(IrViolation::IdNotDense {
3118 what: "buffer",
3119 index: i,
3120 });
3121 }
3122 }
3123 for (i, e) in self.edges.iter().enumerate() {
3124 if e.id.idx() != i {
3125 return Err(IrViolation::IdNotDense {
3126 what: "edge",
3127 index: i,
3128 });
3129 }
3130 }
3131 for (i, s) in self.semaphores.iter().enumerate() {
3132 if s.id.idx() != i {
3133 return Err(IrViolation::IdNotDense {
3134 what: "semaphore",
3135 index: i,
3136 });
3137 }
3138 }
3139 for (i, w) in self.work.items.iter().enumerate() {
3140 if w.id.idx() != i {
3141 return Err(IrViolation::IdNotDense {
3142 what: "work item",
3143 index: i,
3144 });
3145 }
3146 }
3147 if let Some(pl) = &self.placement {
3148 for (i, g) in pl.groups.iter().enumerate() {
3149 if g.id.idx() != i {
3150 return Err(IrViolation::IdNotDense {
3151 what: "group",
3152 index: i,
3153 });
3154 }
3155 }
3156 }
3157 Ok(())
3158 }
3159
3160 fn validate_sections(&self) -> Result<(), IrViolation> {
3163 if self.sections.is_empty() {
3164 return Ok(());
3165 }
3166 type Span = fn(&SectionSpan) -> (usize, usize);
3167 let ends: [(&'static str, usize, Span); 5] = [
3168 ("stages", self.stages.len(), |s| s.stages),
3169 ("buffers", self.buffers.len(), |s| s.buffers),
3170 ("edges", self.edges.len(), |s| s.edges),
3171 ("semaphores", self.semaphores.len(), |s| s.semaphores),
3172 ("items", self.work.items.len(), |s| s.items),
3173 ];
3174 for (what, total, get) in ends {
3175 let mut cursor = 0usize;
3176 for (k, sp) in self.sections.iter().enumerate() {
3177 let (a, b) = get(sp);
3178 if a != cursor || b < a {
3179 return Err(IrViolation::SectionSpanMalformed { section: k, what });
3180 }
3181 cursor = b;
3182 }
3183 if cursor != total {
3184 return Err(IrViolation::SectionSpanMalformed {
3185 section: self.sections.len() - 1,
3186 what,
3187 });
3188 }
3189 }
3190 for (k, sp) in self.sections.iter().enumerate() {
3194 for s in &self.stages[sp.stages.0..sp.stages.1] {
3195 for b in s.ins.iter().chain(s.outs.iter()) {
3196 if !sp.holds_buffer(*b) {
3197 let bk = self.section_of_buffer(*b).unwrap_or(usize::MAX);
3198 return Err(IrViolation::SectionBufferEscape {
3199 stage: s.id,
3200 buffer: *b,
3201 stage_section: k,
3202 buffer_section: bk,
3203 });
3204 }
3205 }
3206 }
3207 }
3208 Ok(())
3209 }
3210
3211 pub fn section_phases(&self) -> Vec<usize> {
3213 let mut phase = 0usize;
3214 self.sections
3215 .iter()
3216 .map(|sp| {
3217 let here = phase;
3218 if sp.barrier_after.is_some() {
3219 phase += 1;
3220 }
3221 here
3222 })
3223 .collect()
3224 }
3225
3226 fn validate_references(&self) -> Result<(), IrViolation> {
3227 let nt = self.tensors.len();
3228 let nb = self.buffers.len();
3229 let ns = self.stages.len();
3230 let ng = self.placement.as_ref().map(|p| p.groups.len()).unwrap_or(0);
3231 for s in &self.stages {
3232 for b in s.ins.iter().chain(s.outs.iter()) {
3233 if b.idx() >= nb {
3234 return Err(IrViolation::DanglingBuffer {
3235 stage: Some(s.id),
3236 buffer: *b,
3237 });
3238 }
3239 }
3240 match &s.kind {
3241 StageKind::Load {
3242 tensor, multicast, ..
3243 } => {
3244 if tensor.idx() >= nt {
3245 return Err(IrViolation::DanglingTensor {
3246 stage: s.id,
3247 tensor: *tensor,
3248 });
3249 }
3250 if let Some(g) = multicast {
3251 if g.idx() >= ng {
3252 return Err(IrViolation::DanglingGroup {
3253 stage: s.id,
3254 group: *g,
3255 });
3256 }
3257 }
3258 }
3259 StageKind::Store { tensor, .. } => {
3260 if tensor.idx() >= nt {
3261 return Err(IrViolation::DanglingTensor {
3262 stage: s.id,
3263 tensor: *tensor,
3264 });
3265 }
3266 }
3267 StageKind::Recv { group } => {
3268 if group.idx() >= ng {
3269 return Err(IrViolation::DanglingGroup {
3270 stage: s.id,
3271 group: *group,
3272 });
3273 }
3274 }
3275 StageKind::Compute(ComputeOp::MatmulTiles {
3276 interm: Some(i), ..
3277 }) => {
3278 if i.idx() >= nb {
3279 return Err(IrViolation::DanglingBuffer {
3280 stage: Some(s.id),
3281 buffer: *i,
3282 });
3283 }
3284 }
3285 StageKind::Compute(ComputeOp::Bias { bias }) => {
3286 if bias.idx() >= nb {
3287 return Err(IrViolation::DanglingBuffer {
3288 stage: Some(s.id),
3289 buffer: *bias,
3290 });
3291 }
3292 }
3293 _ => {}
3294 }
3295 if let Some(it) = s.item {
3296 if it.idx() >= self.work.items.len() {
3297 return Err(IrViolation::DanglingItem {
3298 stage: s.id,
3299 item: it,
3300 });
3301 }
3302 }
3303 }
3304 for e in &self.edges {
3305 if e.from.idx() >= ns {
3306 return Err(IrViolation::DanglingStage {
3307 edge: e.id,
3308 stage: e.from,
3309 });
3310 }
3311 for t in &e.to {
3312 if t.idx() >= ns {
3313 return Err(IrViolation::DanglingStage {
3314 edge: e.id,
3315 stage: *t,
3316 });
3317 }
3318 }
3319 if e.buffer.idx() >= nb {
3320 return Err(IrViolation::DanglingBuffer {
3321 stage: None,
3322 buffer: e.buffer,
3323 });
3324 }
3325 if let Handshake::Semaphore { ready, valid, .. } = &e.handshake {
3326 for s in [ready, valid] {
3327 if s.idx() >= self.semaphores.len() {
3328 return Err(IrViolation::DanglingSemaphore { sem: *s });
3329 }
3330 }
3331 }
3332 }
3333 for b in &self.buffers {
3334 if let Some(a) = b.alias_of {
3335 if a.idx() >= nb {
3336 return Err(IrViolation::DanglingBuffer {
3337 stage: None,
3338 buffer: a,
3339 });
3340 }
3341 }
3342 }
3343 Ok(())
3344 }
3345
3346 fn validate_ports(&self) -> Result<(), IrViolation> {
3347 let mut producers: BTreeMap<BufferId, usize> = BTreeMap::new();
3348 let mut consumers: BTreeMap<BufferId, usize> = BTreeMap::new();
3349 for e in &self.edges {
3350 let from = self.stage(e.from);
3351 if !e.handshake.is_local() {
3352 if !from.outs.contains(&e.buffer) {
3354 return Err(IrViolation::EdgeBufferNotInStagePorts {
3355 edge: e.id,
3356 stage: e.from,
3357 buffer: e.buffer,
3358 });
3359 }
3360 continue;
3361 }
3362 if !from.outs.contains(&e.buffer) {
3363 return Err(IrViolation::EdgeBufferNotInStagePorts {
3364 edge: e.id,
3365 stage: e.from,
3366 buffer: e.buffer,
3367 });
3368 }
3369 *producers.entry(e.buffer).or_default() += 1;
3370 for t in &e.to {
3371 if !self.stage(*t).ins.contains(&e.buffer) {
3372 return Err(IrViolation::EdgeBufferNotInStagePorts {
3373 edge: e.id,
3374 stage: *t,
3375 buffer: e.buffer,
3376 });
3377 }
3378 *consumers.entry(e.buffer).or_default() += 1;
3379 }
3380 }
3381 for b in &self.buffers {
3382 if let BufferKind::Fifo { .. } = b.kind {
3383 let p = producers.get(&b.id).copied().unwrap_or(0);
3384 let c = consumers.get(&b.id).copied().unwrap_or(0);
3385 if p != 1 {
3386 return Err(IrViolation::FifoProducerCount {
3387 buffer: b.id,
3388 count: p,
3389 });
3390 }
3391 if c != 1 {
3392 return Err(IrViolation::FifoConsumerCount {
3393 buffer: b.id,
3394 count: c,
3395 });
3396 }
3397 }
3398 }
3399 Ok(())
3400 }
3401
3402 fn validate_carries(&self) -> Result<(), IrViolation> {
3403 let carries = self.carries();
3404 if carries.is_empty() {
3405 return Ok(());
3406 }
3407 let levels = self.fold_levels();
3408 for (buffer, combine) in carries {
3409 let b = self.buffer(buffer);
3410 if !matches!(b.kind, BufferKind::Interm { depth_pages: 2, .. }) {
3411 return Err(IrViolation::CarryNotARing { buffer });
3412 }
3413 let updates = self
3414 .stages
3415 .iter()
3416 .filter(|s| s.ins.contains(&buffer) && s.outs.contains(&buffer))
3417 .count();
3418 let inits: Vec<&Stage> = self
3419 .stages
3420 .iter()
3421 .filter(|s| s.id != combine && s.outs.contains(&buffer))
3422 .collect();
3423 if updates > 1 || inits.len() > 1 {
3424 return Err(IrViolation::CarryUpdatedTwice { buffer });
3425 }
3426 let Some(init) = inits.first() else {
3427 return Err(IrViolation::CarryWithoutInit { buffer });
3428 };
3429 if levels[init.id.idx()] != FoldLevel::Prologue {
3430 return Err(IrViolation::CarryInitMovesWithTheBlocks { buffer });
3431 }
3432 }
3433 Ok(())
3434 }
3435
3436 fn validate_tensors(&self) -> Result<(), IrViolation> {
3437 let mut loaded = vec![false; self.tensors.len()];
3438 let mut stores: Vec<Vec<&Stage>> = vec![Vec::new(); self.tensors.len()];
3439 for s in &self.stages {
3440 match &s.kind {
3441 StageKind::Load { tensor, .. } => {
3442 if let Some(l) = loaded.get_mut(tensor.idx()) {
3443 *l = true;
3444 }
3445 }
3446 StageKind::Store { tensor, .. } => {
3447 if let Some(v) = stores.get_mut(tensor.idx()) {
3448 v.push(s);
3449 }
3450 }
3451 _ => {}
3452 }
3453 }
3454 for t in &self.tensors {
3455 let rows_t = t.elems.rows.div_ceil(self.tile.rows);
3456 let cols_t = t.elems.cols.div_ceil(self.tile.cols);
3457 if t.tiles != (TileGrid { rows_t, cols_t }) {
3458 return Err(IrViolation::TileGridMismatch { tensor: t.id });
3459 }
3460 let (loaded, stored) = (loaded[t.id.idx()], !stores[t.id.idx()].is_empty());
3476 match t.role {
3477 TensorRole::Activation if self.sections.is_empty() || !(loaded && stored) => {
3478 return Err(IrViolation::TensorRoleDeferred {
3479 tensor: t.id,
3480 role: t.role,
3481 until: "B-P7",
3482 })
3483 }
3484 TensorRole::State if !stored => {
3485 return Err(IrViolation::TensorRoleDeferred {
3486 tensor: t.id,
3487 role: t.role,
3488 until: "B-P6",
3489 })
3490 }
3491 _ => {}
3492 }
3493 if t.space == TensorSpace::Sram
3497 && !matches!(t.role, TensorRole::Activation | TensorRole::State)
3498 {
3499 return Err(IrViolation::TensorSpaceDeferred {
3500 tensor: t.id,
3501 space: t.space,
3502 until: "B-P6",
3503 });
3504 }
3505 }
3506 for s in &self.stages {
3507 match &s.kind {
3508 StageKind::Load { tensor, .. } => {
3509 if !self.tensor(*tensor).role.is_read() {
3510 return Err(IrViolation::LoadReadsOutput {
3511 stage: s.id,
3512 tensor: *tensor,
3513 });
3514 }
3515 }
3516 StageKind::Store { tensor, .. } => {
3517 if !self.tensor(*tensor).role.is_written() {
3518 return Err(IrViolation::StoreWritesInput {
3519 stage: s.id,
3520 tensor: *tensor,
3521 });
3522 }
3523 }
3524 _ => {}
3525 }
3526 }
3527 let site_items = std::cell::OnceCell::new();
3538 for t in self.tensors.iter().filter(|t| t.role.is_written()) {
3539 let mut counts: BTreeMap<Option<usize>, Vec<usize>> = BTreeMap::new();
3540 let mut parts: BTreeMap<(Option<usize>, Option<Site>, Option<WorkItemId>), usize> = BTreeMap::new();
3542 for s in &stores[t.id.idx()] {
3543 *parts.entry((self.section_of_stage(s.id), s.site, s.item)).or_default() += 1;
3544 }
3545 for s in &stores[t.id.idx()] {
3546 let section = self.section_of_stage(s.id);
3547 let in_parts = parts[&(section, s.site, s.item)] > 1;
3548 let addressed = match &s.kind {
3549 StageKind::Store {
3550 pattern,
3551 sub_blocks,
3552 ..
3553 } if self.section_result(section).map(|r| r.id) != Some(t.id) || in_parts => {
3554 Some((pattern, *sub_blocks))
3555 }
3556 _ => None,
3557 };
3558 let count = counts
3559 .entry(if t.role == TensorRole::State {
3560 section
3561 } else {
3562 None
3563 })
3564 .or_insert_with(|| vec![0usize; t.tiles.count()]);
3565 let items: Vec<&WorkItem> = match (s.item, s.site) {
3566 (Some(it), _) => vec![&self.work.items[it.idx()]],
3567 (None, Some(site)) if self.placement.is_some() => site_items
3568 .get_or_init(|| self.site_items())
3569 .at(section, site)
3570 .iter()
3571 .map(|i| &self.work.items[i.idx()])
3572 .collect(),
3573 _ => match section.and_then(|k| self.sections.get(k)) {
3574 Some(sp) => self.work.items[sp.items.0..sp.items.1].iter().collect(),
3575 None => self.work.items.iter().collect(),
3576 },
3577 };
3578 for it in items {
3579 let tiles = addressed
3580 .and_then(|(pattern, sub_blocks)| pattern.tiles_at(it.coord, sub_blocks))
3581 .unwrap_or_else(|| it.out.tiles(t.tiles.cols_t));
3582 for tile in tiles {
3583 if tile >= count.len() {
3584 return Err(IrViolation::OutputOverlap { tensor: t.id, tile });
3585 }
3586 count[tile] += 1;
3587 }
3588 }
3589 }
3590 let uncovered = vec![0usize; t.tiles.count()];
3591 for count in counts
3592 .values()
3593 .chain(counts.is_empty().then_some(&uncovered))
3594 {
3595 for (tile, c) in count.iter().enumerate() {
3596 if *c == 0 && t.role.is_fully_covered() {
3597 return Err(IrViolation::OutputUncovered { tensor: t.id, tile });
3598 }
3599 if *c > 1 {
3600 return Err(IrViolation::OutputOverlap { tensor: t.id, tile });
3601 }
3602 }
3603 }
3604 }
3605 Ok(())
3606 }
3607
3608 fn validate_compute(&self) -> Result<(), IrViolation> {
3609 for s in &self.stages {
3610 if let StageKind::Compute(ComputeOp::MatmulTiles {
3611 block,
3612 operand,
3613 acc,
3614 interm,
3615 ..
3616 }) = &s.kind
3617 {
3618 for b in s.ins.iter().take(2) {
3619 let found = self.buffer(*b).dtype;
3620 if found != *operand {
3621 return Err(IrViolation::OperandDtypeMismatch {
3622 stage: s.id,
3623 buffer: *b,
3624 expected: *operand,
3625 found,
3626 });
3627 }
3628 }
3629 if let Some(i) = interm {
3630 let found = self.buffer(*i).dtype;
3631 if self.numerics == NumericsPolicy::Exact
3632 && precision_bytes(found) < precision_bytes(acc.interm_precision())
3633 {
3634 return Err(IrViolation::IntermNarrowerThanAcc {
3635 stage: s.id,
3636 interm: *i,
3637 acc: *acc,
3638 found,
3639 });
3640 }
3641 }
3642 if block.acc_tiles() > self.acc_tiles_used.max(1) {
3643 return Err(IrViolation::AccTilesExceeded {
3644 stage: s.id,
3645 needed: block.acc_tiles(),
3646 available: self.acc_tiles_used,
3647 });
3648 }
3649 }
3650 if let StageKind::Compute(ComputeOp::Generic {
3651 iters,
3652 ins: maps,
3653 out,
3654 body,
3655 }) = &s.kind
3656 {
3657 let operands = s.ins.len();
3658 if maps.len() != operands {
3659 return Err(IrViolation::GenericOperandCount {
3660 stage: s.id,
3661 operands,
3662 maps: maps.len(),
3663 });
3664 }
3665 if let Some(i) = body.max_operand().filter(|i| *i >= operands) {
3666 return Err(IrViolation::GenericOperandOutOfRange {
3667 stage: s.id,
3668 operand: i,
3669 operands,
3670 });
3671 }
3672 let malformed = |reason| {
3676 Err(IrViolation::ContractionMalformed {
3677 stage: s.id,
3678 reason,
3679 })
3680 };
3681 match iters.get(2) {
3682 Some(_) if iters.len() > 3 => return malformed("a fourth iteration axis — a Generic has rows, columns and one contraction"),
3683 Some(k) if k.var != crate::spatial::passes::LV_GEN_K => return malformed("a contraction axis over a variable other than LV_GEN_K"),
3684 Some(k) if !matches!(k.kind, IterKind::Reduce { pool: ReducePool::Sum }) => return malformed("a contraction axis that is not a sum fold"),
3685 Some(_) if out.varies_along(crate::spatial::passes::LV_GEN_K) => return malformed("an output map that steps along the contraction axis it folds"),
3686 None if maps.iter().chain(std::iter::once(out)).any(|m| m.varies_along(crate::spatial::passes::LV_GEN_K)) => {
3687 return malformed("a map stepping along a contraction axis the iteration space does not have")
3688 }
3689 _ => {}
3690 }
3691 for m in maps.iter().chain(std::iter::once(out)) {
3692 if let Some(ind) = m.indirect {
3693 if ind.operand >= operands {
3694 return Err(IrViolation::IndirectOperandOutOfRange {
3695 stage: s.id,
3696 operand: ind.operand,
3697 operands,
3698 });
3699 }
3700 if ind.modulus == 0 {
3701 return Err(IrViolation::IndirectMalformed {
3702 stage: s.id,
3703 reason: "an indirect term's modulus must be nonzero",
3704 });
3705 }
3706 }
3707 }
3708 if out.indirect.is_some() {
3711 let scattered = s.outs.iter().any(|b| {
3712 self.stages.iter().any(|st| matches!(&st.kind, StageKind::Store { pattern, .. } if st.ins.contains(b) && pattern.origin.indirect.is_some()))
3713 });
3714 if !scattered {
3715 return Err(IrViolation::IndirectOutWithoutScatter { stage: s.id });
3716 }
3717 }
3718 if iters.iter().any(|a| {
3720 matches!(
3721 a.kind,
3722 IterKind::Reduce {
3723 pool: ReducePool::ArgMax
3724 }
3725 )
3726 }) {
3727 for b in &s.outs {
3728 let found = self.buffer(*b).dtype;
3729 if precision_int_bits(found).is_none() {
3730 return Err(IrViolation::ArgMaxOutputNotInteger { stage: s.id, found });
3731 }
3732 }
3733 }
3734 }
3735 }
3736 Ok(())
3737 }
3738
3739 fn validate_indirect(&self) -> Result<(), IrViolation> {
3742 for s in &self.stages {
3743 let pattern = match &s.kind {
3744 StageKind::Load { pattern, .. } | StageKind::Store { pattern, .. } => pattern,
3745 _ => continue,
3746 };
3747 let Some(ind) = pattern.origin.indirect else {
3748 continue;
3749 };
3750 if ind.index.idx() >= self.buffers.len() {
3751 return Err(IrViolation::DanglingBuffer {
3752 stage: Some(s.id),
3753 buffer: ind.index,
3754 });
3755 }
3756 if !s.ins.contains(&ind.index) {
3757 return Err(IrViolation::IndirectIndexNotAnOperand {
3758 stage: s.id,
3759 buffer: ind.index,
3760 });
3761 }
3762 if ind.tile_stride == 0 || ind.elems_per_tile == 0 {
3763 return Err(IrViolation::IndirectMalformed {
3764 stage: s.id,
3765 reason: "an indirect tile's stride and elements-per-tile must be nonzero",
3766 });
3767 }
3768 if ind.bound.0 > ind.bound.1 {
3769 return Err(IrViolation::IndirectMalformed {
3770 stage: s.id,
3771 reason: "an indirect tile's bound is inverted",
3772 });
3773 }
3774 }
3775 Ok(())
3776 }
3777
3778 pub fn has_indirect_access(&self) -> bool {
3782 fn op_has(op: &ComputeOp) -> bool {
3783 matches!(op, ComputeOp::Generic { ins, out, .. } if ins.iter().chain(std::iter::once(out)).any(|m| m.indirect.is_some()))
3784 }
3785 let in_stages = self.stages.iter().any(|s| match &s.kind {
3786 StageKind::Load { pattern, .. } | StageKind::Store { pattern, .. } => {
3787 pattern.origin.indirect.is_some()
3788 }
3789 StageKind::Compute(op) => op_has(op),
3790 _ => false,
3791 });
3792 if in_stages {
3793 return true;
3794 }
3795 let mut found = false;
3796 if let Some(sch) = &self.schedule {
3797 for p in &sch.programs {
3798 walk_tagged(&p.ops, 1, &mut |t, _| {
3799 found |= match &t.op {
3800 ChoreoOp::ReadTiles { at, .. } | ChoreoOp::WriteTiles { at, .. } => {
3801 at.indirect.is_some()
3802 }
3803 ChoreoOp::Compute { op, .. } => op_has(op),
3804 _ => false,
3805 };
3806 });
3807 }
3808 }
3809 found
3810 }
3811
3812 pub fn generic_is_plain_pointwise(
3818 &self,
3819 iters: &[IterSpec],
3820 maps: &[ElemAccess],
3821 out: &ElemAccess,
3822 ins: &[BufferId],
3823 ) -> bool {
3824 let [row_ax, col_ax] = iters else {
3825 return false;
3826 };
3827 if !matches!(row_ax.kind, IterKind::Parallel) || !matches!(col_ax.kind, IterKind::Parallel)
3828 {
3829 return false;
3830 }
3831 if maps.len() != ins.len()
3832 || out.indirect.is_some()
3833 || !out.varies_along(row_ax.var)
3834 || !out.varies_along(col_ax.var)
3835 {
3836 return false;
3837 }
3838 maps.iter().zip(ins).all(|(m, b)| {
3839 let width = self.buffer(*b).kind.page_tiles() * self.tile.cols;
3840 m.broadcast_class(row_ax.var, col_ax.var, width) == Some((true, true))
3841 })
3842 }
3843
3844 fn validate_plan(&self) -> Result<(), IrViolation> {
3845 let Some(p) = self.work.plan else {
3846 return Ok(());
3847 };
3848 if p.per_site_m == 0 || p.block.osh == 0 || p.per_site_m % p.block.osh != 0 {
3849 return Err(IrViolation::BlockDivisibility {
3850 what: "per_site_m % osh",
3851 value: p.per_site_m,
3852 divisor: p.block.osh,
3853 });
3854 }
3855 if p.per_site_n == 0 || p.block.osw == 0 || p.per_site_n % p.block.osw != 0 {
3856 return Err(IrViolation::BlockDivisibility {
3857 what: "per_site_n % osw",
3858 value: p.per_site_n,
3859 divisor: p.block.osw,
3860 });
3861 }
3862 if p.block.ibw == 0 || self.work.k_tiles % p.block.ibw != 0 {
3863 return Err(IrViolation::BlockDivisibility {
3864 what: "k_tiles % ibw",
3865 value: self.work.k_tiles,
3866 divisor: p.block.ibw,
3867 });
3868 }
3869 if p.inner.num_blocks != self.work.k_tiles / p.block.ibw
3870 || p.inner.block_k_tiles != p.block.ibw
3871 {
3872 return Err(IrViolation::BlockDivisibility {
3873 what: "num_blocks",
3874 value: p.inner.num_blocks,
3875 divisor: p.block.ibw,
3876 });
3877 }
3878 let last = self.sections.len().checked_sub(1);
3886 let items = match last.and_then(|k| self.sections.get(k)) {
3887 Some(sp) => sp.items.1 - sp.items.0,
3888 None => self.work.items.len(),
3889 };
3890 if let Some(out) = self.section_result(last) {
3891 let expected = if out.role.is_fully_covered() {
3892 (
3893 out.tiles.rows_t.div_ceil(p.per_site_m),
3894 out.tiles.cols_t.div_ceil(p.per_site_n),
3895 )
3896 } else {
3897 (1, items)
3898 };
3899 if p.grid != expected {
3900 return Err(IrViolation::GridMismatch {
3901 expected,
3902 found: p.grid,
3903 });
3904 }
3905 }
3906 Ok(())
3907 }
3908
3909 fn validate_phase_order(&self) -> Result<(), IrViolation> {
3910 if self.placement.is_some() && self.work.plan.is_none() {
3911 return Err(IrViolation::PhaseOrder {
3912 have: "placement",
3913 missing: "plan",
3914 });
3915 }
3916 if self.memory.is_some() && self.placement.is_none() {
3917 return Err(IrViolation::PhaseOrder {
3918 have: "memory",
3919 missing: "placement",
3920 });
3921 }
3922 if self.schedule.is_some() && self.memory.is_none() {
3923 return Err(IrViolation::PhaseOrder {
3924 have: "schedule",
3925 missing: "memory",
3926 });
3927 }
3928 if self.certificate.is_some() && self.schedule.is_none() {
3929 return Err(IrViolation::PhaseOrder {
3930 have: "certificate",
3931 missing: "schedule",
3932 });
3933 }
3934 Ok(())
3935 }
3936
3937 fn validate_placement(&self) -> Result<(), IrViolation> {
3938 let Some(pl) = &self.placement else {
3939 return Ok(());
3940 };
3941 if pl.site_of.len() != self.stages.len() || pl.engine_of.len() != self.stages.len() {
3942 return Err(IrViolation::StageUnplaced {
3943 stage: StageId::new(pl.site_of.len().min(pl.engine_of.len())),
3944 });
3945 }
3946 for s in &self.stages {
3947 let Some(site) = s.site else {
3948 return Err(IrViolation::StageUnplaced { stage: s.id });
3949 };
3950 if pl.site_of[s.id.idx()] != site || pl.engine_of[s.id.idx()].site != site {
3951 return Err(IrViolation::EngineSiteMismatch { stage: s.id });
3952 }
3953 let role_ok = match (s.role, pl.engine_of[s.id.idx()].engine) {
3954 (SiteRole::Movement, Engine::Movement(_)) => true,
3955 (SiteRole::Compute, Engine::Compute) => true,
3956 _ => false,
3957 };
3958 if !role_ok || s.role != s.kind.role() {
3959 return Err(IrViolation::EngineRoleMismatch { stage: s.id });
3960 }
3961 }
3962 for e in &self.edges {
3963 let from_site = self.stage(e.from).site;
3964 for t in &e.to {
3965 let to_site = self.stage(*t).site;
3966 if e.handshake.is_local() && from_site != to_site {
3967 return Err(IrViolation::CrossSiteLocalEdge { edge: e.id });
3968 }
3969 if !e.handshake.is_local()
3970 && from_site == to_site
3971 && !matches!(e.handshake, Handshake::Barrier { .. })
3972 {
3973 return Err(IrViolation::LocalCrossSiteEdge { edge: e.id });
3974 }
3975 }
3976 }
3977 let mut members_of: BTreeMap<(GroupId, Site), Vec<&Stage>> = BTreeMap::new();
3978 if !pl.groups.is_empty() {
3979 for s in &self.stages {
3980 if let (Some(g), Some(site)) = (self.stage_group(s), s.site) {
3981 members_of.entry((g, site)).or_default().push(s);
3982 }
3983 }
3984 }
3985 let site_items = if pl.groups.is_empty() {
3986 None
3987 } else {
3988 Some(self.site_items())
3989 };
3990 for g in &pl.groups {
3991 if !g.rect.contains(g.sender) {
3992 return Err(IrViolation::GroupMalformed {
3993 group: g.id,
3994 reason: "sender outside rect",
3995 });
3996 }
3997 match g.direction {
3998 Direction::Row => {
3999 if g.rect.y0 != g.rect.y1 || g.sender.x != g.rect.x0 {
4000 return Err(IrViolation::GroupMalformed {
4001 group: g.id,
4002 reason: "row group must be one row with the sender at its left",
4003 });
4004 }
4005 }
4006 Direction::Col => {
4007 if g.rect.x0 != g.rect.x1 || g.sender.y != g.rect.y0 {
4008 return Err(IrViolation::GroupMalformed {
4009 group: g.id,
4010 reason: "column group must be one column with the sender at its top",
4011 });
4012 }
4013 }
4014 }
4015 if g.rect.len() < 2 {
4016 return Err(IrViolation::GroupMalformed {
4017 group: g.id,
4018 reason: "a multicast group needs at least one receiver",
4019 });
4020 }
4021 let mut member_kind: Option<(BufferKind, Precision)> = None;
4022 let mut member_offset: Option<u32> = None;
4023 for site in g.rect.sites() {
4024 let members: &[&Stage] = members_of
4025 .get(&(g.id, site))
4026 .map(Vec::as_slice)
4027 .unwrap_or(&[]);
4028 if members.len() != 1 {
4029 return Err(IrViolation::GroupMemberMismatch {
4030 group: g.id,
4031 site,
4032 reason: "exactly one member stage per site",
4033 });
4034 }
4035 let m = members[0];
4036 let is_sender = site == g.sender;
4037 match &m.kind {
4038 StageKind::Load {
4039 multicast: Some(_), ..
4040 } if is_sender => {}
4041 StageKind::Recv { .. } if !is_sender => {}
4042 _ => {
4043 return Err(IrViolation::GroupMemberMismatch {
4044 group: g.id,
4045 site,
4046 reason: "sender must be a multicasting Load, receivers must be Recv",
4047 })
4048 }
4049 }
4050 let Some(out) = m.outs.first() else {
4051 return Err(IrViolation::GroupMemberMismatch {
4052 group: g.id,
4053 site,
4054 reason: "member has no output buffer",
4055 });
4056 };
4057 let b = self.buffer(*out);
4058 match member_kind {
4059 None => member_kind = Some((b.kind, b.dtype)),
4060 Some(k) => {
4061 if k != (b.kind, b.dtype) {
4062 return Err(IrViolation::GroupMemberMismatch {
4063 group: g.id,
4064 site,
4065 reason: "member buffers differ in kind or dtype",
4066 });
4067 }
4068 }
4069 }
4070 if let Some(mm) = &self.memory {
4071 let at = mm
4072 .buffer_at
4073 .get(out.idx())
4074 .ok_or(IrViolation::MemoryUnmapped { buffer: *out })?;
4075 match member_offset {
4076 None => member_offset = Some(at.offset_bytes),
4077 Some(o) => {
4078 if o != at.offset_bytes {
4079 return Err(IrViolation::GroupMemberMismatch {
4080 group: g.id,
4081 site,
4082 reason:
4083 "multicast lands at one L1 address: member offsets differ",
4084 });
4085 }
4086 }
4087 }
4088 }
4089 let section = self.section_of_stage(m.id);
4092 let at = |x: Site| site_items.as_ref().map_or(0, |ix| ix.at(section, x).len());
4093 let (n_items, n_sender) = (at(site), at(g.sender));
4094 if n_items != n_sender {
4095 return Err(IrViolation::ItemCountMismatch { group: g.id });
4096 }
4097 }
4098 }
4099 Ok(())
4100 }
4101
4102 fn validate_memory(&self) -> Result<(), IrViolation> {
4103 let Some(mm) = &self.memory else {
4104 return Ok(());
4105 };
4106 if mm.buffer_at.len() != self.buffers.len() {
4107 return Err(IrViolation::MemoryUnmapped {
4108 buffer: BufferId::new(mm.buffer_at.len()),
4109 });
4110 }
4111 let phases = self.section_phases();
4118 let mut by_site: BTreeMap<(Option<usize>, Site, MemSpace), Vec<(u32, u32, BufferId)>> =
4119 BTreeMap::new();
4120 for b in &self.buffers {
4121 if b.alias_of.is_some() || b.space == MemSpace::Acc {
4122 continue;
4123 }
4124 let at = mm.buffer_at[b.id.idx()];
4125 let phase = self.section_of_buffer(b.id).map(|k| phases[k]);
4126 by_site
4127 .entry((phase, at.site, at.space))
4128 .or_default()
4129 .push((at.offset_bytes, at.offset_bytes + at.bytes, b.id));
4130 }
4131 for ((_, site, _), mut v) in by_site {
4132 v.sort();
4133 for w in v.windows(2) {
4134 if w[1].0 < w[0].1 {
4135 return Err(IrViolation::MemoryOverlap {
4136 site,
4137 a: w[0].2,
4138 b: w[1].2,
4139 });
4140 }
4141 }
4142 }
4143 Ok(())
4144 }
4145
4146 fn validate_schedule(&self) -> Result<(), IrViolation> {
4147 let Some(sch) = &self.schedule else {
4148 return Ok(());
4149 };
4150 let mut seen: BTreeSet<(Site, Engine)> = BTreeSet::new();
4151 for p in &sch.programs {
4152 if !seen.insert((p.site, p.engine)) {
4153 return Err(IrViolation::DuplicateEngineProgram {
4154 site: p.site,
4155 engine: p.engine,
4156 });
4157 }
4158 }
4159 let writes_tiles: BTreeSet<StageId> = {
4182 let mut v = BTreeSet::new();
4183 fn walk(ops: &[TaggedOp], out: &mut BTreeSet<StageId>) {
4184 for t in ops {
4185 if matches!(t.op, ChoreoOp::WriteTiles { .. }) {
4186 out.insert(t.stage);
4187 }
4188 if let ChoreoOp::Loop { body, .. } = &t.op {
4189 walk(body, out);
4190 }
4191 }
4192 }
4193 for p in &sch.programs {
4194 walk(&p.ops, &mut v);
4195 }
4196 v
4197 };
4198 for st in &self.stages {
4199 if matches!(st.kind, StageKind::Store { .. }) && !writes_tiles.contains(&st.id) {
4200 return Err(IrViolation::StageNotScheduled {
4201 stage: st.id,
4202 op: "WriteTiles",
4203 });
4204 }
4205 }
4206 let mut reserved: BTreeMap<BufferId, u64> = BTreeMap::new();
4207 let mut pushed: BTreeMap<BufferId, u64> = BTreeMap::new();
4208 let mut popped: BTreeMap<BufferId, u64> = BTreeMap::new();
4209 let mut incs: BTreeMap<SemId, u64> = BTreeMap::new();
4210 let mut waits: BTreeMap<SemId, u64> = BTreeMap::new();
4211 let mut sets_valid: BTreeMap<(SemId, Site), u64> = BTreeMap::new();
4212 let mut resets: BTreeMap<SemId, u64> = BTreeMap::new();
4213 let mut wait_ops: BTreeMap<SemId, u64> = BTreeMap::new();
4217 let mut invalidations: BTreeMap<(SemId, Site), u64> = BTreeMap::new();
4218 let mut valid_waits: BTreeMap<(SemId, Site), u64> = BTreeMap::new();
4219 let mut posts_at: BTreeMap<(SemId, Site), u64> = BTreeMap::new();
4222 let mut barrier_waits: BTreeMap<SemId, Vec<(Site, SemCond, u64)>> = BTreeMap::new();
4223 for p in &sch.programs {
4224 let mut acq = 0u64;
4225 let mut rel = 0u64;
4226 let mut acq_over: Option<usize> = None;
4227 let mut computed_since_acquire = true;
4228 let mut pack_without_compute = false;
4229 walk_tagged_ops(&p.ops, 1, &mut |op, m| match op {
4230 ChoreoOp::Compute { .. } => computed_since_acquire = true,
4231 ChoreoOp::PackAcc { .. } if !computed_since_acquire => pack_without_compute = true,
4232 ChoreoOp::Reserve { buf, pages } => {
4233 *reserved.entry(*buf).or_default() += m * *pages as u64
4234 }
4235 ChoreoOp::Push { buf, pages } => {
4236 *pushed.entry(*buf).or_default() += m * *pages as u64
4237 }
4238 ChoreoOp::Pop { buf, pages } => {
4239 *popped.entry(*buf).or_default() += m * *pages as u64
4240 }
4241 ChoreoOp::SemInc { sem, at, by } => {
4242 *incs.entry(*sem).or_default() += m * *by as u64;
4243 if matches!(self.semaphores[sem.idx()].role, SemRole::Barrier { .. }) {
4244 *posts_at.entry((*sem, *at)).or_default() += m * *by as u64;
4245 }
4246 }
4247 ChoreoOp::SemSet {
4255 sem,
4256 value: SemValue::Count(0),
4257 } => *resets.entry(*sem).or_default() += m,
4258 ChoreoOp::SemSet {
4261 sem,
4262 value: SemValue::Invalid,
4263 } => *invalidations.entry((*sem, p.site)).or_default() += m,
4264 ChoreoOp::SemWait { sem, until } => match self.semaphores[sem.idx()].role {
4265 SemRole::ValidFlag => *valid_waits.entry((*sem, p.site)).or_default() += m,
4266 role => {
4267 let n = match until {
4268 SemCond::Equals(n) | SemCond::AtLeast(n) => *n as u64,
4269 };
4270 *waits.entry(*sem).or_default() += m * n;
4271 *wait_ops.entry(*sem).or_default() += m;
4272 if matches!(role, SemRole::Barrier { .. }) {
4273 barrier_waits
4274 .entry(*sem)
4275 .or_default()
4276 .push((p.site, *until, m));
4277 }
4278 }
4279 },
4280 ChoreoOp::McastSetValid { sem, .. } => {
4281 *sets_valid.entry((*sem, p.site)).or_default() += m
4282 }
4283 ChoreoOp::AcquireAcc { tiles } => {
4284 computed_since_acquire = false;
4285 acq += m;
4286 if *tiles > self.acc_tiles_used.max(1) {
4287 acq_over = Some(*tiles);
4288 }
4289 }
4290 ChoreoOp::ReleaseAcc => rel += m,
4291 _ => {}
4292 });
4293 if acq != rel {
4294 return Err(IrViolation::AccImbalance {
4295 site: p.site,
4296 acquires: acq,
4297 releases: rel,
4298 });
4299 }
4300 if pack_without_compute {
4301 return Err(IrViolation::AccumulatorNeverComputed {
4302 site: p.site,
4303 engine: p.engine,
4304 });
4305 }
4306 if let Some(t) = acq_over {
4307 return Err(IrViolation::AccTilesExceeded {
4308 stage: p.stages.first().copied().unwrap_or(StageId(0)),
4309 needed: t,
4310 available: self.acc_tiles_used,
4311 });
4312 }
4313 }
4314 for b in &self.buffers {
4315 if b.space == MemSpace::Acc {
4316 continue;
4317 }
4318 let r = reserved.get(&b.id).copied().unwrap_or(0);
4319 let pu = pushed.get(&b.id).copied().unwrap_or(0);
4320 let po = popped.get(&b.id).copied().unwrap_or(0);
4321 if r != pu || pu != po {
4322 return Err(IrViolation::CountingImbalance {
4323 buffer: b.id,
4324 reserved: r,
4325 pushed: pu,
4326 popped: po,
4327 });
4328 }
4329 }
4330 if !self.trips.is_empty() {
4334 fn per_round(
4335 ops: &[TaggedOp],
4336 mult: u64,
4337 trip: Option<TripId>,
4338 f: &mut dyn FnMut(&ChoreoOp, u64, Option<TripId>),
4339 ) {
4340 for t in ops {
4341 f(&t.op, mult, trip);
4342 if let ChoreoOp::Loop { count, body, .. } = &t.op {
4343 match count.trip() {
4344 Some(id) => per_round(body, mult, Some(id), f),
4345 None => per_round(body, mult * count.bound() as u64, trip, f),
4346 }
4347 }
4348 }
4349 }
4350 let mut moved: BTreeMap<(BufferId, TripId), [u64; 3]> = BTreeMap::new();
4351 for p in &sch.programs {
4352 let mut acc: BTreeMap<TripId, [u64; 2]> = BTreeMap::new();
4353 per_round(&p.ops, 1, None, &mut |op, m, trip| {
4354 let Some(trip) = trip else { return };
4355 match op {
4356 ChoreoOp::Reserve { buf, pages } => {
4357 moved.entry((*buf, trip)).or_default()[0] += m * *pages as u64
4358 }
4359 ChoreoOp::Push { buf, pages } => {
4360 moved.entry((*buf, trip)).or_default()[1] += m * *pages as u64
4361 }
4362 ChoreoOp::Pop { buf, pages } => {
4363 moved.entry((*buf, trip)).or_default()[2] += m * *pages as u64
4364 }
4365 ChoreoOp::AcquireAcc { .. } => acc.entry(trip).or_default()[0] += m,
4366 ChoreoOp::ReleaseAcc => acc.entry(trip).or_default()[1] += m,
4367 _ => {}
4368 }
4369 });
4370 if let Some([a, r]) = acc.values().find(|[a, r]| a != r) {
4371 return Err(IrViolation::AccImbalance {
4372 site: p.site,
4373 acquires: *a,
4374 releases: *r,
4375 });
4376 }
4377 }
4378 for ((buffer, _), [r, pu, po]) in &moved {
4379 if self.buffer(*buffer).space != MemSpace::Acc && (r != pu || pu != po) {
4380 return Err(IrViolation::CountingImbalance {
4381 buffer: *buffer,
4382 reserved: *r,
4383 pushed: *pu,
4384 popped: *po,
4385 });
4386 }
4387 }
4388 }
4389 for s in &self.semaphores {
4390 match s.role {
4391 SemRole::Barrier { n, phase } => {
4399 let i = incs.get(&s.id).copied().unwrap_or(0);
4400 let ops = wait_ops.get(&s.id).copied().unwrap_or(0);
4401 let unbalanced = IrViolation::SectionBarrierUnbalanced {
4402 sem: s.id,
4403 participants: n,
4404 sites: s.sites.len(),
4405 incs: i,
4406 wait_ops: ops,
4407 };
4408 let sites: BTreeSet<Site> = s.sites.iter().copied().collect();
4409 let (post_each, want_wait_ops, wait_for) = match phase {
4410 BarrierPhase::Gather => (n as u64, 1u64, n),
4411 BarrierPhase::Release => (1u64, n as u64, 1),
4412 };
4413 let shape_ok = sites.len() == s.sites.len()
4414 && !sites.is_empty()
4415 && (phase == BarrierPhase::Release || sites.len() == 1);
4416 let posts_ok = i == post_each * sites.len() as u64
4417 && sites
4418 .iter()
4419 .all(|x| posts_at.get(&(s.id, *x)).copied().unwrap_or(0) == post_each);
4420 let waiting = barrier_waits.get(&s.id).map(Vec::as_slice).unwrap_or(&[]);
4421 let waits_ok = ops == want_wait_ops
4422 && waiting.iter().all(|(site, until, _)| {
4423 *until == SemCond::AtLeast(wait_for) && sites.contains(site)
4424 });
4425 if !(shape_ok && posts_ok && waits_ok) {
4426 return Err(unbalanced);
4427 }
4428 }
4429 SemRole::ReadyCounter { .. } => {
4430 let i = incs.get(&s.id).copied().unwrap_or(0);
4431 let w = waits.get(&s.id).copied().unwrap_or(0);
4432 if i != w {
4433 return Err(IrViolation::SemaphoreImbalance {
4434 sem: s.id,
4435 incs: i,
4436 waits: w,
4437 });
4438 }
4439 if matches!(s.role, SemRole::ReadyCounter { .. }) {
4451 let ops = wait_ops.get(&s.id).copied().unwrap_or(0);
4452 let r = resets.get(&s.id).copied().unwrap_or(0);
4453 if ops != r {
4454 return Err(IrViolation::SemaphoreNotReset {
4455 sem: s.id,
4456 waits: ops,
4457 resets: r,
4458 });
4459 }
4460 }
4461 }
4462 SemRole::ValidFlag => {
4463 for site in &s.sites {
4466 let w = valid_waits.get(&(s.id, *site)).copied().unwrap_or(0);
4467 if w == 0 {
4468 continue;
4469 }
4470 let setter = self.placement.as_ref().and_then(|pl| {
4471 s.operand.and_then(|op| {
4472 pl.groups
4473 .iter()
4474 .find(|g| g.operand == op && g.rect.contains(*site))
4475 .map(|g| g.sender)
4476 })
4477 });
4478 let sets = match setter {
4479 Some(sender) => sets_valid.get(&(s.id, sender)).copied().unwrap_or(0),
4480 None => sets_valid
4481 .iter()
4482 .filter(|((id, _), _)| *id == s.id)
4483 .map(|(_, v)| *v)
4484 .sum(),
4485 };
4486 if w != sets {
4487 return Err(IrViolation::SemaphoreImbalance {
4488 sem: s.id,
4489 incs: sets,
4490 waits: w,
4491 });
4492 }
4493 let cleared = invalidations.get(&(s.id, *site)).copied().unwrap_or(0);
4499 if cleared != w {
4500 return Err(IrViolation::SemaphoreNotReset {
4501 sem: s.id,
4502 waits: w,
4503 resets: cleared,
4504 });
4505 }
4506 }
4507 }
4508 }
4509 }
4510 Ok(())
4511 }
4512
4513 fn validate_rendezvous(&self) -> Result<(), IrViolation> {
4525 let Some(sch) = &self.schedule else {
4526 return Ok(());
4527 };
4528 let mut closing: BTreeMap<SemId, (usize, BarrierPhase)> = BTreeMap::new();
4529 for (k, sp) in self.sections.iter().enumerate() {
4530 let Some(r) = sp.barrier_after else { continue };
4531 let role = |id: SemId| self.semaphores.get(id.idx()).map(|s| s.role);
4532 let owned = |id: SemId| id.idx() >= sp.semaphores.0 && id.idx() < sp.semaphores.1;
4533 let well_formed = matches!(
4534 role(r.gather),
4535 Some(SemRole::Barrier {
4536 phase: BarrierPhase::Gather,
4537 ..
4538 })
4539 ) && matches!(
4540 role(r.release),
4541 Some(SemRole::Barrier {
4542 phase: BarrierPhase::Release,
4543 ..
4544 })
4545 ) && owned(r.gather)
4546 && owned(r.release);
4547 if !well_formed {
4548 return Err(IrViolation::SectionSpanMalformed {
4549 section: k,
4550 what:
4551 "a rendezvous whose semaphores are not its own section's gather and release",
4552 });
4553 }
4554 closing.insert(r.gather, (k, BarrierPhase::Gather));
4555 closing.insert(r.release, (k, BarrierPhase::Release));
4556 }
4557 if let Some(stray) = self
4558 .semaphores
4559 .iter()
4560 .find(|s| matches!(s.role, SemRole::Barrier { .. }) && !closing.contains_key(&s.id))
4561 {
4562 let section = self
4563 .sections
4564 .iter()
4565 .position(|sp| {
4566 stray.id.idx() >= sp.semaphores.0 && stray.id.idx() < sp.semaphores.1
4567 })
4568 .unwrap_or(0);
4569 return Err(IrViolation::SectionSpanMalformed {
4570 section,
4571 what: "a barrier semaphore that closes no section",
4572 });
4573 }
4574 if closing.is_empty() {
4575 return Ok(());
4576 }
4577 let n = self.sections.len();
4578 let closes: Vec<bool> = self
4579 .sections
4580 .iter()
4581 .map(|sp| sp.barrier_after.is_some())
4582 .collect();
4583 let mut section_of = vec![usize::MAX; self.stages.len()];
4584 for (k, sp) in self.sections.iter().enumerate() {
4585 for slot in section_of.iter_mut().take(sp.stages.1).skip(sp.stages.0) {
4586 *slot = k;
4587 }
4588 }
4589 let rendezvous_sem = |op: &ChoreoOp| match op {
4590 ChoreoOp::SemInc { sem, .. } | ChoreoOp::SemWait { sem, .. } => {
4591 closing.get(sem).copied()
4592 }
4593 _ => None,
4594 };
4595 let root_of: BTreeMap<usize, (Site, Engine)> = sch
4596 .programs
4597 .iter()
4598 .flat_map(|p| {
4599 p.ops
4600 .iter()
4601 .filter_map(move |t| match (&t.op, rendezvous_sem(&t.op)) {
4602 (ChoreoOp::SemWait { .. }, Some((k, BarrierPhase::Gather))) => {
4603 Some((k, (p.site, p.engine)))
4604 }
4605 _ => None,
4606 })
4607 })
4608 .collect();
4609 for p in &sch.programs {
4610 let misplaced =
4611 |section: usize, what: &'static str| IrViolation::SectionBarrierMisplaced {
4612 section,
4613 site: p.site,
4614 engine: p.engine,
4615 what,
4616 };
4617 let mut nested: Option<usize> = None;
4618 for t in &p.ops {
4619 if let ChoreoOp::Loop { body, .. } = &t.op {
4620 walk_tagged_ops(body, 1, &mut |op, _| {
4621 if let Some((k, _)) = rendezvous_sem(op) {
4622 nested.get_or_insert(k);
4623 }
4624 });
4625 }
4626 }
4627 if let Some(k) = nested {
4628 return Err(misplaced(k, "a rendezvous op inside a loop body"));
4629 }
4630 let mut arrive: Vec<Option<usize>> = vec![None; n];
4632 let mut gather_wait: Vec<Option<usize>> = vec![None; n];
4633 let mut posts: Vec<Option<(usize, usize)>> = vec![None; n];
4634 let mut release_wait: Vec<Option<usize>> = vec![None; n];
4635 for (i, t) in p.ops.iter().enumerate() {
4636 let Some((k, phase)) = rendezvous_sem(&t.op) else {
4637 continue;
4638 };
4639 let slot = match (&t.op, phase) {
4640 (ChoreoOp::SemInc { .. }, BarrierPhase::Gather) => &mut arrive[k],
4641 (ChoreoOp::SemWait { .. }, BarrierPhase::Gather) => &mut gather_wait[k],
4642 (ChoreoOp::SemWait { .. }, BarrierPhase::Release) => &mut release_wait[k],
4643 _ => {
4644 let span = posts[k].get_or_insert((i, i));
4645 span.1 = i;
4646 continue;
4647 }
4648 };
4649 if slot.replace(i).is_some() {
4650 return Err(misplaced(k, "a program that passes one rendezvous twice"));
4651 }
4652 }
4653 let mut next_arrival: Vec<Option<usize>> = vec![None; n + 1];
4656 for k in (0..n).rev() {
4657 next_arrival[k] = if closes[k] {
4658 arrive[k]
4659 } else {
4660 next_arrival[k + 1]
4661 };
4662 }
4663 let mut previous_release: Option<usize> = None;
4664 let mut release_before: Vec<Option<usize>> = vec![None; n];
4665 for k in 0..n {
4666 release_before[k] = previous_release;
4667 if !closes[k] {
4668 continue;
4669 }
4670 let (Some(a), Some(w)) = (arrive[k], release_wait[k]) else {
4671 return Err(misplaced(
4672 k,
4673 "a program that does not arrive at, or wait for, a rendezvous",
4674 ));
4675 };
4676 if a > w {
4677 return Err(misplaced(
4678 k,
4679 "a release wait before the program's own arrival",
4680 ));
4681 }
4682 if previous_release.is_some_and(|prev| prev > a) {
4683 return Err(misplaced(
4684 k,
4685 "a rendezvous entered before the previous one was released",
4686 ));
4687 }
4688 let is_root = root_of.get(&k) == Some(&(p.site, p.engine));
4689 match (is_root, gather_wait[k], posts[k]) {
4690 (true, Some(g), Some((first, last))) => {
4691 if !(a < g && g < first && last < w) {
4692 return Err(misplaced(k, "a root that does not arrive, gather, release and then wait, in that order"));
4693 }
4694 }
4695 (true, _, _) => {
4696 return Err(misplaced(k, "a root that gathers without releasing"))
4697 }
4698 (false, None, None) => {}
4699 (false, _, _) => {
4700 return Err(misplaced(
4701 k,
4702 "a release posted by a program that does not hold the gather",
4703 ))
4704 }
4705 }
4706 previous_release = Some(w);
4707 }
4708 for (i, t) in p.ops.iter().enumerate() {
4709 if rendezvous_sem(&t.op).is_some() {
4710 continue;
4711 }
4712 let Some(&s) = section_of.get(t.stage.idx()).filter(|s| **s < n) else {
4713 continue;
4714 };
4715 if next_arrival[s].is_some_and(|a| i > a) {
4716 return Err(misplaced(
4717 s,
4718 "an op of the section after the arrival that announces it finished",
4719 ));
4720 }
4721 if release_before[s].is_some_and(|w| i < w) {
4722 return Err(misplaced(
4723 s,
4724 "an op of the next section before the release that admits it",
4725 ));
4726 }
4727 }
4728 }
4729 Ok(())
4730 }
4731
4732 fn validate_requirements(&self) -> Result<(), IrViolation> {
4733 for e in &self.edges {
4734 let k = e.handshake.kind();
4735 if !self.requires.handshakes.contains(&k) {
4736 return Err(IrViolation::HandshakeNotRequired { kind: k });
4737 }
4738 }
4739 Ok(())
4740 }
4741}