logicaffeine_compile/codegen_ana/
xbar.rs1use super::cell::{ASrc, AnalogCell, AnalogNetlist};
16
17pub 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
115fn 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
140pub 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
202fn 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}