logicaffeine_compile/spatial/
describe.rs

1//! `describe` — what a lowered program actually DOES, as facts derived from the IR.
2//!
3//! ★ WHY THIS EXISTS: A KERNEL MUST NOT BE ABLE TO ASSERT ITS OWN CORRECTNESS IN PROSE.
4//!
5//! The tests this replaces asked the emitted C++ whether it did the right thing, by substring. That is
6//! a question generated text can always answer favourably. The concrete case:
7//! `fused_row_family_lock` asserted `contains("per-row rmsnorm (weighted)")` and passed for as long as
8//! the retired template PRINTED THAT PHRASE IN A COMMENT — while its body declared `cb_w` and never
9//! read it. The kernel claimed to apply a weight, in English, and did not apply one.
10//!
11//! Swapping the needle for an instruction name (`contains("mul_tiles_bcast_rows")`) does not fix that;
12//! it moves the same hole one level down, because a comment can contain an instruction name just as
13//! easily as a sentence. The only durable answer is to stop asking the TEXT.
14//!
15//! So these facts are computed from the [`SpatialProgram`] — from iterator kinds and access maps, the
16//! things the compiler actually decided — and never from a string. A reduction is counted because an
17//! axis is marked [`IterKind::Reduce`]. A per-column weight is counted because an operand's access map
18//! OMITS the row axis, which is what makes it constant down the rows. Nothing a printer emits, and no
19//! comment anyone writes, can change either number.
20
21use std::collections::BTreeSet;
22
23use super::ir::{ComputeOp, IterKind, SfpuKind, SpatialProgram, StageKind, TensorBody};
24
25/// What a program's compute stages do, in the vocabulary of the algebra rather than of a backend.
26#[derive(Clone, Debug, Default, PartialEq, Eq)]
27pub struct ComputeFacts {
28    /// Compute stages in the program.
29    pub stages: usize,
30    /// Stages that fold an axis away — the row reductions.
31    pub reductions: usize,
32    /// Stages applying a value that is constant DOWN THE ROWS: a per-COLUMN vector, which is what a
33    /// normalization's weight and bias are. The access map omits the row axis.
34    pub per_column_operands: usize,
35    /// Stages applying a value that is constant ACROSS A ROW: a per-ROW scalar, which is what a
36    /// reduction's result is when it is broadcast back. The access map omits the column axis.
37    pub per_row_operands: usize,
38    /// Every SFPU transform the program applies, anywhere in a body (including nested inside
39    /// arithmetic, which is where fusion puts them).
40    pub sfpu: BTreeSet<SfpuKind>,
41}
42
43impl ComputeFacts {
44    /// Whether the program applies a per-column vector — a weight or a bias.
45    pub fn applies_a_per_column_operand(&self) -> bool {
46        self.per_column_operands > 0
47    }
48    /// Whether the program broadcasts a per-row scalar back across its row — a rescale.
49    pub fn rescales_by_a_per_row_scalar(&self) -> bool {
50        self.per_row_operands > 0
51    }
52    /// Whether the program applies a given SFPU transform anywhere.
53    pub fn uses(&self, kind: SfpuKind) -> bool {
54        self.sfpu.contains(&kind)
55    }
56}
57
58fn collect_sfpu(b: &TensorBody, acc: &mut BTreeSet<SfpuKind>) {
59    match b {
60        TensorBody::Unary { kind, arg } => {
61            acc.insert(*kind);
62            collect_sfpu(arg, acc);
63        }
64        TensorBody::Add(x, y)
65        | TensorBody::Sub(x, y)
66        | TensorBody::Mul(x, y)
67        | TensorBody::Min(x, y)
68        | TensorBody::Max(x, y)
69        | TensorBody::FloorDiv(x, y) => {
70            collect_sfpu(x, acc);
71            collect_sfpu(y, acc);
72        }
73        TensorBody::Shr { arg, .. } => collect_sfpu(arg, acc),
74        TensorBody::Select {
75            then, otherwise, ..
76        } => {
77            collect_sfpu(then, acc);
78            collect_sfpu(otherwise, acc);
79        }
80        TensorBody::Operand(_) | TensorBody::Const(_) => {}
81    }
82}
83
84/// The facts of a program, computed from its stages. Cheap and total — it never fails, because every
85/// question it answers is structural.
86pub fn compute_facts(p: &SpatialProgram) -> ComputeFacts {
87    let mut f = ComputeFacts::default();
88    for s in &p.stages {
89        let StageKind::Compute(op) = &s.kind else {
90            continue;
91        };
92        f.stages += 1;
93        let ComputeOp::Generic {
94            iters, ins, body, ..
95        } = op
96        else {
97            continue;
98        };
99        let (row_ax, col_ax) = (iters.first(), iters.get(1));
100        if iters
101            .iter()
102            .any(|a| matches!(a.kind, IterKind::Reduce { .. }))
103        {
104            f.reductions += 1;
105        }
106        collect_sfpu(body, &mut f.sfpu);
107        // An operand CONSTANT along an axis is broadcast along it. The scaler tile a reduction takes is
108        // constant along BOTH axes and is not a broadcast of data, so it is excluded from both counts.
109        for m in ins {
110            let varies_row = row_ax.map(|a| m.varies_along(a.var)).unwrap_or(true);
111            let varies_col = col_ax.map(|a| m.varies_along(a.var)).unwrap_or(true);
112            match (varies_row, varies_col) {
113                (false, true) => f.per_column_operands += 1,
114                (true, false) => f.per_row_operands += 1,
115                _ => {}
116            }
117        }
118    }
119    f
120}
logicaffeine_compile/spatial/describe.rs

logicaffeine_compile/spatial/
describe.rs
