logicaffeine_forge/
vectorize.rs

1//! Element-wise map vectorization — the general, sound SIMD lever.
2//!
3//! Recognizes a region that is a PURE ELEMENT-WISE MAP — a loop `for i in 1..=n`
4//! (1-based, the Logos array convention) whose body reads one or more arrays at
5//! index `i`, computes a straight-line float expression, and writes the result
6//! to an array at index `i`, with the
7//! induction variable `i` as the ONLY loop-carried value. Such a loop is
8//! bit-exactly vectorizable: lane `i` computes exactly `f(a[i], b[i], ...)`
9//! whether scalar or packed, so 2-wide (or wider) SIMD changes nothing about the
10//! per-element result — unlike a REDUCTION (`sum += a[i]`), whose 2-lane form
11//! reassociates the float adds and is therefore NOT bit-identical (and is
12//! rejected here).
13//!
14//! This is the recognizer half (a pure analysis). The codegen half lowers each
15//! body op to its packed form (`AddF`->`addpd`, …) over 2 lanes with a scalar
16//! tail; the packed primitives live in [`crate::x64asm`].
17
18use crate::jit::{Cmp, MicroOp, Slot};
19use std::collections::HashSet;
20
21/// A recognized element-wise map loop. The straight-line load/compute/store body
22/// is `ops[body_start..body_end]`; the guard, increment, and back-edge surround it.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct MapPlan {
25    /// The induction variable slot `i` (a 1-based array index).
26    pub induction: Slot,
27    /// The loop-limit slot `n` (the guard `Branch` right operand).
28    pub limit: Slot,
29    /// The guard comparison: `Lt` (`i < n`) or `LtEq` (`i <= n`). Decides how many
30    /// elements a full SIMD pair needs and where the scalar tail begins.
31    pub cmp: Cmp,
32    /// The straight-line load/compute/store body is `ops[body_start..body_end]`.
33    pub body_start: usize,
34    pub body_end: usize,
35}
36
37/// Returns `Some(MapPlan)` iff `ops` is a pure element-wise map region in the
38/// top-tested `while` shape that `adapt_region` produces:
39///
40/// ```text
41///   [0]          Branch { cmp(Lt|LtEq), lhs:i, rhs:n, .. }   // loop guard (exit when cmp FALSE)
42///   [1..j-2]     body: ArrLoad{idx==i,checked:false} / float arith / ArrStore{idx==i,checked:false}
43///   [j-2]        LoadConst { dst:step, value:1 }
44///   [j-1]        Add { dst:i, lhs:i, rhs:step }              // increment
45///   [j]          Jump { target:0 }                           // back-edge to head
46///   [j+1..]      exit (Return)
47/// ```
48///
49/// Any reduction (loop-carried operand), cross-index access, bounds-checked
50/// access, non-unit stride, call, or other op makes it `None` — the
51/// conservative, sound side.
52pub fn recognize_elementwise_map(ops: &[MicroOp]) -> Option<MapPlan> {
53    // guard + >=1 body op + step LoadConst + increment + back-Jump.
54    if ops.len() < 5 {
55        return None;
56    }
57
58    // [0]: the loop guard — `Branch` jumps to its target when the comparison is
59    // FALSE, i.e. it falls through (continues) while `i </<= n`.
60    let (induction, limit, cmp) = match &ops[0] {
61        MicroOp::Branch {
62            cmp: c @ (Cmp::Lt | Cmp::LtEq),
63            lhs,
64            rhs,
65            ..
66        } => (*lhs, *rhs, *c),
67        _ => return None,
68    };
69
70    // The back-edge `Jump { target: 0 }`.
71    let j = ops
72        .iter()
73        .position(|o| matches!(o, MicroOp::Jump { target: 0 }))?;
74    if j < 4 {
75        return None; // need [0]guard [1..]body [j-2]LoadConst [j-1]increment
76    }
77
78    // [j-1]: the induction increment `i = i + step`.
79    let step = match &ops[j - 1] {
80        MicroOp::Add { dst, lhs, rhs } if *dst == induction && *lhs == induction => *rhs,
81        _ => return None,
82    };
83    // [j-2]: the stride-1 constant feeding it. The kernel processes EVERY index,
84    // so a non-unit / unknown stride (a loop over a subset) must be rejected.
85    match &ops[j - 2] {
86        MicroOp::LoadConst { dst, value } if *dst == step && *value == 1 => {}
87        _ => return None,
88    }
89
90    // The straight-line body is everything between the guard and the increment
91    // machinery: `ops[1..j-2]` (load/compute/store only — no LoadConst/Add here).
92    let (body_start, body_end) = (1usize, j - 2);
93    let body = &ops[body_start..body_end];
94    if body.is_empty() {
95        return None;
96    }
97
98    // The body must be straight-line: every slot an op reads is either the
99    // induction variable, a constant, or a value DEFINED EARLIER IN THE BODY.
100    // A read of a slot not yet defined in the body is a loop-carried value
101    // (e.g. a reduction accumulator) — NOT bit-exactly vectorizable -> reject.
102    let mut defined: HashSet<Slot> = HashSet::new();
103    defined.insert(induction);
104    let mut wrote_an_array = false;
105
106    let known = |s: &Slot, defined: &HashSet<Slot>| defined.contains(s);
107
108    for op in body {
109        match op {
110            MicroOp::LoadConst { dst, .. } => {
111                defined.insert(*dst);
112            }
113            MicroOp::ArrLoad { dst, idx, checked, .. } => {
114                if *idx != induction || *checked {
115                    // cross-index read, OR a bounds-checked access whose length the
116                    // kernel can't prove >= n (the kernel elides per-element checks,
117                    // so it is sound ONLY when the Oracle already proved in-bounds,
118                    // i.e. `checked == false`).
119                    return None;
120                }
121                defined.insert(*dst);
122            }
123            MicroOp::AddF { dst, lhs, rhs }
124            | MicroOp::SubF { dst, lhs, rhs }
125            | MicroOp::MulF { dst, lhs, rhs }
126            | MicroOp::DivF { dst, lhs, rhs } => {
127                if !known(lhs, &defined) || !known(rhs, &defined) {
128                    return None; // loop-carried operand (reduction) — unsound
129                }
130                defined.insert(*dst);
131            }
132            MicroOp::SqrtF { dst, src } => {
133                if !known(src, &defined) {
134                    return None;
135                }
136                defined.insert(*dst);
137            }
138            MicroOp::ArrStore { src, idx, checked, .. } => {
139                if *idx != induction || *checked || !known(src, &defined) {
140                    return None;
141                }
142                wrote_an_array = true;
143            }
144            // Any other op (Call, Move, integer arithmetic, comparisons,
145            // list/map/byte ops, a second induction) means this is not a pure
146            // element-wise float map — reject conservatively.
147            _ => return None,
148        }
149    }
150
151    // A map must actually STORE a result; a body that only loads is not a map.
152    if !wrote_an_array {
153        return None;
154    }
155
156    Some(MapPlan {
157        induction,
158        limit,
159        cmp,
160        body_start,
161        body_end,
162    })
163}
164
165/// Lower one straight-line float-arithmetic body op to its 2-wide PACKED form,
166/// computing `dst = lhs OP rhs` over both lanes. `xmm(slot)` maps a body slot to
167/// its assigned lane register; `scratch` is a free XMM for the awkward
168/// `dst == rhs` non-commutative case. SSE packed ops are 2-operand (`x OP= y`),
169/// so this materializes the 3-operand `dst = lhs OP rhs` form, never clobbering
170/// an operand it still needs. Loads/stores/LoadConst are handled by the loop
171/// driver, not here.
172#[cfg(target_arch = "x86_64")]
173pub fn emit_packed_arith(
174    asm: &mut crate::x64asm::Asm,
175    op: &MicroOp,
176    xmm: impl Fn(Slot) -> crate::x64asm::Xmm,
177    scratch: crate::x64asm::Xmm,
178) {
179    use crate::x64asm::Asm;
180    match op {
181        MicroOp::AddF { dst, lhs, rhs } => {
182            emit_bin(asm, xmm(*dst), xmm(*lhs), xmm(*rhs), scratch, Asm::addpd_rr, true)
183        }
184        MicroOp::MulF { dst, lhs, rhs } => {
185            emit_bin(asm, xmm(*dst), xmm(*lhs), xmm(*rhs), scratch, Asm::mulpd_rr, true)
186        }
187        MicroOp::SubF { dst, lhs, rhs } => {
188            emit_bin(asm, xmm(*dst), xmm(*lhs), xmm(*rhs), scratch, Asm::subpd_rr, false)
189        }
190        MicroOp::DivF { dst, lhs, rhs } => {
191            emit_bin(asm, xmm(*dst), xmm(*lhs), xmm(*rhs), scratch, Asm::divpd_rr, false)
192        }
193        MicroOp::SqrtF { dst, src } => asm.sqrtpd_rr(xmm(*dst), xmm(*src)),
194        _ => {}
195    }
196}
197
198/// Materialize `d = l OP r` from a 2-operand packed `op(x, y): x OP= y`, without
199/// destroying an operand before it is consumed.
200#[cfg(target_arch = "x86_64")]
201fn emit_bin(
202    asm: &mut crate::x64asm::Asm,
203    d: crate::x64asm::Xmm,
204    l: crate::x64asm::Xmm,
205    r: crate::x64asm::Xmm,
206    scratch: crate::x64asm::Xmm,
207    op: fn(&mut crate::x64asm::Asm, crate::x64asm::Xmm, crate::x64asm::Xmm),
208    commutative: bool,
209) {
210    if d == l {
211        op(asm, d, r); // d already holds l
212    } else if d == r {
213        if commutative {
214            op(asm, d, l); // d = r OP l == l OP r
215        } else {
216            asm.movupd_rr(scratch, l);
217            op(asm, scratch, r); // scratch = l OP r (r untouched until read)
218            asm.movupd_rr(d, scratch);
219        }
220    } else {
221        asm.movupd_rr(d, l);
222        op(asm, d, r);
223    }
224}
225
226/// Emit a complete vectorized loop kernel for a recognized element-wise map,
227/// frame-ABI `extern "C" fn(*mut i64) -> i64`: it reads the induction `i`, limit
228/// `n`, and each array base pointer from their frame slots (`[rdi + slot*8]`),
229/// runs a 2-wide packed loop (`i += 2`) over full pairs, then a single-element
230/// scalar tail that REUSES the packed body (load one element with `movsd` so lane
231/// 1 is zero, run the same packed ops, store lane 0 with `movsd` — the junk lane
232/// is never stored and masked SSE faults are harmless), writes `i = n` back, and
233/// returns 0. Returns `None` (caller falls back to the scalar region) if the body
234/// needs more than 4 distinct arrays or 14 lane temps, or uses an op the kernel
235/// does not lower (e.g. `LoadConst`).
236///
237/// Register plan (all SysV caller-saved, no callee-save needed): rdi=frame,
238/// rsi=i, rdx=n, rcx=i*8 offset, rax=address scratch, r8..r11=array bases,
239/// xmm0..xmm13=lane temps, xmm15=arith scratch.
240#[cfg(target_arch = "x86_64")]
241pub fn emit_map_kernel(body: &[MicroOp], plan: &MapPlan) -> Option<Vec<u8>> {
242    use crate::x64asm::{Asm, Cond, Reg, Xmm};
243    const ARRAY_REGS: [Reg; 4] = [Reg::R8, Reg::R9, Reg::R10, Reg::R11];
244    const LANES: [Xmm; 14] = [
245        Xmm::Xmm0, Xmm::Xmm1, Xmm::Xmm2, Xmm::Xmm3, Xmm::Xmm4, Xmm::Xmm5, Xmm::Xmm6,
246        Xmm::Xmm7, Xmm::Xmm8, Xmm::Xmm9, Xmm::Xmm10, Xmm::Xmm11, Xmm::Xmm12, Xmm::Xmm13,
247    ];
248    let scratch = Xmm::Xmm15;
249
250    // Assign each distinct array pointer slot a GP base register.
251    let mut arr_reg: Vec<(Slot, Reg)> = Vec::new();
252    for op in body {
253        let ptr = match op {
254            MicroOp::ArrLoad { ptr_slot, .. } | MicroOp::ArrStore { ptr_slot, .. } => Some(*ptr_slot),
255            _ => None,
256        };
257        if let Some(p) = ptr {
258            if !arr_reg.iter().any(|(s, _)| *s == p) {
259                if arr_reg.len() >= ARRAY_REGS.len() {
260                    return None;
261                }
262                arr_reg.push((p, ARRAY_REGS[arr_reg.len()]));
263            }
264        }
265    }
266    let reg_of = |slot: Slot| arr_reg.iter().find(|(s, _)| *s == slot).map(|(_, r)| *r);
267
268    // Assign each load/arith destination slot a lane XMM register.
269    let mut lane: Vec<(Slot, Xmm)> = Vec::new();
270    let mut add_lane = |slot: Slot, lane: &mut Vec<(Slot, Xmm)>| -> Option<()> {
271        if !lane.iter().any(|(s, _)| *s == slot) {
272            if lane.len() >= LANES.len() {
273                return None;
274            }
275            lane.push((slot, LANES[lane.len()]));
276        }
277        Some(())
278    };
279    for op in body {
280        match op {
281            MicroOp::ArrLoad { dst, .. }
282            | MicroOp::AddF { dst, .. }
283            | MicroOp::SubF { dst, .. }
284            | MicroOp::MulF { dst, .. }
285            | MicroOp::DivF { dst, .. }
286            | MicroOp::SqrtF { dst, .. } => add_lane(*dst, &mut lane)?,
287            MicroOp::ArrStore { .. } => {}
288            _ => return None, // unsupported op (e.g. LoadConst) — bail to scalar
289        }
290    }
291    let xmm_of = |slot: Slot| lane.iter().find(|(s, _)| *s == slot).map(|(_, x)| *x);
292
293    let off = |slot: Slot| (slot as i32) * 8;
294    let mut a = Asm::new();
295    // Prologue: load `i`, `n`, and point each array register at `&buffer[i-1]`
296    // (the 1-based element `i`). The inner loop then advances these pointers
297    // instead of recomputing `base + (i-1)*8` every iteration.
298    a.mov_rm(Reg::Rsi, Reg::Rdi, off(plan.induction)); // i
299    a.mov_rm(Reg::Rdx, Reg::Rdi, off(plan.limit)); // n
300    a.mov_rr(Reg::Rcx, Reg::Rsi);
301    a.sub_ri(Reg::Rcx, 1);
302    a.shl_ri(Reg::Rcx, 3); // off0 = (i-1)*8
303    for (slot, reg) in &arr_reg {
304        a.mov_rm(*reg, Reg::Rdi, off(*slot)); // base
305        a.add_rr(*reg, Reg::Rcx); // base + off0 = &buffer[i-1]
306    }
307
308    // Emit the load/arith/store body at a given element width (`packed`). The
309    // array pointers already hold `&buffer[i-1]`, so each access is a direct
310    // `movupd [ptr]` / `movsd [ptr]` — no per-iteration address arithmetic.
311    let emit_body = |a: &mut Asm, packed: bool| {
312        for op in body {
313            match op {
314                MicroOp::ArrLoad { dst, ptr_slot, .. } => {
315                    let p = reg_of(*ptr_slot).unwrap();
316                    let x = xmm_of(*dst).unwrap();
317                    if packed {
318                        a.movupd_rm(x, p, 0);
319                    } else {
320                        a.movsd_rm(x, p, 0);
321                    }
322                }
323                MicroOp::ArrStore { src, ptr_slot, .. } => {
324                    let p = reg_of(*ptr_slot).unwrap();
325                    let x = xmm_of(*src).unwrap();
326                    if packed {
327                        a.movupd_mr(p, 0, x);
328                    } else {
329                        a.movsd_mr(p, 0, x);
330                    }
331                }
332                _ => emit_packed_arith(a, op, |s| xmm_of(s).unwrap(), scratch),
333            }
334        }
335    };
336
337    // Exit the loop when the (1-based) index exceeds the valid range: for
338    // `i <= n` (LtEq) that is `index > n`; for `i < n` (Lt) it is `index >= n`.
339    let exit_cond = if plan.cmp == Cmp::LtEq { Cond::Gt } else { Cond::Ge };
340
341    // Pair loop: process elements i and i+1 while BOTH are in range (i+1 valid).
342    let pair_top = a.new_label();
343    let tail = a.new_label();
344    let done = a.new_label();
345    a.bind(pair_top);
346    a.mov_rr(Reg::Rax, Reg::Rsi);
347    a.add_ri(Reg::Rax, 1); // i+1
348    a.cmp_rr(Reg::Rax, Reg::Rdx);
349    a.jcc(exit_cond, tail); // i+1 out of range -> no full pair left
350    emit_body(&mut a, true);
351    for (_, reg) in &arr_reg {
352        a.add_ri(*reg, 16); // advance each pointer by 2 f64 elements
353    }
354    a.add_ri(Reg::Rsi, 2);
355    a.jmp(pair_top);
356
357    // Scalar tail: at most one element (i still in range, i+1 not).
358    a.bind(tail);
359    a.cmp_rr(Reg::Rsi, Reg::Rdx);
360    a.jcc(exit_cond, done);
361    emit_body(&mut a, false);
362
363    a.bind(done);
364    // Leave the induction at its scalar terminal value: n+1 for `i <= n`, n for
365    // `i < n` (the first value that fails the loop test), so post-loop bytecode
366    // that reads `i` sees exactly what the scalar loop would have left.
367    a.mov_rr(Reg::Rax, Reg::Rdx);
368    if plan.cmp == Cmp::LtEq {
369        a.add_ri(Reg::Rax, 1);
370    }
371    a.mov_mr(Reg::Rdi, off(plan.induction), Reg::Rax);
372    a.xor_rr(Reg::Rax, Reg::Rax);
373    a.ret();
374    Some(a.resolve())
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    // `c[i] = a[i] + b[i]` — the canonical map. Slots: i=0, n=1, a=2, b=3, c=4,
382    // alen=5, blen=6, clen=7, ta=8, tb=9, tc=10, one=11.
383    fn load(dst: Slot, idx: Slot, ptr_slot: Slot, len_slot: Slot) -> MicroOp {
384        MicroOp::ArrLoad { dst, idx, ptr_slot, len_slot, byte: false, narrow32: false, checked: false }
385    }
386    fn store(src: Slot, idx: Slot, ptr_slot: Slot, len_slot: Slot) -> MicroOp {
387        MicroOp::ArrStore { src, idx, ptr_slot, len_slot, byte: false, narrow32: false, checked: false }
388    }
389    fn load_checked(dst: Slot, idx: Slot, ptr_slot: Slot, len_slot: Slot) -> MicroOp {
390        MicroOp::ArrLoad { dst, idx, ptr_slot, len_slot, byte: false, narrow32: false, checked: true }
391    }
392
393    // Wrap a load/compute/store `body` in the top-tested while-loop shape the
394    // recognizer expects: guard / body / LoadConst(step) / increment / back-Jump /
395    // Return. Induction = slot 0, limit = slot 1, step slot = 90.
396    fn loop_region(body: Vec<MicroOp>, cmp: Cmp, step_val: i64) -> Vec<MicroOp> {
397        let mut ops = vec![MicroOp::Branch { cmp, lhs: 0, rhs: 1, target: 0 }];
398        ops.extend(body);
399        ops.push(MicroOp::LoadConst { dst: 90, value: step_val });
400        ops.push(MicroOp::Add { dst: 0, lhs: 0, rhs: 90 });
401        let jump_idx = ops.len();
402        ops.push(MicroOp::Jump { target: 0 });
403        ops.push(MicroOp::Return { src: 0 });
404        if let MicroOp::Branch { target, .. } = &mut ops[0] {
405            *target = jump_idx + 1; // exit -> the Return
406        }
407        ops
408    }
409
410    fn map_add_region() -> Vec<MicroOp> {
411        loop_region(
412            vec![
413                load(8, 0, 2, 5),
414                load(9, 0, 3, 6),
415                MicroOp::AddF { dst: 10, lhs: 8, rhs: 9 },
416                store(10, 0, 4, 7),
417            ],
418            Cmp::LtEq,
419            1,
420        )
421    }
422
423    #[test]
424    fn recognizes_elementwise_add_map() {
425        let ops = map_add_region();
426        let plan = recognize_elementwise_map(&ops).expect("c[i]=a[i]+b[i] is a map");
427        assert_eq!(plan.induction, 0);
428        assert_eq!(plan.limit, 1);
429        assert_eq!(plan.cmp, Cmp::LtEq);
430        // body = the 4 load/arith/store ops at [1..5] (guard at 0, control after).
431        assert_eq!((plan.body_start, plan.body_end), (1, 5));
432    }
433
434    #[test]
435    fn recognizes_multiop_map() {
436        // c[i] = a[i]*b[i] + a[i] — still a pure map (all reads at i, no carry).
437        let ops = loop_region(
438            vec![
439                load(8, 0, 2, 5),
440                load(9, 0, 3, 6),
441                MicroOp::MulF { dst: 10, lhs: 8, rhs: 9 },
442                MicroOp::AddF { dst: 10, lhs: 10, rhs: 8 },
443                store(10, 0, 4, 7),
444            ],
445            Cmp::Lt,
446            1,
447        );
448        assert!(recognize_elementwise_map(&ops).is_some());
449    }
450
451    #[test]
452    fn rejects_reduction_loop_carried_accumulator() {
453        // `sum` (slot 9) is read before being defined in the body — loop-carried,
454        // NOT bit-exactly vectorizable. Must reject.
455        let ops = loop_region(
456            vec![
457                load(8, 0, 2, 5),
458                MicroOp::AddF { dst: 9, lhs: 9, rhs: 8 }, // slot 9 loop-carried
459                store(9, 0, 4, 7),
460            ],
461            Cmp::Lt,
462            1,
463        );
464        assert_eq!(recognize_elementwise_map(&ops), None);
465    }
466
467    #[test]
468    fn rejects_cross_index_access_scan() {
469        // out[i] = a[j] where j != i (a scan/stencil read at a different index).
470        let ops = loop_region(
471            vec![load(8, 9, 2, 5), store(8, 0, 4, 7)], // idx=9 != i(0)
472            Cmp::Lt,
473            1,
474        );
475        assert_eq!(recognize_elementwise_map(&ops), None);
476    }
477
478    #[test]
479    fn rejects_body_with_load_only_no_store() {
480        // A body that never stores is not a map (e.g. it feeds a reduction).
481        let ops = loop_region(vec![load(8, 0, 2, 5)], Cmp::Lt, 1);
482        assert_eq!(recognize_elementwise_map(&ops), None);
483    }
484
485    #[test]
486    fn rejects_non_unit_stride() {
487        // i += 2 — the kernel visits every index, so a stride-2 loop (a subset)
488        // must be rejected (the step LoadConst is not 1).
489        let ops = loop_region(vec![load(8, 0, 2, 5), store(8, 0, 4, 7)], Cmp::Lt, 2);
490        assert_eq!(recognize_elementwise_map(&ops), None);
491    }
492
493    #[test]
494    fn rejects_bounds_checked_access() {
495        // A `checked: true` access — the Oracle did NOT prove length >= n, so the
496        // kernel (which elides per-element checks) would be memory-unsafe. Reject.
497        let ops = loop_region(vec![load_checked(8, 0, 2, 5), store(8, 0, 4, 7)], Cmp::Lt, 1);
498        assert_eq!(recognize_elementwise_map(&ops), None);
499    }
500
501    #[test]
502    fn rejects_non_back_edge_region() {
503        // No back-edge `Jump { target: 0 }` -> not a recognized loop region.
504        let ops = vec![
505            MicroOp::Branch { cmp: Cmp::Lt, lhs: 0, rhs: 1, target: 5 },
506            load(8, 0, 2, 5),
507            store(8, 0, 4, 7),
508            MicroOp::LoadConst { dst: 90, value: 1 },
509            MicroOp::Add { dst: 0, lhs: 0, rhs: 90 },
510        ];
511        assert_eq!(recognize_elementwise_map(&ops), None);
512    }
513
514    #[cfg(target_arch = "x86_64")]
515    fn run_frame(code: &[u8], frame: &mut [i64]) -> i64 {
516        let page = crate::JitPage::new(code).unwrap();
517        let f: extern "C" fn(*mut i64) -> i64 =
518            unsafe { std::mem::transmute(page.as_ptr()) };
519        f(frame.as_mut_ptr())
520    }
521
522    #[cfg(target_arch = "x86_64")]
523    #[test]
524    fn packed_body_lowering_matches_scalar_both_lanes() {
525        use crate::x64asm::{Asm, Reg, Xmm};
526        // Lane assignment for body slots: ta=xmm0, tb=xmm1, tc=xmm2.
527        let lane = |s: Slot| match s {
528            0 => Xmm::Xmm0,
529            1 => Xmm::Xmm1,
530            2 => Xmm::Xmm2,
531            _ => unreachable!(),
532        };
533        // Body: tc = ta * tb ; tc = tc + ta  (a pure map expression, 2 ops).
534        let body = [
535            MicroOp::MulF { dst: 2, lhs: 0, rhs: 1 },
536            MicroOp::AddF { dst: 2, lhs: 2, rhs: 0 },
537        ];
538        let cases: [(f64, f64, f64, f64); 3] = [
539            (2.5, 0.5, 1.0, 3.0),
540            (-1.0, 7.25, 4.0, -2.5),
541            (0.1, 0.2, 1.0 / 3.0, 9.0),
542        ];
543        for (a0, a1, b0, b1) in cases {
544            let mut asm = Asm::new();
545            asm.movupd_rm(Xmm::Xmm0, Reg::Rdi, 0); // ta = {a0,a1}
546            asm.movupd_rm(Xmm::Xmm1, Reg::Rdi, 16); // tb = {b0,b1}
547            for op in &body {
548                emit_packed_arith(&mut asm, op, lane, Xmm::Xmm15);
549            }
550            asm.movupd_mr(Reg::Rdi, 32, Xmm::Xmm2); // store tc
551            asm.ret();
552            let mut frame = [
553                a0.to_bits() as i64, a1.to_bits() as i64,
554                b0.to_bits() as i64, b1.to_bits() as i64,
555                0, 0,
556            ];
557            run_frame(&asm.resolve(), &mut frame);
558            assert_eq!(f64::from_bits(frame[4] as u64).to_bits(), (a0 * b0 + a0).to_bits(), "lane0");
559            assert_eq!(f64::from_bits(frame[5] as u64).to_bits(), (a1 * b1 + a1).to_bits(), "lane1");
560        }
561    }
562
563    #[cfg(target_arch = "x86_64")]
564    fn run_map(body: &[MicroOp], plan: &MapPlan, a: &[f64], b: &[f64]) -> Vec<f64> {
565        // 1-based loop `for i in 1..=n` (LtEq), accessing buffer[i-1] — the shape
566        // the recognizer fires on. i starts at 1; n = the length.
567        let code = emit_map_kernel(body, plan).expect("kernel emits");
568        let mut c = vec![0.0f64; a.len()];
569        let mut frame = [
570            1i64,
571            a.len() as i64,
572            a.as_ptr() as i64,
573            b.as_ptr() as i64,
574            c.as_mut_ptr() as i64,
575            0, 0, 0,
576        ];
577        run_frame(&code, &mut frame);
578        // The kernel leaves i at the scalar terminal value: n+1 for `i <= n`.
579        assert_eq!(frame[0], a.len() as i64 + 1, "induction terminal n+1");
580        c
581    }
582
583    #[cfg(target_arch = "x86_64")]
584    #[test]
585    fn map_kernel_add_matches_scalar_all_lengths() {
586        // c[i] = a[i] + b[i] over various n (even, odd, 1, 0) — exercises the pair
587        // loop AND the scalar tail.
588        let body = vec![
589            load(5, 0, 2, 0),
590            load(6, 0, 3, 0),
591            MicroOp::AddF { dst: 7, lhs: 5, rhs: 6 },
592            store(7, 0, 4, 0),
593        ];
594        let plan = MapPlan { induction: 0, limit: 1, cmp: Cmp::LtEq, body_start: 0, body_end: 4 };
595        for n in [0usize, 1, 2, 3, 7, 8, 15] {
596            let a: Vec<f64> = (0..n).map(|i| i as f64 * 1.5 - 3.0).collect();
597            let b: Vec<f64> = (0..n).map(|i| i as f64 * -0.25 + 1.0).collect();
598            let got = run_map(&body, &plan, &a, &b);
599            for i in 0..n {
600                assert_eq!(got[i].to_bits(), (a[i] + b[i]).to_bits(), "n={n} i={i}");
601            }
602        }
603    }
604
605    #[cfg(target_arch = "x86_64")]
606    #[test]
607    fn map_kernel_multiop_matches_scalar() {
608        // c[i] = a[i]*b[i] + a[i] — multi-op map, pair loop + tail.
609        let body = vec![
610            load(5, 0, 2, 0),
611            load(6, 0, 3, 0),
612            MicroOp::MulF { dst: 7, lhs: 5, rhs: 6 },
613            MicroOp::AddF { dst: 7, lhs: 7, rhs: 5 },
614            store(7, 0, 4, 0),
615        ];
616        let plan = MapPlan { induction: 0, limit: 1, cmp: Cmp::LtEq, body_start: 0, body_end: 5 };
617        for n in [1usize, 4, 5, 16, 17] {
618            let a: Vec<f64> = (0..n).map(|i| 0.1 * i as f64 + 0.3).collect();
619            let b: Vec<f64> = (0..n).map(|i| 2.0 - 0.07 * i as f64).collect();
620            let got = run_map(&body, &plan, &a, &b);
621            for i in 0..n {
622                assert_eq!(got[i].to_bits(), (a[i] * b[i] + a[i]).to_bits(), "n={n} i={i}");
623            }
624        }
625    }
626
627    #[cfg(target_arch = "x86_64")]
628    #[test]
629    fn map_kernel_bails_on_unsupported_loadconst() {
630        // A body with a LoadConst is not lowered by the kernel -> None (the caller
631        // falls back to the scalar region, staying sound).
632        let body = vec![
633            load(5, 0, 2, 0),
634            MicroOp::LoadConst { dst: 6, value: 2 },
635            MicroOp::MulF { dst: 7, lhs: 5, rhs: 6 },
636            store(7, 0, 4, 0),
637        ];
638        let plan = MapPlan { induction: 0, limit: 1, cmp: Cmp::LtEq, body_start: 0, body_end: 4 };
639        assert!(emit_map_kernel(&body, &plan).is_none());
640    }
641
642    #[cfg(target_arch = "x86_64")]
643    #[test]
644    fn packed_body_lowering_non_commutative_dst_eq_rhs_uses_scratch() {
645        use crate::x64asm::{Asm, Reg, Xmm};
646        // SubF dst==rhs: td(xmm1) = ta(xmm0) - td(xmm1). The 2-operand subpd would
647        // compute xmm1 = xmm1 - xmm0 (wrong order); emit_bin must route through the
648        // scratch to preserve `lhs - rhs`.
649        let lane = |s: Slot| match s {
650            0 => Xmm::Xmm0,
651            1 => Xmm::Xmm1,
652            _ => unreachable!(),
653        };
654        let op = MicroOp::SubF { dst: 1, lhs: 0, rhs: 1 };
655        let cases: [(f64, f64, f64, f64); 2] = [(2.5, 0.5, 1.0, 3.0), (-1.0, 3.0, 4.0, -2.5)];
656        for (a0, a1, b0, b1) in cases {
657            let mut asm = Asm::new();
658            asm.movupd_rm(Xmm::Xmm0, Reg::Rdi, 0); // ta
659            asm.movupd_rm(Xmm::Xmm1, Reg::Rdi, 16); // td (the rhs, gets overwritten)
660            emit_packed_arith(&mut asm, &op, lane, Xmm::Xmm15);
661            asm.movupd_mr(Reg::Rdi, 32, Xmm::Xmm1);
662            asm.ret();
663            let mut frame = [
664                a0.to_bits() as i64, a1.to_bits() as i64,
665                b0.to_bits() as i64, b1.to_bits() as i64,
666                0, 0,
667            ];
668            run_frame(&asm.resolve(), &mut frame);
669            assert_eq!(f64::from_bits(frame[4] as u64).to_bits(), (a0 - b0).to_bits(), "lane0 a-b");
670            assert_eq!(f64::from_bits(frame[5] as u64).to_bits(), (a1 - b1).to_bits(), "lane1 a-b");
671        }
672    }
673}
logicaffeine_forge/vectorize.rs

logicaffeine_forge/
vectorize.rs
