logicaffeine_compile/tensor_graph/
link.rs1use super::graph::{Op, Sym, TensorGraph};
8use super::shim::{allows, Allow, Shim};
9use super::symbols::ClassId;
10
11#[derive(Clone, Debug)]
13pub enum Linked<'g> {
14 Whole(&'g Op),
15 Parts(Vec<Linked<'g>>),
16}
17
18impl Linked<'_> {
19 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#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct Unresolved {
31 pub op: usize,
33 pub class: Option<ClassId>,
35 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
57pub 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}