Skip to main content

logicaffeine_compile/codegen_ana/
xbar.rs

1//! Crossbar canonicalization — folding a gain/sum tree into single Kirchhoff columns.
2//!
3//! `Sum` is linear, so a `GainConst` feeding a `Sum` input is the same node with the weight
4//! product (`w·(k·x) = (w·k)·x` in ℝ), and a `Sum` feeding a `Sum` splices its addends with
5//! distributed weights. Driven to fixpoint, an unrolled constant-weight dot product collapses to
6//! ONE `Sum { inputs, weights }` — one physical current-summing node whose programmed
7//! conductances ARE the weight vector: the crossbar column (`I = V·G`).
8//!
9//! The fold RE-ASSOCIATES float arithmetic, so the folded netlist's ideal evaluation is no
10//! longer bit-identical to the VM's op order — it is equivalent in ℝ and compared through
11//! `logos_approx_eq` / the error budget (which carries an explicit f64-rounding term). Callers
12//! choose: the unfolded netlist for the bit-exact ideal tier, the folded one for the physical
13//! crossbar shape. Both evaluate with the same [`super::eval`] entry points.
14
15use super::cell::{ASrc, AnalogCell, AnalogNetlist};
16
17/// Fold to fixpoint, prune, then TREE the columns back down to the physical summing fan-in:
18/// a folded column wider than `max_fanin` splits into balanced sub-sums (weights ride on the
19/// leaves; the joining nodes take unit weights), because fan-in is a physical resource — the
20/// summing node's loading, settling, and noise all grow with it.
21pub fn fold_crossbar_for(netlist: &AnalogNetlist, max_fanin: usize) -> AnalogNetlist {
22    assert!(max_fanin >= 2, "a summing node needs at least two inputs");
23    let folded = fold_crossbar(netlist);
24    let mut cells: Vec<AnalogCell> = Vec::new();
25    let mut remap: Vec<ASrc> = Vec::with_capacity(folded.cells.len());
26    for cell in &folded.cells {
27        let placed = match cell {
28            AnalogCell::Sum { inputs, weights } => {
29                let leaves: Vec<(ASrc, f64)> = inputs
30                    .iter()
31                    .zip(weights)
32                    .map(|(x, w)| (relabel(x, &remap), *w))
33                    .collect();
34                tree_sum(&leaves, max_fanin, &mut cells)
35            }
36            AnalogCell::Mul { a, b } => {
37                let c = AnalogCell::Mul {
38                    a: relabel(a, &remap),
39                    b: relabel(b, &remap),
40                };
41                cells.push(c);
42                ASrc::Val(cells.len() - 1)
43            }
44            AnalogCell::GainConst { a, k } => {
45                let c = AnalogCell::GainConst {
46                    a: relabel(a, &remap),
47                    k: *k,
48                };
49                cells.push(c);
50                ASrc::Val(cells.len() - 1)
51            }
52            AnalogCell::Select { c, t, f } => {
53                let cell = AnalogCell::Select {
54                    c: relabel(c, &remap),
55                    t: relabel(t, &remap),
56                    f: relabel(f, &remap),
57                };
58                cells.push(cell);
59                ASrc::Val(cells.len() - 1)
60            }
61            AnalogCell::Exp { a } => {
62                let cell = AnalogCell::Exp {
63                    a: relabel(a, &remap),
64                };
65                cells.push(cell);
66                ASrc::Val(cells.len() - 1)
67            }
68            AnalogCell::Log { a } => {
69                let cell = AnalogCell::Log {
70                    a: relabel(a, &remap),
71                };
72                cells.push(cell);
73                ASrc::Val(cells.len() - 1)
74            }
75            AnalogCell::Tanh { a } => {
76                let cell = AnalogCell::Tanh {
77                    a: relabel(a, &remap),
78                };
79                cells.push(cell);
80                ASrc::Val(cells.len() - 1)
81            }
82            AnalogCell::Sqrt { a } => {
83                let cell = AnalogCell::Sqrt {
84                    a: relabel(a, &remap),
85                };
86                cells.push(cell);
87                ASrc::Val(cells.len() - 1)
88            }
89            AnalogCell::Div { a, b } => {
90                let cell = AnalogCell::Div {
91                    a: relabel(a, &remap),
92                    b: relabel(b, &remap),
93                };
94                cells.push(cell);
95                ASrc::Val(cells.len() - 1)
96            }
97        };
98        remap.push(placed);
99    }
100    let out = relabel(&folded.out, &remap);
101    AnalogNetlist {
102        n_in: folded.n_in,
103        cells,
104        out,
105    }
106}
107
108fn relabel(s: &ASrc, remap: &[ASrc]) -> ASrc {
109    match s {
110        ASrc::Val(v) => remap[*v],
111        other => *other,
112    }
113}
114
115/// Emit a ≤`max_fanin` tree over weighted leaves; joining nodes carry unit weights.
116fn tree_sum(leaves: &[(ASrc, f64)], max_fanin: usize, cells: &mut Vec<AnalogCell>) -> ASrc {
117    if leaves.len() == 1 {
118        let (x, w) = leaves[0];
119        cells.push(AnalogCell::Sum {
120            inputs: vec![x],
121            weights: vec![w],
122        });
123        return ASrc::Val(cells.len() - 1);
124    }
125    if leaves.len() <= max_fanin {
126        cells.push(AnalogCell::Sum {
127            inputs: leaves.iter().map(|(x, _)| *x).collect(),
128            weights: leaves.iter().map(|(_, w)| *w).collect(),
129        });
130        return ASrc::Val(cells.len() - 1);
131    }
132    let chunk = leaves.len().div_ceil(max_fanin).max(1);
133    let subs: Vec<(ASrc, f64)> = leaves
134        .chunks(chunk.max(1))
135        .map(|c| (tree_sum(c, max_fanin, cells), 1.0))
136        .collect();
137    tree_sum(&subs, max_fanin, cells)
138}
139
140/// Fold gain-into-sum and sum-into-sum to fixpoint, then prune dead tiles. Order-preserving:
141/// spliced addends keep their producer's addend order at the consumer's position, so the
142/// evaluation order stays deterministic.
143pub fn fold_crossbar(netlist: &AnalogNetlist) -> AnalogNetlist {
144    let mut cells = netlist.cells.clone();
145    loop {
146        let mut changed = false;
147        for i in 0..cells.len() {
148            let AnalogCell::Sum { inputs, weights } = &cells[i] else {
149                continue;
150            };
151            let mut new_inputs: Vec<ASrc> = Vec::with_capacity(inputs.len());
152            let mut new_weights: Vec<f64> = Vec::with_capacity(weights.len());
153            let mut folded_here = false;
154            for (x, w) in inputs.iter().zip(weights) {
155                match x {
156                    ASrc::Val(j) => match &cells[*j] {
157                        AnalogCell::GainConst { a, k } => {
158                            new_inputs.push(*a);
159                            new_weights.push(w * k);
160                            folded_here = true;
161                        }
162                        AnalogCell::Sum {
163                            inputs: is2,
164                            weights: ws2,
165                        } => {
166                            for (x2, w2) in is2.iter().zip(ws2) {
167                                new_inputs.push(*x2);
168                                new_weights.push(w * w2);
169                            }
170                            folded_here = true;
171                        }
172                        _ => {
173                            new_inputs.push(*x);
174                            new_weights.push(*w);
175                        }
176                    },
177                    _ => {
178                        new_inputs.push(*x);
179                        new_weights.push(*w);
180                    }
181                }
182            }
183            if folded_here {
184                cells[i] = AnalogCell::Sum {
185                    inputs: new_inputs,
186                    weights: new_weights,
187                };
188                changed = true;
189            }
190        }
191        if !changed {
192            break;
193        }
194    }
195    prune(&AnalogNetlist {
196        n_in: netlist.n_in,
197        cells,
198        out: netlist.out,
199    })
200}
201
202/// Drop tiles unreachable from the output and compact the SSA numbering (the digital fabric's
203/// `prune_dead_cells` discipline, over analog tiles).
204fn prune(netlist: &AnalogNetlist) -> AnalogNetlist {
205    let mut live = vec![false; netlist.cells.len()];
206    let mut stack: Vec<usize> = Vec::new();
207    let mark = |s: &ASrc, stack: &mut Vec<usize>| {
208        if let ASrc::Val(v) = s {
209            stack.push(*v);
210        }
211    };
212    mark(&netlist.out, &mut stack);
213    while let Some(v) = stack.pop() {
214        if live[v] {
215            continue;
216        }
217        live[v] = true;
218        match &netlist.cells[v] {
219            AnalogCell::Sum { inputs, .. } => {
220                for x in inputs {
221                    mark(x, &mut stack);
222                }
223            }
224            AnalogCell::Mul { a, b } => {
225                mark(a, &mut stack);
226                mark(b, &mut stack);
227            }
228            AnalogCell::GainConst { a, .. } => mark(a, &mut stack),
229            AnalogCell::Select { c, t, f } => {
230                mark(c, &mut stack);
231                mark(t, &mut stack);
232                mark(f, &mut stack);
233            }
234            AnalogCell::Exp { a }
235            | AnalogCell::Log { a }
236            | AnalogCell::Tanh { a }
237            | AnalogCell::Sqrt { a } => mark(a, &mut stack),
238            AnalogCell::Div { a, b } => {
239                mark(a, &mut stack);
240                mark(b, &mut stack);
241            }
242        }
243    }
244    let mut remap = vec![usize::MAX; netlist.cells.len()];
245    let mut kept: Vec<AnalogCell> = Vec::new();
246    for (i, c) in netlist.cells.iter().enumerate() {
247        if live[i] {
248            remap[i] = kept.len();
249            kept.push(c.clone());
250        }
251    }
252    let rl = |s: &ASrc| match s {
253        ASrc::Val(v) => ASrc::Val(remap[*v]),
254        other => *other,
255    };
256    for c in kept.iter_mut() {
257        match c {
258            AnalogCell::Sum { inputs, .. } => {
259                for x in inputs.iter_mut() {
260                    *x = rl(x);
261                }
262            }
263            AnalogCell::Mul { a, b } => {
264                *a = rl(a);
265                *b = rl(b);
266            }
267            AnalogCell::GainConst { a, .. } => *a = rl(a),
268            AnalogCell::Select { c, t, f } => {
269                *c = rl(c);
270                *t = rl(t);
271                *f = rl(f);
272            }
273            AnalogCell::Exp { a }
274            | AnalogCell::Log { a }
275            | AnalogCell::Tanh { a }
276            | AnalogCell::Sqrt { a } => *a = rl(a),
277            AnalogCell::Div { a, b } => {
278                *a = rl(a);
279                *b = rl(b);
280            }
281        }
282    }
283    AnalogNetlist {
284        n_in: netlist.n_in,
285        cells: kept,
286        out: rl(&netlist.out),
287    }
288}