Skip to main content

logicaffeine_compile/tensor_graph/
link.rs

1//! The link's resolver: every op of a graph answered by a chip's shim, before the chip's own link builds anything.
2//!
3//! An op the chip links whole resolves to itself; an op of a class the chip runs as its parts resolves to its parts,
4//! each resolved in turn; an op the chip runs neither way is unresolved, by its place, its class and the chip's reason.
5//! Every refusal of the graph is reported together, and nothing is dropped: a graph resolves only when every op does.
6
7use super::graph::{Op, Sym, TensorGraph};
8use super::shim::{allows, Allow, Shim};
9use super::symbols::ClassId;
10
11/// An op as the chip runs it: whole, or as its parts.
12#[derive(Clone, Debug)]
13pub enum Linked<'g> {
14    Whole(&'g Op),
15    Parts(Vec<Linked<'g>>),
16}
17
18impl Linked<'_> {
19    /// The ops the chip's link builds, in order.
20    pub fn whole(&self) -> Vec<&Op> {
21        match self {
22            Linked::Whole(op) => vec![op],
23            Linked::Parts(parts) => parts.iter().flat_map(Linked::whole).collect(),
24        }
25    }
26}
27
28/// An op the chip runs neither whole nor as its parts.
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct Unresolved {
31    /// The op's place in the graph.
32    pub op: usize,
33    /// Its class, for a fused op.
34    pub class: Option<ClassId>,
35    /// Its kind: the lifted op's, or the class id.
36    pub kind: String,
37    pub why: String,
38}
39
40impl std::fmt::Display for Unresolved {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        write!(f, "op {} ({}): {}", self.op, self.kind, self.why)
43    }
44}
45
46fn link<'g>(shim: &dyn Shim, op: &'g Op, g: &TensorGraph) -> Result<Linked<'g>, String> {
47    match allows(shim, op, g) {
48        Allow::Native => Ok(Linked::Whole(op)),
49        Allow::Decompose => {
50            let Sym::Fused(f) = &op.sym else { unreachable!("only a fused op decomposes") };
51            f.parts.iter().map(|p| link(shim, p, g)).collect::<Result<_, _>>().map(Linked::Parts)
52        }
53        Allow::Refuse(why) => Err(why),
54    }
55}
56
57/// Every op of `g` as `shim`'s chip runs it, or every op it cannot run.
58pub fn resolve<'g>(g: &'g TensorGraph, shim: &dyn Shim) -> Result<Vec<Linked<'g>>, Vec<Unresolved>> {
59    let mut linked = Vec::with_capacity(g.ops.len());
60    let mut unresolved = Vec::new();
61    for (i, op) in g.ops.iter().enumerate() {
62        match link(shim, op, g) {
63            Ok(l) => linked.push(l),
64            Err(why) => unresolved.push(Unresolved {
65                op: i,
66                class: match &op.sym {
67                    Sym::Fused(f) => Some(f.class),
68                    _ => None,
69                },
70                kind: op.kind().to_string(),
71                why,
72            }),
73        }
74    }
75    if unresolved.is_empty() {
76        Ok(linked)
77    } else {
78        Err(unresolved)
79    }
80}