logicaffeine_compile/tensor_graph/
shim.rs

1//! What a chip allows: the one interface every chip answers through.
2//!
3//! The shared optimizer decides a fusion once, for every chip, by dataflow alone; a chip's shim says whether it runs the
4//! fused op whole. The answer is data the chip declares, consulted before a rule commits, so a chip never receives a
5//! fused op it cannot run:
6//!
7//!   • **native** — the chip has a primitive that computes the fused op whole;
8//!   • **decompose** — the chip runs the fused op as its parts, each of which it allows;
9//!   • **refuse** — the chip runs neither, and says why.
10
11use super::fuse::ChipRule;
12use super::graph::{Op, QuantFormat, Sym, TensorGraph};
13use super::numerics::{Bound, Contract, Keep};
14use super::symbols::ClassId;
15
16/// A chip's answer for one op.
17#[derive(Clone, Debug, PartialEq, Eq)]
18pub enum Allow {
19    Native,
20    Decompose,
21    Refuse(String),
22}
23
24impl Allow {
25    pub fn is_native(&self) -> bool {
26        matches!(self, Allow::Native)
27    }
28}
29
30/// What a chip's link does with a class wherever it meets one.
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub enum Support {
33    /// A primitive of the chip's own computes it whole; the chip then judges each op's shape ([`Shim::fits`]).
34    Native,
35    /// The chip runs its parts, each as it runs it alone.
36    Decompose,
37    /// Neither, and why.
38    Refuse(&'static str),
39}
40
41/// What a chip allows.
42pub trait Shim {
43    /// The chip and the configuration it answers for, e.g. `cuda-lane:rtx3060:TwoPassInt8`.
44    fn name(&self) -> String;
45
46    /// What the chip's link does with `class`. Every chip answers by a match over every class with no catch-all arm,
47    /// so a class added to the registry fails the build until each chip answers it.
48    fn supports(&self, class: ClassId) -> Support;
49
50    /// Whether the chip's link builds `op` as it stands — its shapes, its accesses, its destination: `Ok`, or why not.
51    fn fits(&self, op: &Op, g: &TensorGraph) -> Result<(), String>;
52
53    /// The quantized form `consumer` reads its operand `operand` in, or `None` for one it reads as its values.
54    fn activation(&self, _consumer: &Op, _operand: usize, _g: &TensorGraph) -> Option<QuantFormat> {
55        None
56    }
57
58    /// The chip's own rules, each linked into the table where its position says.
59    fn rules(&self) -> Vec<(Position, ChipRule<'_>)> {
60        Vec::new()
61    }
62
63    /// The numerics contract of the knob point the chip answers for, or `None` while the chip has not declared it.
64    fn numerics(&self) -> Option<Contract>;
65}
66
67/// Whether the chip runs `op`: a fused op of a class it links whole that its link builds as it stands (`Native`), one
68/// of a class it runs as its parts, every part allowed (`Decompose`), or neither, by name — a refused part names the
69/// part and its own reason. A lifted op, a view or a quantize is judged by its shape alone.
70pub fn allows(shim: &dyn Shim, op: &Op, g: &TensorGraph) -> Allow {
71    let fit = |op: &Op| shim.fits(op, g).map_or_else(Allow::Refuse, |()| Allow::Native);
72    let Sym::Fused(f) = &op.sym else { return fit(op) };
73    match shim.supports(f.class) {
74        Support::Native => fit(op),
75        Support::Decompose => {
76            for part in &f.parts {
77                if let Allow::Refuse(why) = allows(shim, part, g) {
78                    return Allow::Refuse(format!("{} runs as its parts, and its {} is refused: {why}", f.class, part.kind()));
79                }
80            }
81            Allow::Decompose
82        }
83        Support::Refuse(why) => Allow::Refuse(format!("{}: {why}", f.class)),
84    }
85}
86
87/// Where a chip's own rule runs in the table: before a shared rule, after one, or after them all.
88#[derive(Clone, Copy, Debug, PartialEq, Eq)]
89pub enum Position {
90    Before(&'static str),
91    After(&'static str),
92    End,
93}
94
95/// The shim that allows everything: the meaning alone, for graph-level tests. It answers no class by name — every class
96/// is one it links whole, whatever the registry adds.
97#[derive(Clone, Copy, Debug, Default)]
98pub struct ReferenceShim;
99
100impl Shim for ReferenceShim {
101    fn name(&self) -> String {
102        "reference".to_string()
103    }
104
105    fn supports(&self, _: ClassId) -> Support {
106        Support::Native
107    }
108
109    fn fits(&self, _: &Op, _: &TensorGraph) -> Result<(), String> {
110        Ok(())
111    }
112
113    fn numerics(&self) -> Option<Contract> {
114        Some(Contract { point: "the meaning".to_string(), activations: None, cache: Keep::Exact, state: Keep::Exact, attention: None, fused: Bound::EXACT, accumulator_bits: 64 })
115    }
116}
117
118/// A chip as if it linked `classes` whole, whatever their shape, and answered everything else as it does: where a rule
119/// building those classes would fire on the chip once its link has them.
120pub struct Linking<'a> {
121    pub base: &'a dyn Shim,
122    pub classes: &'a [ClassId],
123}
124
125impl Shim for Linking<'_> {
126    fn name(&self) -> String {
127        format!("{} linking {:?}", self.base.name(), self.classes)
128    }
129
130    fn supports(&self, class: ClassId) -> Support {
131        if self.classes.contains(&class) {
132            Support::Native
133        } else {
134            self.base.supports(class)
135        }
136    }
137
138    fn fits(&self, op: &Op, g: &TensorGraph) -> Result<(), String> {
139        match &op.sym {
140            Sym::Fused(f) if self.classes.contains(&f.class) => Ok(()),
141            _ => self.base.fits(op, g),
142        }
143    }
144
145    fn activation(&self, consumer: &Op, operand: usize, g: &TensorGraph) -> Option<QuantFormat> {
146        self.base.activation(consumer, operand, g)
147    }
148
149    fn rules(&self) -> Vec<(Position, ChipRule<'_>)> {
150        self.base.rules()
151    }
152
153    fn numerics(&self) -> Option<Contract> {
154        self.base.numerics()
155    }
156}
logicaffeine_compile/tensor_graph/shim.rs

logicaffeine_compile/tensor_graph/
shim.rs
