1use std::collections::BTreeMap;
11
12use super::lane::{LInst, LOp, LReg, LStmt, LaneShape, Routine};
13
14#[derive(Clone, Debug, PartialEq, Eq)]
16pub struct Mw(pub Vec<LReg>);
17
18impl Mw {
19 pub fn len(&self) -> usize {
20 self.0.len()
21 }
22
23 pub fn is_empty(&self) -> bool {
24 self.0.is_empty()
25 }
26
27 pub fn slice(&self, from: usize, to: usize) -> Mw {
29 Mw(self.0[from.min(self.len())..to.min(self.len())].to_vec())
30 }
31}
32
33pub struct Builder {
35 shape: LaneShape,
36 regs: u32,
37 consts: BTreeMap<u32, LReg>,
38 prelude: Vec<LStmt>,
39 bodies: Vec<Vec<LStmt>>,
40}
41
42impl Builder {
43 pub fn new(shape: LaneShape) -> Builder {
44 Builder {
45 shape,
46 regs: 0,
47 consts: BTreeMap::new(),
48 prelude: Vec::new(),
49 bodies: vec![Vec::new()],
50 }
51 }
52
53 pub fn shape(&self) -> LaneShape {
54 self.shape
55 }
56
57 pub fn w(&self) -> u32 {
59 self.shape.word
60 }
61
62 pub fn fresh(&mut self) -> LReg {
64 let r = self.regs;
65 self.regs += 1;
66 r
67 }
68
69 fn push(&mut self, inst: LInst) {
70 self.bodies
71 .last_mut()
72 .expect("a body under construction")
73 .push(LStmt::Inst(inst));
74 }
75
76 fn emit(&mut self, op: LOp) -> LReg {
77 let dst = self.fresh();
78 self.push(LInst {
79 dst,
80 dst2: None,
81 op,
82 });
83 dst
84 }
85
86 fn emit2(&mut self, op: LOp) -> (LReg, LReg) {
87 let dst = self.fresh();
88 let dst2 = self.fresh();
89 self.push(LInst {
90 dst,
91 dst2: Some(dst2),
92 op,
93 });
94 (dst, dst2)
95 }
96
97 pub fn k(&mut self, v: u32) -> LReg {
99 assert!(
100 v & !self.shape.mask() == 0,
101 "constant {v:#x} is wider than a {}-bit word",
102 self.shape.word
103 );
104 if let Some(r) = self.consts.get(&v) {
105 return *r;
106 }
107 let dst = self.fresh();
108 self.prelude.push(LStmt::Inst(LInst {
109 dst,
110 dst2: None,
111 op: LOp::Const(v),
112 }));
113 self.consts.insert(v, dst);
114 dst
115 }
116
117 pub fn assign(&mut self, dst: LReg, src: LReg) {
119 if dst != src {
120 self.push(LInst {
121 dst,
122 dst2: None,
123 op: LOp::Mov(src),
124 });
125 }
126 }
127
128 pub fn repeat(&mut self, times: u32, f: impl FnOnce(&mut Builder)) {
130 self.bodies.push(Vec::new());
131 f(self);
132 let body = self.bodies.pop().expect("the repeat body");
133 self.bodies
134 .last_mut()
135 .expect("the enclosing body")
136 .push(LStmt::Repeat { times, body });
137 }
138
139 pub fn finish(
141 mut self,
142 name: impl Into<String>,
143 inputs: Vec<LReg>,
144 outputs: Vec<LReg>,
145 ) -> Routine {
146 let body = self.bodies.pop().expect("the routine body");
147 assert!(self.bodies.is_empty(), "every repeat closed");
148 let mut all = std::mem::take(&mut self.prelude);
149 all.extend(body);
150 Routine {
151 name: name.into(),
152 shape: self.shape,
153 regs: self.regs,
154 inputs,
155 outputs,
156 body: all,
157 }
158 }
159
160 pub fn add(&mut self, a: LReg, b: LReg, cin: LReg) -> (LReg, LReg) {
163 self.emit2(LOp::Add { a, b, cin })
164 }
165
166 pub fn sub(&mut self, a: LReg, b: LReg, bin: LReg) -> (LReg, LReg) {
167 self.emit2(LOp::Sub { a, b, bin })
168 }
169
170 pub fn mul(&mut self, a: LReg, b: LReg) -> LReg {
171 self.emit(LOp::Mul { a, b })
172 }
173
174 pub fn and(&mut self, a: LReg, b: LReg) -> LReg {
175 self.emit(LOp::And { a, b })
176 }
177
178 pub fn or(&mut self, a: LReg, b: LReg) -> LReg {
179 self.emit(LOp::Or { a, b })
180 }
181
182 pub fn xor(&mut self, a: LReg, b: LReg) -> LReg {
183 self.emit(LOp::Xor { a, b })
184 }
185
186 pub fn not(&mut self, a: LReg) -> LReg {
187 self.emit(LOp::Not { a })
188 }
189
190 pub fn shl(&mut self, a: LReg, k: u32) -> LReg {
191 if k == 0 {
192 return a;
193 }
194 assert!(k < self.w());
195 self.emit(LOp::Shl { a, k })
196 }
197
198 pub fn shr(&mut self, a: LReg, k: u32) -> LReg {
199 if k == 0 {
200 return a;
201 }
202 assert!(k < self.w());
203 self.emit(LOp::Shr { a, k })
204 }
205
206 pub fn ltu(&mut self, a: LReg, b: LReg) -> LReg {
207 self.emit(LOp::Ltu { a, b })
208 }
209
210 pub fn eq(&mut self, a: LReg, b: LReg) -> LReg {
211 self.emit(LOp::Eq { a, b })
212 }
213
214 pub fn sel(&mut self, c: LReg, t: LReg, f: LReg) -> LReg {
215 if t == f {
216 return t;
217 }
218 self.emit(LOp::Sel { c, t, f })
219 }
220
221 pub fn bit_not(&mut self, b: LReg) -> LReg {
223 let one = self.k(1);
224 self.xor(b, one)
225 }
226
227 pub fn word_bit(&mut self, a: LReg, i: u32) -> LReg {
229 let s = self.shr(a, i);
230 let one = self.k(1);
231 self.and(s, one)
232 }
233
234 pub fn nonzero(&mut self, a: LReg) -> LReg {
236 let z = self.k(0);
237 let e = self.eq(a, z);
238 self.bit_not(e)
239 }
240
241 pub fn mw_input(&mut self, n: usize) -> Mw {
245 Mw((0..n).map(|_| self.fresh()).collect())
246 }
247
248 pub fn mw_const(&mut self, v: u128, n: usize) -> Mw {
250 let w = self.w();
251 let mask = self.shape.mask() as u128;
252 Mw((0..n)
253 .map(|i| {
254 let at = i as u32 * w;
255 let word = if at >= 128 {
256 0
257 } else {
258 ((v >> at) & mask) as u32
259 };
260 self.k(word)
261 })
262 .collect())
263 }
264
265 pub fn mw_zero(&mut self, n: usize) -> Mw {
266 self.mw_const(0, n)
267 }
268
269 pub fn mw_state(&mut self, a: &Mw) -> Mw {
271 let regs: Vec<LReg> = a.0.iter().map(|_| self.fresh()).collect();
272 for (d, s) in regs.iter().zip(&a.0) {
273 self.push(LInst {
274 dst: *d,
275 dst2: None,
276 op: LOp::Mov(*s),
277 });
278 }
279 Mw(regs)
280 }
281
282 pub fn state(&mut self, a: LReg) -> LReg {
284 let r = self.fresh();
285 self.push(LInst {
286 dst: r,
287 dst2: None,
288 op: LOp::Mov(a),
289 });
290 r
291 }
292
293 pub fn mw_assign(&mut self, dst: &Mw, src: &Mw) {
297 assert_eq!(dst.len(), src.len());
298 let staged: Vec<LReg> = src
299 .0
300 .iter()
301 .zip(&dst.0)
302 .map(|(s, d)| {
303 if s != d && dst.0.contains(s) {
304 self.state(*s)
305 } else {
306 *s
307 }
308 })
309 .collect();
310 for (d, s) in dst.0.iter().zip(&staged) {
311 self.assign(*d, *s);
312 }
313 }
314
315 pub fn mw_resize(&mut self, a: &Mw, n: usize) -> Mw {
317 let z = self.k(0);
318 Mw((0..n).map(|i| a.0.get(i).copied().unwrap_or(z)).collect())
319 }
320
321 pub fn mw_sext(&mut self, a: &Mw, n: usize) -> Mw {
323 if n <= a.len() {
324 return a.slice(0, n);
325 }
326 let s = self.mw_top_bit(a);
327 let ones = self.k(self.shape.mask());
328 let z = self.k(0);
329 let fill = self.sel(s, ones, z);
330 Mw((0..n)
331 .map(|i| a.0.get(i).copied().unwrap_or(fill))
332 .collect())
333 }
334
335 pub fn mw_from_word(&mut self, a: LReg, n: usize) -> Mw {
337 self.mw_resize(&Mw(vec![a]), n)
338 }
339
340 pub fn mw_add(&mut self, a: &Mw, b: &Mw, cin: LReg) -> (Mw, LReg) {
341 assert_eq!(a.len(), b.len());
342 let mut c = cin;
343 let mut out = Vec::with_capacity(a.len());
344 for (x, y) in a.0.iter().zip(&b.0) {
345 let (s, c2) = self.add(*x, *y, c);
346 out.push(s);
347 c = c2;
348 }
349 (Mw(out), c)
350 }
351
352 pub fn mw_plus(&mut self, a: &Mw, b: &Mw) -> Mw {
354 let z = self.k(0);
355 self.mw_add(a, b, z).0
356 }
357
358 pub fn mw_sub(&mut self, a: &Mw, b: &Mw, bin: LReg) -> (Mw, LReg) {
359 assert_eq!(a.len(), b.len());
360 let mut c = bin;
361 let mut out = Vec::with_capacity(a.len());
362 for (x, y) in a.0.iter().zip(&b.0) {
363 let (d, c2) = self.sub(*x, *y, c);
364 out.push(d);
365 c = c2;
366 }
367 (Mw(out), c)
368 }
369
370 pub fn mw_minus(&mut self, a: &Mw, b: &Mw) -> Mw {
372 let z = self.k(0);
373 self.mw_sub(a, b, z).0
374 }
375
376 pub fn mw_neg(&mut self, a: &Mw) -> Mw {
378 let zero = self.mw_zero(a.len());
379 self.mw_minus(&zero, a)
380 }
381
382 pub fn mw_inc(&mut self, a: &Mw, bit: LReg) -> (Mw, LReg) {
384 let zero = self.mw_zero(a.len());
385 self.mw_add(a, &zero, bit)
386 }
387
388 pub fn mw_and(&mut self, a: &Mw, b: &Mw) -> Mw {
389 Mw(a.0
390 .iter()
391 .zip(&b.0)
392 .map(|(x, y)| self.and(*x, *y))
393 .collect())
394 }
395
396 pub fn mw_or(&mut self, a: &Mw, b: &Mw) -> Mw {
397 Mw(a.0.iter().zip(&b.0).map(|(x, y)| self.or(*x, *y)).collect())
398 }
399
400 pub fn mw_xor(&mut self, a: &Mw, b: &Mw) -> Mw {
401 Mw(a.0
402 .iter()
403 .zip(&b.0)
404 .map(|(x, y)| self.xor(*x, *y))
405 .collect())
406 }
407
408 pub fn mw_not(&mut self, a: &Mw) -> Mw {
409 Mw(a.0.iter().map(|x| self.not(*x)).collect())
410 }
411
412 pub fn mw_sel(&mut self, c: LReg, t: &Mw, f: &Mw) -> Mw {
413 assert_eq!(t.len(), f.len());
414 Mw(t.0
415 .iter()
416 .zip(&f.0)
417 .map(|(x, y)| self.sel(c, *x, *y))
418 .collect())
419 }
420
421 pub fn mw_or_reduce(&mut self, a: &Mw) -> LReg {
423 let mut acc = self.k(0);
424 for x in &a.0 {
425 acc = self.or(acc, *x);
426 }
427 acc
428 }
429
430 pub fn mw_is_zero(&mut self, a: &Mw) -> LReg {
432 let r = self.mw_or_reduce(a);
433 let z = self.k(0);
434 self.eq(r, z)
435 }
436
437 pub fn mw_eq(&mut self, a: &Mw, b: &Mw) -> LReg {
439 let x = self.mw_xor(a, b);
440 self.mw_is_zero(&x)
441 }
442
443 pub fn mw_ltu(&mut self, a: &Mw, b: &Mw) -> LReg {
445 let z = self.k(0);
446 let (_, borrow) = self.mw_sub(a, b, z);
447 borrow
448 }
449
450 pub fn mw_lts(&mut self, a: &Mw, b: &Mw) -> LReg {
452 let sa = self.mw_top_bit(a);
453 let sb = self.mw_top_bit(b);
454 let differ = self.xor(sa, sb);
455 let u = self.mw_ltu(a, b);
456 self.sel(differ, sa, u)
457 }
458
459 pub fn mw_top_bit(&mut self, a: &Mw) -> LReg {
461 let top = *a.0.last().expect("a nonempty value");
462 let w = self.w();
463 self.word_bit(top, w - 1)
464 }
465
466 pub fn mw_bit(&mut self, a: &Mw, i: u32) -> LReg {
468 let w = self.w();
469 match a.0.get((i / w) as usize) {
470 Some(word) => self.word_bit(*word, i % w),
471 None => self.k(0),
472 }
473 }
474
475 pub fn mw_shl(&mut self, a: &Mw, s: u32) -> Mw {
477 let w = self.w();
478 let n = a.len();
479 let (words, bits) = ((s / w) as usize, s % w);
480 let z = self.k(0);
481 let word = |i: isize| {
482 if i < 0 || i as usize >= n {
483 z
484 } else {
485 a.0[i as usize]
486 }
487 };
488 let mut out = Vec::with_capacity(n);
489 for i in 0..n {
490 let src = i as isize - words as isize;
491 if bits == 0 {
492 out.push(word(src));
493 } else {
494 let hi = word(src);
495 let lo = word(src - 1);
496 let h = self.shl(hi, bits);
497 let l = self.shr(lo, w - bits);
498 out.push(self.or(h, l));
499 }
500 }
501 Mw(out)
502 }
503
504 pub fn mw_shr(&mut self, a: &Mw, s: u32) -> Mw {
506 let w = self.w();
507 let n = a.len();
508 let (words, bits) = ((s / w) as usize, s % w);
509 let z = self.k(0);
510 let word = |i: usize| if i >= n { z } else { a.0[i] };
511 let mut out = Vec::with_capacity(n);
512 for i in 0..n {
513 let src = i + words;
514 if bits == 0 {
515 out.push(word(src));
516 } else {
517 let lo = word(src);
518 let hi = word(src + 1);
519 let l = self.shr(lo, bits);
520 let h = self.shl(hi, w - bits);
521 out.push(self.or(l, h));
522 }
523 }
524 Mw(out)
525 }
526
527 pub fn mw_shr_sticky(&mut self, a: &Mw, s: u32) -> (Mw, LReg) {
529 let total = a.len() as u32 * self.w();
530 let lost = if s >= total {
531 self.mw_or_reduce(a)
532 } else {
533 let kept_back = self.mw_shl(a, total - s);
534 self.mw_or_reduce(&kept_back)
535 };
536 let sticky = self.nonzero(lost);
537 (self.mw_shr(a, s), sticky)
538 }
539
540 pub fn mw_sar(&mut self, a: &Mw, s: u32) -> Mw {
542 let n = a.len();
543 let wide = self.mw_sext(a, 2 * n);
544 let shifted = self.mw_shr(&wide, s);
545 shifted.slice(0, n)
546 }
547
548 pub fn mw_shl_var(&mut self, a: &Mw, amt: &Mw, bits: u32) -> Mw {
550 let mut cur = a.clone();
551 for j in 0..bits {
552 let t = self.mw_shl(&cur, 1u32 << j);
553 let b = self.mw_bit(amt, j);
554 cur = self.mw_sel(b, &t, &cur);
555 }
556 cur
557 }
558
559 pub fn mw_shr_var_sticky(&mut self, a: &Mw, amt: &Mw, bits: u32) -> (Mw, LReg) {
561 let mut cur = a.clone();
562 let mut sticky = self.k(0);
563 for j in 0..bits {
564 let (t, lost) = self.mw_shr_sticky(&cur, 1u32 << j);
565 let b = self.mw_bit(amt, j);
566 cur = self.mw_sel(b, &t, &cur);
567 let took = self.and(b, lost);
568 sticky = self.or(sticky, took);
569 }
570 (cur, sticky)
571 }
572
573 pub fn mw_sar_var(&mut self, a: &Mw, amt: &Mw, bits: u32) -> Mw {
575 let mut cur = a.clone();
576 for j in 0..bits {
577 let t = self.mw_sar(&cur, 1u32 << j);
578 let b = self.mw_bit(amt, j);
579 cur = self.mw_sel(b, &t, &cur);
580 }
581 cur
582 }
583
584 pub fn mw_clz(&mut self, a: &Mw, n: usize) -> Mw {
587 let total = a.len() as u32 * self.w();
588 assert!(
589 total.is_power_of_two(),
590 "a leading-zero count over a power-of-two width"
591 );
592 let mut x = a.clone();
593 let mut count = self.mw_zero(n);
594 let mut s = total / 2;
595 while s >= 1 {
596 let top = self.mw_shr(&x, total - s);
597 let empty = self.mw_is_zero(&top);
598 let shifted = self.mw_shl(&x, s);
599 x = self.mw_sel(empty, &shifted, &x);
600 let step = self.mw_const(s as u128, n);
601 let zero = self.mw_zero(n);
602 let add = self.mw_sel(empty, &step, &zero);
603 count = self.mw_or(&count, &add);
604 s /= 2;
605 }
606 let top = self.mw_top_bit(&x);
607 let missing = self.bit_not(top);
608 self.mw_inc(&count, missing).0
609 }
610
611 fn digits(&mut self, a: &Mw) -> Vec<LReg> {
613 let (w, m) = (self.w(), self.shape.mul);
614 let dmask = self.k((1u32 << m) - 1);
615 let mut out = Vec::with_capacity(a.len() * (w / m) as usize);
616 for x in &a.0 {
617 for t in 0..(w / m) {
618 let s = self.shr(*x, t * m);
619 out.push(self.and(s, dmask));
620 }
621 }
622 out
623 }
624
625 pub fn mw_mul(&mut self, a: &Mw, b: &Mw) -> Mw {
628 let (w, m) = (self.w(), self.shape.mul);
629 let n = a.len() + b.len();
630 let da = self.digits(a);
631 let db = self.digits(b);
632 let mut columns: Vec<Vec<LReg>> = vec![Vec::new(); n + 1];
633 for (i, x) in da.iter().enumerate() {
634 for (j, y) in db.iter().enumerate() {
635 let p = self.mul(*x, *y);
636 let at = (i + j) as u32 * m;
637 let (k, o) = ((at / w) as usize, at % w);
638 if o == 0 {
639 columns[k].push(p);
640 } else {
641 let lo = self.shl(p, o);
642 columns[k].push(lo);
643 if o + 2 * m > w {
644 let hi = self.shr(p, w - o);
645 columns[k + 1].push(hi);
646 }
647 }
648 }
649 }
650 let most = columns.iter().map(Vec::len).max().unwrap_or(0) as u64;
651 let acc_words = self.shape.words_for(w + 64 - (most + 1).leading_zeros());
652 let mut result = self.mw_zero(n);
653 for (k, col) in columns.iter().enumerate() {
654 if col.is_empty() || k >= n {
655 continue;
656 }
657 let mut acc = self.mw_zero(acc_words);
658 for x in col {
659 let xw = self.mw_from_word(*x, acc_words);
660 acc = self.mw_plus(&acc, &xw);
661 }
662 let placed: Vec<LReg> = (0..n)
663 .map(|i| {
664 if i >= k && i - k < acc_words {
665 acc.0[i - k]
666 } else {
667 self.k(0)
668 }
669 })
670 .collect();
671 result = self.mw_plus(&result, &Mw(placed));
672 }
673 result
674 }
675
676 pub fn mw_divrem(&mut self, a: &Mw, b: &Mw) -> (Mw, Mw) {
679 assert_eq!(a.len(), b.len());
680 let n = a.len();
681 let bits = n as u32 * self.w();
682 let q = self.mw_state(a);
683 let zero = self.mw_zero(n + 1);
684 let r = self.mw_state(&zero);
685 let bw = self.mw_resize(b, n + 1);
686 self.repeat(bits, |s| {
687 let top = s.mw_top_bit(&q);
688 let q2 = s.mw_shl(&q, 1);
689 let r_shift = s.mw_shl(&r, 1);
690 let low = r_shift.0[0];
691 let low2 = s.or(low, top);
692 let mut r2 = r_shift.clone();
693 r2.0[0] = low2;
694 let z = s.k(0);
695 let (d, borrow) = s.mw_sub(&r2, &bw, z);
696 let fits = s.bit_not(borrow);
697 let r3 = s.mw_sel(fits, &d, &r2);
698 let q_low = s.or(q2.0[0], fits);
699 let mut q3 = q2.clone();
700 q3.0[0] = q_low;
701 s.mw_assign(&q, &q3);
702 s.mw_assign(&r, &r3);
703 });
704 (q, r.slice(0, n))
705 }
706
707 pub fn mw_popcount(&mut self, a: &Mw, n: usize) -> Mw {
709 let total = a.len() as u32 * self.w();
710 let mut count = self.mw_zero(n);
711 for i in 0..total {
712 let b = self.mw_bit(a, i);
713 count = self.mw_inc(&count, b).0;
714 }
715 count
716 }
717}