earthsci-ast 0.6.0

//! Step 4: kernel fusion on the tape.
//!
//! A build-time post-pass over the lowered (pre-coloring) [`TapeProgram`] that
//! merges DAG segments of same-box elementwise instructions
//! (`Bin`/`Un`/`Neg`/`Select`/`Fill`/`Copy`) into [`Instr::Fused`] groups: one
//! straight-line micro-program executed per element in a single pass over the
//! box, with only the group's live-out values materialized to the slab. A
//! [`Instr::Gather`] whose plan is a pure constant shift/wrap of an
//! equal-shape source (no fixed axes, no broadcast axes, no transposition) is
//! *folded into* its consumers as a shifted read: the group's precompiled run
//! schedule partitions the box into contiguous runs on which every folded
//! input is either one constant flat source offset or entirely the gather's
//! Dirichlet ghost (`+0.0`) — so the per-element value is byte-identical to
//! reading the materialized gather output.
//!
//! ## Bit identity by construction
//!
//! Per element, a fused group applies exactly the same scalar kernels in the
//! same order as the instructions it replaces; all fused ops are elementwise
//! maps (no reductions, scans, or cross-element reads other than the
//! precompiled constant-shift gathers, whose sources are always values
//! defined OUTSIDE the group), so element iteration order cannot change any
//! output bit. `Select` keeps evaluate-both semantics (both operands are
//! computed per element, then picked), exactly like the vectorized
//! `vec_select`.
//!
//! ## Soundness of the reordering
//!
//! The pass buffers group members while passing other instructions through,
//! so a group's single `Fused` instruction executes at the position of the
//! member that *forced* the flush. This is sound because:
//!
//! * slots are SSA (defined once; the phi double-definition only occurs
//!   inside `JmpIfZero` regions, which are hard barriers copied verbatim);
//! * any instruction that READS a slot defined by an open group flushes that
//!   group first, so no reader ever precedes the `Fused` that defines it;
//! * `JmpIfZero` regions flush every open group, so their exact skip counts
//!   are preserved;
//! * passed-through instructions keep their original relative order.
//!
//! Coloring runs AFTER fusion, so liveness (and the alias discipline — a
//! `Fused` is treated like a `Gather`: its outputs never recycle a source
//! buffer dying at the same instruction) is computed on the fused program,
//! and the internal values that no longer materialize shrink the slab.

use super::super::BinCode;
use super::ir::*;
use crate::precision::Precision;
use rustc_hash::{FxHashMap, FxHashSet};
use smallvec::SmallVec;
use std::collections::BTreeSet;

/// A [`ChunkGather`] with rows shorter than this is read whole
/// ([`ChunkGather::whole`]).
const MIN_CHUNK_GATHER_ROW: usize = 12;

/// So is one with rows shorter than this when its group has fewer than
/// [`LIGHT_GROUP_OPS`] micro-ops.
const SHORT_CHUNK_GATHER_ROW: usize = 32;
const LIGHT_GROUP_OPS: usize = 8;

/// Maximum simultaneously open groups (oldest is flushed beyond this).
const MAX_OPEN_GROUPS: usize = 4;

/// Row-major strides (elements) of `shape`.
fn rm_strides(shape: &[usize]) -> SmallVec<[i64; 4]> {
    let n = shape.len();
    let mut st = SmallVec::from_elem(0i64, n);
    let mut acc = 1i64;
    for d in (0..n).rev() {
        st[d] = acc;
        acc *= shape[d] as i64;
    }
    st
}

// ---------------------------------------------------------------------------
// Operand / source keys for deduplication.
// ---------------------------------------------------------------------------

#[derive(Hash, PartialEq, Eq, Clone, Copy)]
enum ScalKey {
    Lit(u64),
    Param(u32),
    Time,
    Slot(SlotId),
    State(u32),
}

/// Dedupe key for ALIGNED array inputs (observed operands are never fused,
/// so no `Obs` arm is needed).
#[derive(Hash, PartialEq, Eq, Clone, Copy)]
enum SrcKey {
    Slot(SlotId),
    State(u32),
}

// ---------------------------------------------------------------------------
// Group builder.
// ---------------------------------------------------------------------------

type AxisSegs = SmallVec<[SmallVec<[(usize, usize, usize); 2]>; 4]>;

/// Geometry of one folded (shifted-read) gather input.
// `Segs` is deliberately far larger than `Linear`: the SmallVec-of-SmallVec is
// inline storage, chosen so the common segment counts never touch the heap on
// this hot fold path. Boxing it to even the variants out -- clippy's suggested
// fix -- would reintroduce exactly the indirection that layout is avoiding.
#[allow(clippy::large_enum_variant)]
#[derive(Clone, PartialEq)]
enum ShiftedGeom {
    /// Piecewise per-axis segments with ghost gaps: along output axis `d` the
    /// source advances by `strides[d]` per element (0 for an axis the source
    /// is broadcast along), from `base` (the fixed axes' offset). Within a run
    /// the source advances by the innermost axis's stride.
    Segs {
        segs: AxisSegs,
        strides: SmallVec<[i64; 4]>,
        base: i64,
    },
    /// Globally linear: `flat_src = a * flat_out + b` over the whole box
    /// (e.g. a level slice of a deeper box). Single run, element stride `a`.
    Linear { a: i64, b: i64 },
}

struct GBuilder {
    shape: super::super::DimU,
    /// Original instruction indices absorbed (in original order).
    member_instrs: Vec<u32>,
    members: FxHashSet<u32>,
    /// Group-defined slot -> value reference (Reg for computed members,
    /// In for folded gathers).
    val_of: FxHashMap<SlotId, MRef>,
    micro: Vec<MicroOp>,
    inputs: Vec<FusedInput>,
    /// Per SHIFTED input: its fold geometry.
    shifted_segs: Vec<ShiftedGeom>,
    aligned_ix: FxHashMap<SrcKey, GroupIx>,
    scalars: Vec<Operand>,
    scalar_ix: FxHashMap<ScalKey, GroupIx>,
    /// Computed member slots in definition order: `(slot, ssa register)`.
    defs: Vec<(SlotId, GroupIx)>,
    /// Folded gathers: `(original instr index, out slot, input index)`.
    folded: Vec<(u32, SlotId, GroupIx)>,
    /// Folded data-subscript gathers, the same way.
    folded_index: Vec<(u32, SlotId, GroupIx)>,
    /// Rule ordinal of the first member (provenance of the emitted Fused).
    prov: u32,
    /// The precision every member runs at (`None` when the program carries
    /// no per-instruction precision). A group never spans two: fusion stops
    /// at a change of precision.
    prec: Option<Precision>,
    /// An absorbed [`Instr::Reduce`] folding one of the group's values
    /// (see [`FusedReduce`]); the group is flushed as soon as it is set.
    reduce: Option<ReduceTail>,
    /// Value numbering over the group's pure micro-ops: an op with the same
    /// kernel and the same operands as an earlier one is that op's register
    /// (the kernels are pure, so it would compute the same bits).
    vn: FxHashMap<MKey, GroupIx>,
    /// Carry slots the absorbed scans use so far (each scan owns one per
    /// lane, the next ones).
    n_carries: u32,
    /// `Export`s and `DyWrite`s of slots this group defines, held back until
    /// it flushes (see `fuse_section`).
    deferred: Vec<u32>,
}

/// An [`Instr::Reduce`] a group absorbed, before register allocation.
struct ReduceTail {
    /// The folded value's SSA register.
    ssa: GroupIx,
    op: BinCode,
    init: f64,
    out: SlotId,
    n_inner: usize,
}

/// A pure micro-op without its destination: the value-numbering key.
#[derive(Hash, PartialEq, Eq)]
enum MKey {
    Bin(BinCode, MRef, MRef),
    Un(super::super::UnCode, MRef),
    Neg(MRef),
    Select(MRef, MRef, MRef),
    Mov(MRef),
}

/// The same source (`SrcRef` has no equality of its own).
fn same_src(a: &SrcRef, b: &SrcRef) -> bool {
    match (a, b) {
        (SrcRef::Slot(x), SrcRef::Slot(y)) => x == y,
        (SrcRef::State(x), SrcRef::State(y)) => x == y,
        _ => false,
    }
}

/// Operand resolved for group membership (pre-commit description).
enum ResOp {
    Existing(MRef),
    NewScalar(ScalKey, Operand),
    NewAligned(SrcKey, SrcRef),
}

impl GBuilder {
    fn new(shape: super::super::DimU, prov: u32, prec: Option<Precision>) -> Self {
        GBuilder {
            shape,
            member_instrs: Vec::new(),
            members: FxHashSet::default(),
            val_of: FxHashMap::default(),
            micro: Vec::new(),
            inputs: Vec::new(),
            shifted_segs: Vec::new(),
            aligned_ix: FxHashMap::default(),
            scalars: Vec::new(),
            scalar_ix: FxHashMap::default(),
            defs: Vec::new(),
            folded: Vec::new(),
            folded_index: Vec::new(),
            prov,
            prec,
            reduce: None,
            vn: FxHashMap::default(),
            n_carries: 0,
            deferred: Vec::new(),
        }
    }

    fn defines(&self, s: SlotId) -> bool {
        self.val_of.contains_key(&s) || self.reduce.as_ref().is_some_and(|r| r.out == s)
    }

    /// Absorb `Reduce { src, axes, .. }` at instruction `ix` when it folds the
    /// LEADING axes of this group's box and reads a value only this group
    /// computes and nothing else reads — so the value never has to be stored.
    /// Returns `false` (group untouched) otherwise.
    fn try_absorb_reduce(&mut self, ix: u32, ins: &Instr, readers: &[SmallVec<[u32; 4]>]) -> bool {
        let Instr::Reduce {
            op,
            init,
            src: SrcRef::Slot(src),
            axes,
            reverse,
            src_shape,
            out,
        } = ins
        else {
            return false;
        };
        // A group's chunks fold ascending; a reverse fold stays standalone,
        // and a group that visits its rows from the last folds nothing.
        if *reverse
            || self
                .micro
                .iter()
                .any(|m| matches!(m, MicroOp::Scan { reverse: true, .. }))
            || self.reduce.is_some()
            || *src_shape != self.shape
            || axes.is_empty()
            || axes.iter().enumerate().any(|(k, &a)| a as usize != k)
            || readers[*src as usize].iter().any(|&r| r != ix)
        {
            return false;
        }
        let Some(&MRef::Reg(ssa)) = self.val_of.get(src) else {
            return false;
        };
        self.reduce = Some(ReduceTail {
            ssa,
            op: *op,
            init: *init,
            out: *out,
            n_inner: src_shape[axes.len()..].iter().product::<usize>().max(1),
        });
        self.members.insert(ix);
        self.member_instrs.push(ix);
        true
    }

    /// Absorb the identity gather at instruction `ix` (`out` reads `src`,
    /// which this group computes, element for element) as another name for
    /// `src`'s register.
    fn alias(&mut self, ix: u32, src: SlotId, out: SlotId) {
        let Some(&MRef::Reg(ssa)) = self.val_of.get(&src) else {
            unreachable!("an alias names a computed value")
        };
        self.val_of.insert(out, MRef::Reg(ssa));
        self.defs.push((out, ssa));
        self.members.insert(ix);
        self.member_instrs.push(ix);
    }

    /// Absorb `Scan` at instruction `ix` when it runs along an axis of this
    /// group's box: a [`MicroOp::Scan`] over its source, which the group may
    /// compute itself or read as an aligned input. Returns `false` (group
    /// untouched) otherwise.
    fn try_add_scan(&mut self, ix: u32, ins: &Instr, prog: &TapeProgram) -> bool {
        let Instr::Scan {
            op,
            init,
            src,
            axis,
            inclusive,
            reverse,
            src_shape,
            out,
        } = ins
        else {
            return false;
        };
        if *src_shape != self.shape || *axis as usize >= self.shape.len() {
            return false;
        }
        // A reverse scan joins a group only along its leading axis, whose
        // rows the group then visits from the last (`FusedSpec::rows_reversed`):
        // never beside a forward scan or a reduction, which fold ascending,
        // with lanes to share a row and not too many rows.
        let scans = |rev: bool| {
            self.micro
                .iter()
                .any(|m| matches!(m, MicroOp::Scan { reverse, .. } if *reverse == rev))
        };
        if *reverse {
            let lanes: usize = self.shape[1..].iter().product();
            if *axis != 0
                || self.reduce.is_some()
                || scans(false)
                || lanes < 2
                || self.shape[0] > MAX_REVERSE_ROWS
            {
                return false;
            }
        } else if scans(true) {
            return false;
        }
        let operand = Operand::from(*src);
        let Some(r @ (ResOp::Existing(MRef::Reg(_) | MRef::In(_)) | ResOp::NewAligned(..))) =
            self.resolve(&operand, prog)
        else {
            return false;
        };
        if let ResOp::Existing(MRef::In(i)) = r
            && (self.inputs[i as usize].shifted_ix.is_some()
                || self.inputs[i as usize].index.is_some())
        {
            return false;
        }
        let Ok(post) = u32::try_from(self.shape[*axis as usize + 1..].iter().product::<usize>())
        else {
            return false;
        };
        if self.n_carries.checked_add(post).is_none() {
            return false;
        }
        let a = self.commit(r);
        let ssa = self.micro.len() as GroupIx;
        self.micro.push(MicroOp::Scan {
            op: *op,
            a,
            init: *init,
            inclusive: *inclusive,
            reverse: *reverse,
            row: self.shape[*axis as usize] as u32,
            post,
            carry: self.n_carries,
            out: ssa,
        });
        self.n_carries += post;
        self.val_of.insert(*out, MRef::Reg(ssa));
        self.defs.push((*out, ssa));
        self.members.insert(ix);
        self.member_instrs.push(ix);
        true
    }

    /// Whether one more member keeps every group-local index inside
    /// [`GroupIx`]. The register file (`n_regs + n_load_regs +
    /// n_splat_regs`) is bounded by `micro + inputs + scalars + 1`, one
    /// member adds at most one micro-op, three inputs and three scalars, and
    /// `GroupIx::MAX` itself is the `load_reg` sentinel. A full group is
    /// flushed and a fresh one opened, which only moves where the slab is
    /// written.
    fn has_room(&self) -> bool {
        const PER_MEMBER: usize = 8;
        self.micro.len() + self.inputs.len() + self.scalars.len() + PER_MEMBER
            < GroupIx::MAX as usize
    }

    /// Resolve an operand for use inside the group; `None` = unfusable
    /// operand (an observed read, whose rank is not statically known here, or
    /// a foreign-box array).
    fn resolve(&self, op: &Operand, prog: &TapeProgram) -> Option<ResOp> {
        match op {
            Operand::Lit(v) => Some(self.scal(ScalKey::Lit(v.to_bits()), *op)),
            Operand::Param(p) => Some(self.scal(ScalKey::Param(*p), *op)),
            Operand::Time => Some(self.scal(ScalKey::Time, *op)),
            Operand::Slot(s) => {
                if let Some(m) = self.val_of.get(s) {
                    return Some(ResOp::Existing(*m));
                }
                let d = &prog.slots[*s as usize];
                if d.scalar {
                    Some(self.scal(ScalKey::Slot(*s), *op))
                } else if d.shape == self.shape {
                    Some(self.aligned(SrcKey::Slot(*s), SrcRef::Slot(*s)))
                } else {
                    None
                }
            }
            Operand::State(ix) => {
                let sv = &prog.state_vars[*ix as usize];
                if sv.shape.is_empty() {
                    Some(self.scal(ScalKey::State(*ix), *op))
                } else if sv.shape == self.shape {
                    Some(self.aligned(SrcKey::State(*ix), SrcRef::State(*ix)))
                } else {
                    None
                }
            }
        }
    }

    fn scal(&self, key: ScalKey, op: Operand) -> ResOp {
        match self.scalar_ix.get(&key) {
            Some(&i) => ResOp::Existing(MRef::Scal(i)),
            None => ResOp::NewScalar(key, op),
        }
    }

    fn aligned(&self, key: SrcKey, src: SrcRef) -> ResOp {
        match self.aligned_ix.get(&key) {
            Some(&i) => ResOp::Existing(MRef::In(i)),
            None => ResOp::NewAligned(key, src),
        }
    }

    fn commit(&mut self, r: ResOp) -> MRef {
        match r {
            ResOp::Existing(m) => m,
            ResOp::NewScalar(key, op) => {
                let i = self.scalars.len() as GroupIx;
                self.scalars.push(op);
                self.scalar_ix.insert(key, i);
                MRef::Scal(i)
            }
            ResOp::NewAligned(key, src) => {
                let i = self.inputs.len() as GroupIx;
                self.inputs.push(FusedInput {
                    src,
                    shifted_ix: None,
                    src_shape: self.shape.clone(),
                    elem_stride: 1,
                    load_reg: GroupIx::MAX,
                    index: None,
                    gather: None,
                });
                self.aligned_ix.insert(key, i);
                MRef::In(i)
            }
        }
    }

    /// Try to absorb an elementwise instruction whose out slot has this
    /// group's shape. Returns `false` (group untouched) if an operand is
    /// unfusable.
    fn try_add(&mut self, ix: u32, ins: &Instr, prog: &TapeProgram) -> bool {
        let (ops, out): (SmallVec<[&Operand; 3]>, SlotId) = match ins {
            Instr::Bin { a, b, out, .. } => (SmallVec::from_slice(&[a, b]), *out),
            Instr::Un { a, out, .. } | Instr::Neg { a, out } => (SmallVec::from_slice(&[a]), *out),
            Instr::Select { cond, a, b, out } => (SmallVec::from_slice(&[cond, a, b]), *out),
            Instr::Fill { v, out } => (SmallVec::from_slice(&[v]), *out),
            Instr::Copy { a, out } => (SmallVec::from_slice(&[a]), *out),
            _ => return false,
        };
        let mut resolved: SmallVec<[ResOp; 3]> = SmallVec::new();
        for op in &ops {
            match self.resolve(op, prog) {
                Some(r) => resolved.push(r),
                None => return false,
            }
        }
        let mut mrefs: SmallVec<[MRef; 3]> = SmallVec::new();
        for r in resolved {
            mrefs.push(self.commit(r));
        }
        let key = match ins {
            Instr::Bin { op, .. } => MKey::Bin(*op, mrefs[0], mrefs[1]),
            Instr::Un { op, .. } => MKey::Un(*op, mrefs[0]),
            Instr::Neg { .. } => MKey::Neg(mrefs[0]),
            Instr::Select { .. } => MKey::Select(mrefs[0], mrefs[1], mrefs[2]),
            Instr::Fill { .. } | Instr::Copy { .. } => MKey::Mov(mrefs[0]),
            _ => unreachable!(),
        };
        if let Some(&ssa) = self.vn.get(&key) {
            self.val_of.insert(out, MRef::Reg(ssa));
            self.defs.push((out, ssa));
            self.members.insert(ix);
            self.member_instrs.push(ix);
            return true;
        }
        let ssa = self.micro.len() as GroupIx;
        self.vn.insert(key, ssa);
        let mop = match ins {
            Instr::Bin { op, .. } => MicroOp::Bin {
                op: *op,
                a: mrefs[0],
                b: mrefs[1],
                out: ssa,
            },
            Instr::Un { op, .. } => MicroOp::Un {
                op: *op,
                a: mrefs[0],
                out: ssa,
            },
            Instr::Neg { .. } => MicroOp::Neg {
                a: mrefs[0],
                out: ssa,
            },
            Instr::Select { .. } => MicroOp::Select {
                cond: mrefs[0],
                a: mrefs[1],
                b: mrefs[2],
                out: ssa,
            },
            Instr::Fill { .. } | Instr::Copy { .. } => MicroOp::Mov {
                a: mrefs[0],
                out: ssa,
            },
            _ => unreachable!(),
        };
        self.micro.push(mop);
        self.val_of.insert(out, MRef::Reg(ssa));
        self.defs.push((out, ssa));
        self.members.insert(ix);
        self.member_instrs.push(ix);
        true
    }

    /// Absorb a rank-1 data-subscript gather whose subscript is `idx`, an
    /// array operand of the group's box defined outside it, as a pre-loaded
    /// input. `false` (group untouched) when the subscript cannot be an
    /// aligned input.
    #[allow(clippy::too_many_arguments)]
    fn add_folded_index_gather(
        &mut self,
        ix: u32,
        src: SrcRef,
        idx: &Operand,
        fold: FoldedIndex,
        src_shape: &super::super::DimU,
        out: SlotId,
        prog: &TapeProgram,
    ) -> bool {
        let resolved = match idx {
            Operand::Slot(_) | Operand::State(_) => self.resolve(idx, prog),
            _ => None,
        };
        // A subscript the group computes itself: read through a `Load`
        // micro-op at its place in the program.
        if let Some(ResOp::Existing(MRef::Reg(r))) = resolved {
            let input_ix = self.inputs.len() as GroupIx;
            self.inputs.push(FusedInput {
                src,
                shifted_ix: None,
                src_shape: src_shape.clone(),
                elem_stride: 1,
                load_reg: GroupIx::MAX,
                index: Some(FoldedIndex {
                    by: FoldedIndex::IN_GROUP,
                    ..fold
                }),
                gather: None,
            });
            let ssa = self.micro.len() as GroupIx;
            self.micro.push(MicroOp::Load {
                a: MRef::Reg(r),
                src: input_ix,
                out: ssa,
            });
            // An ordinary register from here on: stored like any other
            // value of the group when something outside reads it.
            self.val_of.insert(out, MRef::Reg(ssa));
            self.defs.push((out, ssa));
            self.members.insert(ix);
            self.member_instrs.push(ix);
            return true;
        }
        let by = match resolved {
            Some(r @ ResOp::NewAligned(..)) => match self.commit(r) {
                MRef::In(i) => i,
                _ => unreachable!("an aligned input commits to an input"),
            },
            Some(ResOp::Existing(MRef::In(i)))
                if self.inputs[i as usize].shifted_ix.is_none()
                    && self.inputs[i as usize].index.is_none() =>
            {
                i
            }
            _ => return false,
        };
        let input_ix = self.inputs.len() as GroupIx;
        self.inputs.push(FusedInput {
            src,
            shifted_ix: None,
            src_shape: src_shape.clone(),
            elem_stride: 1,
            load_reg: GroupIx::MAX,
            index: Some(FoldedIndex { by, ..fold }),
            gather: None,
        });
        self.val_of.insert(out, MRef::In(input_ix));
        self.folded_index.push((ix, out, input_ix));
        self.members.insert(ix);
        self.member_instrs.push(ix);
        true
    }

    /// Absorb a gather read through its plan one chunk at a time (see
    /// [`ChunkGather`]) — the input an earlier identical one made, if any.
    fn add_chunk_gather(
        &mut self,
        ix: u32,
        src: SrcRef,
        g: ChunkGather,
        src_shape: &super::super::DimU,
        out: SlotId,
    ) {
        let twin = self.inputs.iter().position(|inp| {
            same_src(&inp.src, &src) && inp.gather.as_deref().is_some_and(|h| *h == g)
        });
        let input_ix = match twin {
            Some(t) => t as GroupIx,
            None => {
                self.inputs.push(FusedInput {
                    src,
                    shifted_ix: None,
                    src_shape: src_shape.clone(),
                    elem_stride: 1,
                    load_reg: GroupIx::MAX,
                    index: None,
                    gather: Some(Box::new(g)),
                });
                (self.inputs.len() - 1) as GroupIx
            }
        };
        self.val_of.insert(out, MRef::In(input_ix));
        self.folded_index.push((ix, out, input_ix));
        self.members.insert(ix);
        self.member_instrs.push(ix);
    }

    /// Absorb a foldable gather as a shifted input — the input an earlier
    /// fold of the same source with the same geometry already made, if any.
    fn add_folded_gather(
        &mut self,
        ix: u32,
        src: SrcRef,
        plan: &GatherPlan,
        geom: ShiftedGeom,
        out: SlotId,
    ) {
        let twin = self.inputs.iter().position(|inp| {
            inp.shifted_ix
                .is_some_and(|s| same_src(&inp.src, &src) && self.shifted_segs[s as usize] == geom)
        });
        if let Some(t) = twin {
            let input_ix = t as GroupIx;
            self.val_of.insert(out, MRef::In(input_ix));
            self.folded.push((ix, out, input_ix));
            self.members.insert(ix);
            self.member_instrs.push(ix);
            return;
        }
        let input_ix = self.inputs.len() as GroupIx;
        let shifted_ix = self.shifted_segs.len() as GroupIx;
        let elem_stride = geom.elem_stride();
        self.inputs.push(FusedInput {
            src,
            shifted_ix: Some(shifted_ix),
            src_shape: plan.src_shape.clone(),
            elem_stride,
            load_reg: GroupIx::MAX,
            index: None,
            gather: None,
        });
        self.shifted_segs.push(geom);
        self.val_of.insert(out, MRef::In(input_ix));
        self.folded.push((ix, out, input_ix));
        self.members.insert(ix);
        self.member_instrs.push(ix);
    }
}

/// How fusion absorbs a gather into the group of its output box.
// Returned and consumed at once; boxing the geometry would only add an
// allocation per gather.
#[allow(clippy::large_enum_variant)]
enum GatherFold {
    /// A shifted input read along precompiled runs.
    Shifted(ShiftedGeom),
    /// A same-rank shift whose runs would shatter the box, read through its
    /// plan inside the group chunk by chunk or (for short rows) the whole box
    /// at once, so the program does not depend on the box's extents. Holds
    /// the source stride of each output axis: a broadcast (an output axis the
    /// source lacks) reads that axis with stride 0, the other axes in order.
    Chunk(SmallVec<[i64; 4]>),
}

/// How fusion absorbs a gather of `plan` into a box of `out_shape`, or
/// `None` when the gather passes through materialized.
fn gather_fold(plan: &GatherPlan, out_shape: &[usize]) -> Option<GatherFold> {
    if out_shape != &plan.shape[..] {
        return None;
    }
    let shifted = linear_fold(plan).or_else(|| {
        let g = if foldable_plan(plan) {
            ShiftedGeom::Segs {
                segs: plan.segs.clone(),
                strides: rm_strides(&plan.src_shape),
                base: 0,
            }
        } else {
            broadcast_fold(plan)?
        };
        fold_runs_acceptable(out_shape, &g).then_some(g)
    });
    if let Some(g) = shifted {
        return Some(GatherFold::Shifted(g));
    }
    if foldable_plan(plan) {
        Some(GatherFold::Chunk(rm_strides(&plan.src_shape)))
    } else if plan.fixed_desc.is_empty() {
        match broadcast_fold(plan) {
            Some(ShiftedGeom::Segs { strides, .. }) => Some(GatherFold::Chunk(strides)),
            _ => None,
        }
    } else {
        None
    }
}

/// Whether the gather at `ix`, placed just before its first reader, folds
/// into the fused group its readers compute in (`sink`): fusion absorbs
/// it, and every reader is arithmetic over the gather's box that joins that
/// group. A copy alone leaves a group with nothing to compute, emitted as
/// written, and any other reader (a `DyWrite`, a reduction, a scan, a table
/// read, an instruction of a segmented-reduction chain) reads the gather's
/// materialized box. `chained` is [`chained_tuple_instrs`] of `prog`.
pub(super) fn sunk_gather_fuses(
    prog: &TapeProgram,
    chained: &[bool],
    ix: usize,
    readers: &[usize],
) -> bool {
    let Instr::Gather { plan, out, .. } = prog.instrs[ix] else {
        return false;
    };
    let shape = &prog.slots[out as usize].shape;
    if gather_fold(&prog.plans[plan as usize], shape).is_none() {
        return false;
    }
    readers.iter().all(|&r| {
        let ins = &prog.instrs[r];
        let arith = matches!(
            ins,
            Instr::Bin { .. } | Instr::Un { .. } | Instr::Neg { .. } | Instr::Select { .. }
        );
        arith
            && !chained[r]
            && prec_at(prog, r) == prec_at(prog, ix)
            && ins.out().is_some_and(|o| {
                !prog.slots[o as usize].scalar && prog.slots[o as usize].shape == *shape
            })
    })
}

/// A gather plan is foldable as a shifted read when it is a pure constant
/// per-axis shift/wrap of a same-rank source: no fixed axes, no broadcast
/// axes, no axis permutation (the source box may differ — an interior-subset
/// stencil read). Its per-axis copy segments then give, for every covered
/// run, a constant flat source offset (the source's last axis advances with
/// stride 1, exactly like the output's), and every uncovered position reads
/// the ghost `+0.0` — byte-identical to the materialized gather.
fn foldable_plan(plan: &GatherPlan) -> bool {
    plan.fixed_desc.is_empty()
        && plan.mapped.iter().all(|&m| m)
        && plan.perm.iter().enumerate().all(|(i, &p)| p == i)
        && plan.src_shape.len() == plan.shape.len()
        && !plan.shape.is_empty()
        && !plan.shape.contains(&0)
}

/// A fold that reads every element of an equal-shape source at its own flat
/// offset: one run over the whole box, from source offset 0, contiguous.
fn is_identity_fold(shape: &[usize], geom: &ShiftedGeom) -> bool {
    if geom.elem_stride() != 1 {
        return false;
    }
    let sched = build_schedule(shape, std::slice::from_ref(geom));
    let n_elems: usize = shape.iter().product::<usize>().max(1);
    match &sched.nodes[..] {
        [RunNode::Run(r)] => r.out_off == 0 && r.len as usize == n_elems && r.in_off[..] == [0],
        _ => false,
    }
}

/// Cap on the fold's run-schedule shattering: a fold must keep runs coarse
/// (a shift along the innermost axis of a deep box would shatter the box
/// into per-row runs and eat the chunked executor's dispatch amortization —
/// such gathers stay materialized). What is counted is the runs the executor
/// dispatches, a repetition's body once per repetition, read off the built
/// schedule: a regular row pattern is one node of it however many rows it
/// covers, but each row is still one dispatch.
fn fold_runs_acceptable(shape: &[usize], geom: &ShiftedGeom) -> bool {
    let n_elems: usize = shape.iter().product::<usize>().max(1);
    let runs = build_schedule(shape, std::slice::from_ref(geom)).n_runs;
    runs <= 8 || n_elems / runs >= 64
}

/// Detect the globally-linear fold: `flat_src = a * flat_out + b` over the
/// whole box. Covers full-cover constant shifts with a common stride ratio,
/// fixed source axes (which contribute to `b`), and broadcast output axes of
/// extent 1 (which contribute nothing) — the classic instance is a level
/// slice `index(u, i, j, K)` materialized on an `[nx, ny, 1]` box from an
/// `[nx, ny, nz]` source (`a = nz`, `b = K`).
fn linear_fold(plan: &GatherPlan) -> Option<ShiftedGeom> {
    if plan.shape.is_empty() || plan.shape.contains(&0) || plan.src_shape.len() > 16 {
        return None;
    }
    let nd_out = plan.shape.len();
    let sstr = rm_strides(&plan.src_shape);
    let ostr = rm_strides(&plan.shape);
    let mut b = 0i64;
    let mut fixed_mask = [false; 16];
    for &(d, i0) in &plan.fixed_desc {
        fixed_mask[d] = true;
        b += sstr[d] * i0 as i64;
    }
    // Non-fixed source strides in ascending axis order, permuted into output
    // order exactly as `exec_gather` does.
    let reduced: Vec<i64> = (0..plan.src_shape.len())
        .filter(|d| !fixed_mask[*d])
        .map(|d| sstr[d])
        .collect();
    let mut mpos = 0usize;
    let mut ratio: Option<i64> = None;
    for a in 0..nd_out {
        if plan.mapped[a] {
            let eff = reduced[plan.perm[mpos]];
            mpos += 1;
            let segs = &plan.segs[a];
            if segs.len() != 1 {
                return None;
            }
            let (o, l, so) = segs[0];
            if o != 0 || l != plan.shape[a] {
                return None;
            }
            b += so as i64 * eff;
            if plan.shape[a] > 1 {
                if ostr[a] == 0 || eff % ostr[a] != 0 {
                    return None;
                }
                let r = eff / ostr[a];
                match ratio {
                    None => ratio = Some(r),
                    Some(p) if p == r => {}
                    _ => return None,
                }
            }
        } else if plan.shape[a] != 1 {
            // A broadcast axis of extent > 1 repeats source elements — not
            // expressible as a constant element stride.
            return None;
        }
    }
    Some(ShiftedGeom::Linear {
        a: ratio.unwrap_or(1),
        b,
    })
}

/// Detect the broadcast fold: a plan that REPEATS its source along some
/// output axes — a broadcast axis of any extent, and fixed source axes — but
/// keeps the mapped axes in order (no transposition). Along a broadcast axis
/// the source offset does not move (stride 0), a fixed axis contributes a
/// constant offset, and every mapped axis keeps its own copy segments (so
/// shifted and ghost reads fold exactly as [`foldable_plan`]'s do). The
/// classic instance is a contraction term read by the contracted index alone,
/// `e[j]` over a `window × output` box: constant along each run, one source
/// element per run.
fn broadcast_fold(plan: &GatherPlan) -> Option<ShiftedGeom> {
    if plan.shape.is_empty() || plan.shape.contains(&0) || plan.src_shape.len() > 16 {
        return None;
    }
    if plan.perm.iter().enumerate().any(|(i, &p)| p != i) {
        return None;
    }
    let sstr = rm_strides(&plan.src_shape);
    let mut base = 0i64;
    let mut fixed_mask = [false; 16];
    for &(d, i0) in &plan.fixed_desc {
        fixed_mask[d] = true;
        base += sstr[d] * i0 as i64;
    }
    let reduced: SmallVec<[i64; 4]> = (0..plan.src_shape.len())
        .filter(|d| !fixed_mask[*d])
        .map(|d| sstr[d])
        .collect();
    let mut strides: SmallVec<[i64; 4]> = SmallVec::new();
    let mut mpos = 0usize;
    for a in 0..plan.shape.len() {
        if plan.mapped[a] {
            strides.push(reduced[plan.perm[mpos]]);
            mpos += 1;
        } else {
            strides.push(0);
        }
    }
    Some(ShiftedGeom::Segs {
        segs: plan.segs.clone(),
        strides,
        base,
    })
}

impl ShiftedGeom {
    /// The source advance per element within a run.
    fn elem_stride(&self) -> i64 {
        match self {
            ShiftedGeom::Segs { strides, .. } => strides.last().copied().unwrap_or(1),
            ShiftedGeom::Linear { a, .. } => *a,
        }
    }
}

// ---------------------------------------------------------------------------
// Run-schedule construction.
// ---------------------------------------------------------------------------

/// A run schedule under construction: runs and repetitions in execution
/// (ascending output) order, maximally coalesced. Flattened into a
/// [`RunSchedule`] once built.
#[derive(Clone)]
enum Sched {
    Run(FusedRun),
    /// `body` executed `count` (>= 2) times, the `k`-th time stepped `k`
    /// times by `(out_step, in_step)` (see [`RunNode::Repeat`]).
    Rep {
        count: u32,
        out_step: i64,
        in_step: SmallVec<[i64; 2]>,
        body: Vec<Sched>,
    },
}

/// Whether run `b` continues run `a`: it starts where `a` ends in the output
/// and in every input (each advancing by its element stride), or both are
/// ghost there.
fn continues(a: &FusedRun, b: &FusedRun, elem_stride: &[i64]) -> bool {
    a.out_off as i64 + a.len as i64 == b.out_off as i64
        && a.in_off
            .iter()
            .zip(&b.in_off)
            .zip(elem_stride)
            .all(|((&x, &y), &es)| {
                (x == GHOST_OFF && y == GHOST_OFF)
                    || (x != GHOST_OFF && y != GHOST_OFF && x + a.len as i64 * es == y)
            })
}

/// `a` extended by the run `b` that [`continues`] it.
fn joined(a: FusedRun, b: &FusedRun) -> FusedRun {
    FusedRun {
        len: a.len + b.len,
        ..a
    }
}

/// `seq` moved `k` steps of `(out_step, in_step)`.
fn stepped_seq(seq: &[Sched], k: i64, out_step: i64, in_step: &[i64]) -> Vec<Sched> {
    seq.iter()
        .map(|n| match n {
            Sched::Run(r) => Sched::Run(r.stepped(k, out_step, in_step)),
            Sched::Rep {
                count,
                out_step: os,
                in_step: is,
                body,
            } => Sched::Rep {
                count: *count,
                out_step: *os,
                in_step: is.clone(),
                body: stepped_seq(body, k, out_step, in_step),
            },
        })
        .collect()
}

/// `count` repetitions of `body`: a `Rep`, or the body itself when once.
fn rep_of(count: u32, out_step: i64, in_step: &[i64], body: Vec<Sched>) -> Vec<Sched> {
    if count == 1 {
        body
    } else {
        vec![Sched::Rep {
            count,
            out_step,
            in_step: in_step.into(),
            body,
        }]
    }
}

/// The first run `seq` executes.
fn first_run(seq: &[Sched]) -> FusedRun {
    match &seq[0] {
        Sched::Run(r) => r.clone(),
        Sched::Rep { body, .. } => first_run(body),
    }
}

/// The last run `seq` executes.
fn last_run(seq: &[Sched]) -> FusedRun {
    match seq.last().expect("a schedule piece has a run") {
        Sched::Run(r) => r.clone(),
        Sched::Rep {
            count,
            out_step,
            in_step,
            body,
        } => last_run(body).stepped(*count as i64 - 1, *out_step, in_step),
    }
}

/// Remove and return the first run `seq` executes, peeling it out of any
/// repetition that starts with it.
fn pop_first(seq: &mut Vec<Sched>) -> FusedRun {
    match seq.remove(0) {
        Sched::Run(r) => r,
        Sched::Rep {
            count,
            out_step,
            in_step,
            body,
        } => {
            let rest = rep_of(
                count - 1,
                out_step,
                &in_step,
                stepped_seq(&body, 1, out_step, &in_step),
            );
            let mut head = body;
            let f = pop_first(&mut head);
            head.extend(rest);
            seq.splice(0..0, head);
            f
        }
    }
}

/// Remove and return the last run `seq` executes, peeling it out of any
/// repetition that ends with it.
fn pop_last(seq: &mut Vec<Sched>) -> FusedRun {
    match seq.pop().expect("a schedule piece has a run") {
        Sched::Run(r) => r,
        Sched::Rep {
            count,
            out_step,
            in_step,
            body,
        } => {
            let mut tail = stepped_seq(&body, count as i64 - 1, out_step, &in_step);
            let l = pop_last(&mut tail);
            seq.extend(rep_of(count - 1, out_step, &in_step, body));
            seq.extend(tail);
            l
        }
    }
}

/// Append `more` to `seq`, joining the two runs at the seam when the second
/// continues the first. Both are maximally coalesced, and so is the result:
/// a joined run keeps the first run's start and the second's end.
fn append(seq: &mut Vec<Sched>, mut more: Vec<Sched>, elem_stride: &[i64]) {
    if more.is_empty() {
        return;
    }
    if !seq.is_empty() && continues(&last_run(seq), &first_run(&more), elem_stride) {
        let l = pop_last(seq);
        let f = pop_first(&mut more);
        seq.push(Sched::Run(joined(l, &f)));
    }
    seq.extend(more);
}

/// `p` executed `count` times, stepped `(out_step, in_step)` each time.
///
/// When one repetition's last run is continued by the next one's first run
/// (a row that ends where the next begins), the repetition is rotated so the
/// repeated body ENDS with that joined run: `first, (middle, last+first) x
/// (count - 1), middle, last`. A one-run `p` so continued is one long run.
fn repeat(
    p: Vec<Sched>,
    count: u32,
    out_step: i64,
    in_step: &[i64],
    elem_stride: &[i64],
) -> Vec<Sched> {
    if count == 1 {
        return p;
    }
    let next_first = first_run(&p).stepped(1, out_step, in_step);
    if !continues(&last_run(&p), &next_first, elem_stride) {
        return rep_of(count, out_step, in_step, p);
    }
    if let [Sched::Run(r)] = &p[..] {
        return vec![Sched::Run(FusedRun {
            len: r.len * count,
            ..r.clone()
        })];
    }
    let mut middle = p;
    let first = pop_first(&mut middle);
    let last_rep = stepped_seq(&middle, count as i64 - 1, out_step, in_step);
    let last = pop_last(&mut middle);
    let mut body = middle;
    body.push(Sched::Run(joined(last, &next_first)));
    let mut out = vec![Sched::Run(first)];
    out.extend(rep_of(count - 1, out_step, in_step, body));
    out.extend(last_rep);
    out
}

/// Build the run schedule for a group box `shape` with the given
/// shifted-input geometries: runs ascending in `out_off` and maximally
/// coalesced, with regular repetitions kept as [`RunNode::Repeat`]. Within
/// a run every shifted input advances by its constant element stride through
/// its own row-major source, so a run stores one flat source offset (or
/// ghost) per shifted input.
///
/// The box is cut, per axis, at every segment boundary of every `Segs` input.
/// Inside one interval of a leading axis, every input's source offset (or its
/// ghost) moves by a constant per coordinate, so the rows below each
/// coordinate are the rows below the interval's first coordinate, stepped:
/// the schedule is built once per interval (not per row) and repeated, and
/// its size is set by the number of intervals, not by the box's extents.
fn build_schedule(shape: &[usize], shifted: &[ShiftedGeom]) -> RunSchedule {
    let n_elems: usize = shape.iter().product::<usize>().max(1);
    if shifted.is_empty() {
        return RunSchedule::whole(n_elems);
    }
    let nd = shape.len();
    debug_assert!(nd > 0, "a shifted input has a box of rank >= 1");
    let cx = SchedCx::new(shape, shifted);
    let base: SmallVec<[Option<i64>; 2]> = shifted
        .iter()
        .map(|g| match g {
            ShiftedGeom::Segs { base, .. } => Some(*base),
            ShiftedGeom::Linear { .. } => Some(0),
        })
        .collect();
    let seq = cx.axis(0, 0, &base);
    let mut nodes = Vec::new();
    let (n_runs, depth) = flatten(&seq, &mut nodes);
    debug_assert_eq!(
        {
            let mut n = 0usize;
            let s = RunSchedule {
                nodes: nodes.clone(),
                n_runs,
                depth,
            };
            s.for_each_run(|r| n += r.len as usize);
            n
        },
        n_elems,
        "run schedule must tile the box"
    );
    RunSchedule {
        nodes,
        n_runs,
        depth,
    }
}

/// Write `seq` into `nodes` in pre-order; returns the runs it executes and
/// its repetition depth.
fn flatten(seq: &[Sched], nodes: &mut Vec<RunNode>) -> (usize, usize) {
    let (mut runs, mut depth) = (0usize, 0usize);
    for n in seq {
        match n {
            Sched::Run(r) => {
                nodes.push(RunNode::Run(r.clone()));
                runs += 1;
            }
            Sched::Rep {
                count,
                out_step,
                in_step,
                body,
            } => {
                let at = nodes.len();
                nodes.push(RunNode::Run(FusedRun {
                    out_off: 0,
                    len: 0,
                    in_off: SmallVec::new(),
                }));
                let (r, d) = flatten(body, nodes);
                nodes[at] = RunNode::Repeat {
                    count: *count,
                    body: (nodes.len() - at - 1) as u32,
                    out_step: *out_step,
                    in_step: in_step.clone(),
                };
                runs += *count as usize * r;
                depth = depth.max(d + 1);
            }
        }
    }
    (runs, depth)
}

/// The per-box tables [`build_schedule`] walks.
struct SchedCx<'a> {
    shape: &'a [usize],
    shifted: &'a [ShiftedGeom],
    strides: SmallVec<[i64; 4]>,
    /// Per input: its advance per element within a run.
    elem_stride: SmallVec<[i64; 2]>,
    /// Per axis: `[(start, len, per-input source start on this axis)]`, the
    /// start `None` where the interval is uncovered (a ghost) and for a
    /// `Linear` input (which imposes no cuts).
    #[allow(clippy::type_complexity)]
    axis_ivals: Vec<Vec<(usize, usize, SmallVec<[Option<i64>; 2]>)>>,
}

impl<'a> SchedCx<'a> {
    fn new(shape: &'a [usize], shifted: &'a [ShiftedGeom]) -> Self {
        let nd = shape.len();
        let mut axis_ivals = Vec::with_capacity(nd);
        for d in 0..nd {
            let mut cuts: BTreeSet<usize> = BTreeSet::new();
            cuts.insert(0);
            cuts.insert(shape[d]);
            for g in shifted {
                if let ShiftedGeom::Segs { segs, .. } = g {
                    for &(o, l, _) in &segs[d] {
                        cuts.insert(o);
                        cuts.insert(o + l);
                    }
                }
            }
            let cuts: Vec<usize> = cuts.into_iter().collect();
            let mut ivals = Vec::with_capacity(cuts.len() - 1);
            for w in cuts.windows(2) {
                let (a, b) = (w[0], w[1]);
                let src_start: SmallVec<[Option<i64>; 2]> = shifted
                    .iter()
                    .map(|g| match g {
                        ShiftedGeom::Segs { segs, .. } => segs[d]
                            .iter()
                            .find(|&&(o, l, _)| a >= o && b <= o + l)
                            .map(|&(o, _, so)| so as i64 + (a as i64 - o as i64)),
                        ShiftedGeom::Linear { .. } => None,
                    })
                    .collect();
                ivals.push((a, b - a, src_start));
            }
            axis_ivals.push(ivals);
        }
        SchedCx {
            shape,
            shifted,
            strides: rm_strides(shape),
            elem_stride: shifted.iter().map(ShiftedGeom::elem_stride).collect(),
            axis_ivals,
        }
    }

    /// The schedule of the sub-box of axes `d..` at output offset `off`,
    /// where each `Segs` input's source offset so far is `src[j]` (`None`: a
    /// ghost on some outer axis).
    fn axis(&self, d: usize, off: i64, src: &[Option<i64>]) -> Vec<Sched> {
        let inner = d + 1 == self.shape.len();
        let mut seq: Vec<Sched> = Vec::new();
        for (st, len, starts) in &self.axis_ivals[d] {
            let o = off + *st as i64 * self.strides[d];
            let at: SmallVec<[Option<i64>; 2]> = self
                .shifted
                .iter()
                .enumerate()
                .map(|(j, g)| match (g, src[j], starts[j]) {
                    (ShiftedGeom::Segs { strides, .. }, Some(b), Some(p)) => {
                        Some(b + p * strides[d])
                    }
                    (ShiftedGeom::Segs { .. }, _, _) => None,
                    (ShiftedGeom::Linear { .. }, _, _) => Some(0),
                })
                .collect();
            let piece = if inner {
                let in_off = self
                    .shifted
                    .iter()
                    .zip(&at)
                    .map(|(g, s)| match (g, s) {
                        (ShiftedGeom::Linear { a, b }, _) => a * o + b,
                        (ShiftedGeom::Segs { .. }, Some(s)) => *s,
                        (ShiftedGeom::Segs { .. }, None) => GHOST_OFF,
                    })
                    .collect();
                vec![Sched::Run(FusedRun {
                    out_off: o as u32,
                    len: *len as u32,
                    in_off,
                })]
            } else {
                let in_step: SmallVec<[i64; 2]> = self
                    .shifted
                    .iter()
                    .map(|g| match g {
                        ShiftedGeom::Linear { a, .. } => a * self.strides[d],
                        ShiftedGeom::Segs { strides, .. } => strides[d],
                    })
                    .collect();
                let sub = self.axis(d + 1, o, &at);
                repeat(
                    sub,
                    *len as u32,
                    self.strides[d],
                    &in_step,
                    &self.elem_stride,
                )
            };
            append(&mut seq, piece, &self.elem_stride);
        }
        seq
    }
}

// ---------------------------------------------------------------------------
// Superop peephole.
// ---------------------------------------------------------------------------

fn bin2_arith(op: BinCode) -> bool {
    matches!(
        op,
        BinCode::Add | BinCode::Sub | BinCode::Mul | BinCode::Div
    )
}

/// Which superops the peephole may form. All settings are bit-identical (a
/// superop applies the identical scalar kernels in the identical order); the
/// knobs exist because the WIN is configuration-dependent, and the defaults were
/// set by measurement on one model, one grid and one SIMD width. Re-measure
/// before trusting them for a different shape of workload.
#[derive(Clone, Copy, Debug)]
pub(crate) struct SuperopCfg {
    /// Merge `+ - * /` three-op chains into [`MicroOp::Bin3`]. Default OFF: it
    /// measured slower — the chunked executor is bound by the per-op loads/stores
    /// through the (L2-resident) register file, not by dispatch, and Bin3's
    /// all-pointer form adds a fourth input stream plus splat-register traffic
    /// that outweighs the two intermediate round-trips it saves. Opt-in via
    /// `ESS_TAPE_BIN3=1`: a TUNING THRESHOLD, and under `native` a refusal
    /// boundary rather than a fallback trigger — moving it changes which
    /// program is built, never which evaluator runs
    /// (`esm-libraries-spec.md` §2.5.10).
    pub bin3: bool,
    /// Extend Bin2 beyond the `+ - * /` square to the mask/clamp pairs of
    /// [`bin2_pair_ok`]. Default ON: faster in the generic build, neutral under
    /// the SIMD clones, and fewer element-ops either way.
    /// `ESS_TAPE_EXTPAIR_DISABLE=1` reverts: a TUNING THRESHOLD like
    /// `ESS_TAPE_BIN3`, shaping the program a compiler builds and never
    /// selecting a different evaluator.
    pub ext_pairs: bool,
}

impl SuperopCfg {
    /// The production configuration (env-overridable, read once).
    pub(crate) fn from_env() -> Self {
        use std::sync::OnceLock;
        static CFG: OnceLock<SuperopCfg> = OnceLock::new();
        let truthy = |k: &str| {
            std::env::var(k)
                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
                .unwrap_or(false)
        };
        *CFG.get_or_init(|| SuperopCfg {
            bin3: truthy("ESS_TAPE_BIN3"),
            ext_pairs: !truthy("ESS_TAPE_EXTPAIR_DISABLE"),
        })
    }
}

/// The most rows a group with a reverse scan may visit: it visits them one
/// at a time, each a pass of its own, so a long scanned axis stays a
/// standalone scan (whose kernel sweeps a row's lanes in one loop). The
/// bound is on the scanned axis, not the box, so the choice does not change
/// with the number of columns.
const MAX_REVERSE_ROWS: usize = 4096;

/// The (op1, op2) pairs the executor monomorphizes as [`MicroOp::Bin2`]: the
/// full `+ - * /` square, plus the pairs the simpleclimate adjacency
/// histogram showed material (a multiply feeding a comparison mask; the
/// min/max clamp and limiter idioms). MUST stay in sync with the executor's
/// `Bin2` match arms — an unlisted pair is simply never merged.
fn bin2_pair_ok(op1: BinCode, op2: BinCode, ext_pairs: bool) -> bool {
    use BinCode::{Ge, Gt, Le, Lt, Max, Min, Mul};
    (bin2_arith(op1) && bin2_arith(op2))
        || ext_pairs
            && matches!(
                (op1, op2),
                (Mul, Gt)
                    | (Mul, Ge)
                    | (Mul, Lt)
                    | (Mul, Le)
                    | (Min, Max)
                    | (Max, Min)
                    | (Mul, Min)
                    | (Min, Mul)
                    | (Mul, Max)
                    | (Max, Mul)
            )
}

/// Visit every operand of a micro-op.
fn for_each_operand(op: &MicroOp, mut f: impl FnMut(&MRef)) {
    match op {
        MicroOp::Bin { a, b, .. } => {
            f(a);
            f(b);
        }
        MicroOp::Un { a, .. }
        | MicroOp::Neg { a, .. }
        | MicroOp::Mov { a, .. }
        | MicroOp::Load { a, .. }
        | MicroOp::Scan { a, .. } => f(a),
        MicroOp::Select { cond, a, b, .. } => {
            f(cond);
            f(a);
            f(b);
        }
        MicroOp::Bin2 { a, b, c, .. } => {
            f(a);
            f(b);
            f(c);
        }
        MicroOp::Bin3 { a, b, c, d, .. } => {
            f(a);
            f(b);
            f(c);
            f(d);
        }
    }
}

/// SSA register use counts (live-outs count as a use). Valid only while the
/// micro-program is in SSA form (op `i` defines register `i`).
fn ssa_uses(micro: &[MicroOp], outputs: &[(GroupIx, SlotId)]) -> Vec<u32> {
    let mut uses = vec![0u32; micro.len()];
    for op in micro {
        for_each_operand(op, |m| {
            if let MRef::Reg(r) = m {
                uses[*r as usize] += 1;
            }
        });
    }
    for &(ssa, _) in outputs {
        uses[ssa as usize] += 1;
    }
    uses
}

/// A `+ - * /` `Bin`'s operands, or `None` for any other micro-op.
fn arith_bin(op: &MicroOp) -> Option<(&MRef, &MRef)> {
    match op {
        MicroOp::Bin { op, a, b, .. } if bin2_arith(*op) => Some((a, b)),
        _ => None,
    }
}

/// The [`ScanFuse`]s of an SSA micro-program (`outputs` holding its
/// live-outs, an absorbed reduction's register included), after moving each
/// fusable neighbour next to its scan: the `Bin` computing a scan's operand
/// (read by the scan alone) to just before it, and the one `Bin` reading a
/// scan's value (whose other operand exists before the scan) to just after
/// it. Moving an op past ops it neither reads nor feeds keeps the program's
/// values, and a `Bin` joins one scan at most.
fn scan_fusions(
    micro: &mut Vec<MicroOp>,
    outputs: &mut [(GroupIx, SlotId)],
) -> SmallVec<[ScanFuse; 1]> {
    let mut found: SmallVec<[ScanFuse; 1]> = SmallVec::new();
    if !micro.iter().any(|m| {
        matches!(
            m,
            MicroOp::Scan {
                post: 1,
                op: BinCode::Add | BinCode::Mul,
                reverse: false,
                ..
            }
        )
    }) {
        return found;
    }
    let n = micro.len();
    let uses = ssa_uses(micro, outputs);
    let mut readers: Vec<SmallVec<[usize; 2]>> = vec![SmallVec::new(); n];
    for (i, op) in micro.iter().enumerate() {
        for_each_operand(op, |m| {
            if let MRef::Reg(r) = m {
                readers[*r as usize].push(i);
            }
        });
    }
    let live_out = |r: usize| outputs.iter().filter(|o| o.0 as usize == r).count() as u32;
    // Per scan: the op moved in front of it and the op moved behind it.
    let mut taken = vec![false; n];
    let mut plan: Vec<(usize, Option<usize>, Option<usize>)> = Vec::new();
    for (s, op) in micro.iter().enumerate() {
        // The executor monomorphizes these two folds (sums and products).
        let MicroOp::Scan {
            op: BinCode::Add | BinCode::Mul,
            a,
            post: 1,
            reverse: false,
            ..
        } = op
        else {
            continue;
        };
        let pre = match a {
            MRef::Reg(q) => Some(*q as usize)
                .filter(|&q| !taken[q] && uses[q] == 1 && arith_bin(&micro[q]).is_some()),
            _ => None,
        };
        let post = match &readers[s][..] {
            &[p] if uses[s] == 1 + live_out(s) && !taken[p] => {
                arith_bin(&micro[p]).and_then(|(x, y)| {
                    let me = MRef::Reg(s as GroupIx);
                    let other = match (*x == me, *y == me) {
                        (true, false) => y,
                        (false, true) => x,
                        _ => return None,
                    };
                    let early = !matches!(other, MRef::Reg(r) if *r as usize > s);
                    early.then_some(p)
                })
            }
            _ => None,
        };
        if pre.is_none() && post.is_none() {
            continue;
        }
        for &i in pre.iter().chain(post.iter()) {
            taken[i] = true;
        }
        plan.push((s, pre, post));
    }
    if plan.is_empty() {
        return found;
    }
    // The new order, then every register renamed to its op's new position.
    let mut order: Vec<usize> = Vec::with_capacity(n);
    let mut next = plan.iter().peekable();
    for i in 0..n {
        if taken[i] {
            continue;
        }
        match next.peek() {
            Some(&&(s, pre, post)) if s == i => {
                order.extend(pre);
                order.push(s);
                order.extend(post);
                found.push(ScanFuse {
                    scan: order.len() as u32 - 1 - post.is_some() as u32,
                    pre: pre.is_some(),
                    post: post.is_some(),
                    keep: post.is_some() && live_out(s) > 0,
                });
                next.next();
            }
            _ => order.push(i),
        }
    }
    debug_assert_eq!(order.len(), n);
    let mut new_of = vec![0 as GroupIx; n];
    for (new, &old) in order.iter().enumerate() {
        new_of[old] = new as GroupIx;
    }
    let mut old_micro = std::mem::take(micro);
    for &old in &order {
        let mut op = std::mem::replace(
            &mut old_micro[old],
            MicroOp::Mov {
                a: MRef::Scal(0),
                out: 0,
            },
        );
        for_each_operand_mut(&mut op, |m| {
            if let MRef::Reg(r) = m {
                *r = new_of[*r as usize];
            }
        });
        set_micro_out(&mut op, new_of[old]);
        micro.push(op);
    }
    for o in outputs.iter_mut() {
        o.0 = new_of[o.0 as usize];
    }
    found
}

/// Visit every operand of a micro-op, mutably.
fn for_each_operand_mut(op: &mut MicroOp, mut f: impl FnMut(&mut MRef)) {
    match op {
        MicroOp::Bin { a, b, .. } => {
            f(a);
            f(b);
        }
        MicroOp::Un { a, .. }
        | MicroOp::Neg { a, .. }
        | MicroOp::Mov { a, .. }
        | MicroOp::Load { a, .. }
        | MicroOp::Scan { a, .. } => f(a),
        MicroOp::Select { cond, a, b, .. } => {
            f(cond);
            f(a);
            f(b);
        }
        MicroOp::Bin2 { a, b, c, .. } => {
            f(a);
            f(b);
            f(c);
        }
        MicroOp::Bin3 { a, b, c, d, .. } => {
            f(a);
            f(b);
            f(c);
            f(d);
        }
    }
}

/// Set the register a micro-op writes.
fn set_micro_out(op: &mut MicroOp, r: GroupIx) {
    match op {
        MicroOp::Bin { out, .. }
        | MicroOp::Un { out, .. }
        | MicroOp::Neg { out, .. }
        | MicroOp::Mov { out, .. }
        | MicroOp::Select { out, .. }
        | MicroOp::Bin2 { out, .. }
        | MicroOp::Bin3 { out, .. }
        | MicroOp::Load { out, .. }
        | MicroOp::Scan { out, .. } => *out = r,
    }
}

/// The register a micro-op writes.
pub(super) fn micro_out(op: &MicroOp) -> GroupIx {
    match op {
        MicroOp::Bin { out, .. }
        | MicroOp::Un { out, .. }
        | MicroOp::Neg { out, .. }
        | MicroOp::Mov { out, .. }
        | MicroOp::Select { out, .. }
        | MicroOp::Bin2 { out, .. }
        | MicroOp::Bin3 { out, .. }
        | MicroOp::Load { out, .. }
        | MicroOp::Scan { out, .. } => *out,
    }
}

/// [`FusedSpec::direct`] of an allocated micro-program: each output whose
/// register's last writer is followed by no read of it, which `reduce_reg`
/// (an absorbed reduction's register) is not, and which no other output
/// names, in ascending micro-op order.
fn direct_stores(
    micro: &[MicroOp],
    outputs: &[(GroupIx, SlotId)],
    reduce_reg: Option<GroupIx>,
) -> SmallVec<[(u32, u32); 2]> {
    let mut direct = SmallVec::new();
    for (k, &(reg, _)) in outputs.iter().enumerate() {
        if Some(reg) == reduce_reg || outputs.iter().filter(|o| o.0 == reg).count() > 1 {
            continue;
        }
        let Some(w) = micro.iter().rposition(|op| micro_out(op) == reg) else {
            continue;
        };
        if micro[w + 1..].iter().any(|op| reads_reg(op, reg)) {
            continue;
        }
        direct.push((w as u32, k as u32));
    }
    direct.sort_unstable();
    direct
}

/// `true` when `op` reads register `r`.
pub(super) fn reads_reg(op: &MicroOp, r: GroupIx) -> bool {
    let mut hit = false;
    for_each_operand(op, |m| {
        if *m == MRef::Reg(r) {
            hit = true;
        }
    });
    hit
}

/// Short kind label for the step-4b composition/adjacency histograms.
fn mop_label(op: &MicroOp) -> String {
    match op {
        MicroOp::Bin { op, .. } => format!("Bin({op:?})"),
        MicroOp::Un { op, .. } => format!("Un({op:?})"),
        MicroOp::Neg { .. } => "Neg".to_string(),
        MicroOp::Select { .. } => "Select".to_string(),
        MicroOp::Mov { .. } => "Mov".to_string(),
        MicroOp::Load { .. } => "Load".to_string(),
        MicroOp::Scan { op, .. } => format!("Scan({op:?})"),
        MicroOp::Bin2 { op1, op2, swap, .. } => {
            format!("Bin2({op1:?},{op2:?}{})", if *swap { ",swap" } else { "" })
        }
        MicroOp::Bin3 { op1, op2, op3, .. } => {
            format!("Bin3({op1:?},{op2:?},{op3:?})")
        }
    }
}

/// Merge adjacent `Bin` pairs whose intermediate has exactly one consumer
/// (the very next op) into [`MicroOp::Bin2`]: the intermediate then lives in
/// a CPU register instead of a chunk buffer — one store + one load + one
/// dispatch fewer per element, with the identical two kernels applied in the
/// identical order (bit-identity is untouched). Restricted to the
/// `+ - * /` kernels the executor monomorphizes.
fn merge_superops(micro: &mut Vec<MicroOp>, outputs: &mut [(GroupIx, SlotId)], cfg: SuperopCfg) {
    let n_ssa = micro.len();
    // Use counts per SSA register (live-outs count as a use).
    let mut uses = vec![0u32; n_ssa];
    let count = |m: &MRef, uses: &mut Vec<u32>| {
        if let MRef::Reg(r) = m {
            uses[*r as usize] += 1;
        }
    };
    for op in micro.iter() {
        match op {
            MicroOp::Bin { a, b, .. } => {
                count(a, &mut uses);
                count(b, &mut uses);
            }
            MicroOp::Un { a, .. }
            | MicroOp::Neg { a, .. }
            | MicroOp::Mov { a, .. }
            | MicroOp::Load { a, .. }
            | MicroOp::Scan { a, .. } => count(a, &mut uses),
            MicroOp::Select { cond, a, b, .. } => {
                count(cond, &mut uses);
                count(a, &mut uses);
                count(b, &mut uses);
            }
            MicroOp::Bin2 { .. } | MicroOp::Bin3 { .. } => unreachable!("peephole runs once"),
        }
    }
    for &(ssa, _) in outputs.iter() {
        uses[ssa as usize] += 1;
    }

    /// How the SSA register `t` is consumed by a `Bin` op: `Some((swap,
    /// other))` when exactly one operand is `t` (`swap` = `t` is the RIGHT
    /// operand), `None` when neither or both are `t`.
    fn consumes<'a>(ca: &'a MRef, cb: &'a MRef, t: GroupIx) -> Option<(bool, &'a MRef)> {
        let t = MRef::Reg(t);
        if *ca == t && *cb != t {
            Some((false, cb))
        } else if *cb == t && *ca != t {
            Some((true, ca))
        } else {
            None
        }
    }

    let mut merged: Vec<MicroOp> = Vec::with_capacity(micro.len());
    // old SSA id -> new SSA id (None = merged away, never referenced again).
    let mut new_of: Vec<Option<GroupIx>> = vec![None; n_ssa];
    let remap = |m: &MRef, new_of: &[Option<GroupIx>]| -> MRef {
        match m {
            MRef::Reg(r) => MRef::Reg(new_of[*r as usize].expect("operand defined")),
            other => *other,
        }
    };
    let mut i = 0usize;
    while i < micro.len() {
        // Try to merge micro[i..i+3] into a Bin3 (a `+ - * /` three-op chain
        // whose two intermediates are each single-use, consumed by the next
        // op), then micro[i..i+2] into a Bin2.
        if i + 2 < micro.len() && cfg.bin3 {
            if let (
                MicroOp::Bin {
                    op: op1,
                    a,
                    b,
                    out: t1,
                },
                MicroOp::Bin {
                    op: op2,
                    a: ca,
                    b: cb,
                    out: t2,
                },
                MicroOp::Bin {
                    op: op3,
                    a: da,
                    b: db,
                    ..
                },
            ) = (&micro[i], &micro[i + 1], &micro[i + 2])
            {
                if let (Some((swap2, other2)), Some((swap3, other3))) =
                    (consumes(ca, cb, *t1), consumes(da, db, *t2))
                {
                    if uses[*t1 as usize] == 1
                        && uses[*t2 as usize] == 1
                        && bin2_arith(*op1)
                        && bin2_arith(*op2)
                        && bin2_arith(*op3)
                        && *other3 != MRef::Reg(*t1)
                    {
                        let new_id = merged.len() as GroupIx;
                        let mop = MicroOp::Bin3 {
                            op1: *op1,
                            a: remap(a, &new_of),
                            b: remap(b, &new_of),
                            op2: *op2,
                            c: remap(other2, &new_of),
                            swap2,
                            op3: *op3,
                            d: remap(other3, &new_of),
                            swap3,
                            out: new_id,
                        };
                        merged.push(mop);
                        new_of[i + 2] = Some(new_id); // final consumer's value
                        i += 3;
                        continue;
                    }
                }
            }
        }
        // Try to merge micro[i] (producer) into micro[i+1] (consumer).
        if i + 1 < micro.len() {
            if let (
                MicroOp::Bin {
                    op: op1,
                    a,
                    b,
                    out: t,
                },
                MicroOp::Bin {
                    op: op2,
                    a: ca,
                    b: cb,
                    ..
                },
            ) = (&micro[i], &micro[i + 1])
            {
                if let Some((swap, other)) = consumes(ca, cb, *t) {
                    if uses[*t as usize] == 1 && bin2_pair_ok(*op1, *op2, cfg.ext_pairs) {
                        let new_id = merged.len() as GroupIx;
                        let mop = MicroOp::Bin2 {
                            op1: *op1,
                            a: remap(a, &new_of),
                            b: remap(b, &new_of),
                            op2: *op2,
                            c: remap(other, &new_of),
                            swap,
                            out: new_id,
                        };
                        merged.push(mop);
                        new_of[i + 1] = Some(new_id); // consumer's value
                        i += 2;
                        continue;
                    }
                }
            }
        }
        let new_id = merged.len() as GroupIx;
        let mut op = micro[i].clone();
        match &mut op {
            MicroOp::Bin { a, b, out, .. } => {
                *a = remap(a, &new_of);
                *b = remap(b, &new_of);
                *out = new_id;
            }
            MicroOp::Un { a, out, .. }
            | MicroOp::Neg { a, out }
            | MicroOp::Mov { a, out }
            | MicroOp::Load { a, out, .. }
            | MicroOp::Scan { a, out, .. } => {
                *a = remap(a, &new_of);
                *out = new_id;
            }
            MicroOp::Select { cond, a, b, out } => {
                *cond = remap(cond, &new_of);
                *a = remap(a, &new_of);
                *b = remap(b, &new_of);
                *out = new_id;
            }
            MicroOp::Bin2 { .. } | MicroOp::Bin3 { .. } => unreachable!(),
        }
        merged.push(op);
        new_of[i] = Some(new_id);
        i += 1;
    }
    for (ssa, _) in outputs.iter_mut() {
        *ssa = new_of[*ssa as usize].expect("live-out survives merging");
    }
    *micro = merged;
}

// ---------------------------------------------------------------------------
// Register allocation.
// ---------------------------------------------------------------------------

/// Remap SSA micro registers onto a small recycled physical register file.
/// An op's `out` register is allocated BEFORE its dying operands are freed,
/// so `out` never aliases an operand register (which lets the executor's
/// chunk loops use disjoint slices).
pub(super) fn allocate_registers(
    micro: &mut [MicroOp],
    outputs: &mut [(GroupIx, SlotId)],
) -> GroupIx {
    let n_ssa = micro.len();
    let mut last_use = vec![usize::MAX; n_ssa]; // MAX = never used
    let use_at = |m: &MRef, i: usize, last_use: &mut Vec<usize>| {
        if let MRef::Reg(r) = m {
            let lu = &mut last_use[*r as usize];
            *lu = if *lu == usize::MAX { i } else { (*lu).max(i) };
        }
    };
    for (i, op) in micro.iter().enumerate() {
        for_each_operand(op, |m| use_at(m, i, &mut last_use));
    }
    // Live-outs are used "after the end".
    for &(ssa, _) in outputs.iter() {
        last_use[ssa as usize] = usize::MAX.wrapping_sub(1).max(n_ssa); // sentinel: beyond program
    }
    // Correction: distinguish "never used, not output" (dead) from live-out.
    // Recompute: dead = last_use == usize::MAX and not an output.
    let is_output: Vec<bool> = {
        let mut v = vec![false; n_ssa];
        for &(ssa, _) in outputs.iter() {
            v[ssa as usize] = true;
        }
        v
    };

    let mut phys_of = vec![GroupIx::MAX; n_ssa];
    let mut free: Vec<GroupIx> = Vec::new();
    let mut n_phys: GroupIx = 0;
    for i in 0..n_ssa {
        // Allocate out.
        let p = free.pop().unwrap_or_else(|| {
            n_phys += 1;
            n_phys - 1
        });
        phys_of[i] = p;
        // Free operands dying at i (after allocation, so out != operand).
        let free_if_dies = |m: &MRef, free: &mut Vec<GroupIx>| {
            if let MRef::Reg(r) = m {
                let r = *r as usize;
                if last_use[r] == i && !is_output[r] {
                    free.push(phys_of[r]);
                }
            }
        };
        let mut seen: SmallVec<[MRef; 4]> = SmallVec::new();
        for_each_operand(&micro[i], |m| {
            if !seen.contains(m) {
                seen.push(*m);
            }
        });
        for m in &seen {
            free_if_dies(m, &mut free);
        }
        // A dead value (never read, not an output) frees immediately.
        if last_use[i] == usize::MAX && !is_output[i] {
            free.push(p);
        }
    }

    // Remap.
    let remap = |m: &mut MRef| {
        if let MRef::Reg(r) = m {
            *m = MRef::Reg(phys_of[*r as usize]);
        }
    };
    for op in micro.iter_mut() {
        match op {
            MicroOp::Bin { a, b, out, .. } => {
                remap(a);
                remap(b);
                *out = phys_of[*out as usize];
            }
            MicroOp::Un { a, out, .. }
            | MicroOp::Neg { a, out }
            | MicroOp::Mov { a, out }
            | MicroOp::Load { a, out, .. }
            | MicroOp::Scan { a, out, .. } => {
                remap(a);
                *out = phys_of[*out as usize];
            }
            MicroOp::Select { cond, a, b, out } => {
                remap(cond);
                remap(a);
                remap(b);
                *out = phys_of[*out as usize];
            }
            MicroOp::Bin2 { a, b, c, out, .. } => {
                remap(a);
                remap(b);
                remap(c);
                *out = phys_of[*out as usize];
            }
            MicroOp::Bin3 {
                a, b, c, d, out, ..
            } => {
                remap(a);
                remap(b);
                remap(c);
                remap(d);
                *out = phys_of[*out as usize];
            }
        }
    }
    for (ssa, _) in outputs.iter_mut() {
        *ssa = phys_of[*ssa as usize];
    }
    n_phys
}

// ---------------------------------------------------------------------------
// The pass.
// ---------------------------------------------------------------------------

pub(super) fn fuse_program(prog: &mut TapeProgram, cfg: SuperopCfg) {
    let n = prog.instrs.len();
    // Global reader index sets per slot.
    let mut readers: Vec<SmallVec<[u32; 4]>> = vec![SmallVec::new(); prog.slots.len()];
    for (i, ins) in prog.instrs.iter().enumerate() {
        ins.for_each_read(&prog.tables(), |s| readers[s as usize].push(i as u32));
    }

    let sink = Sink {
        stats: FuseStats {
            instrs_before: n,
            ..Default::default()
        },
        fused: Vec::new(),
        instrs: Vec::with_capacity(n),
        prov: Vec::with_capacity(n),
        prec: Vec::with_capacity(if prog.precision.is_empty() { 0 } else { n }),
        micro_hist: FxHashMap::default(),
        adj_hist: FxHashMap::default(),
        chain_hist: FxHashMap::default(),
    };
    let chained = chained_tuple_instrs(prog);
    let mut fx = FuseCtx {
        prog,
        readers: &readers,
        sink,
        cfg,
        chained,
    };

    let mut section_counts = [0usize; 3];
    for (sec, cad) in [Cadence::Const, Cadence::Segment, Cadence::Continuous]
        .into_iter()
        .enumerate()
    {
        let range = fx.prog.section_range(cad);
        let start_len = fx.sink.instrs.len();
        if cad == Cadence::Continuous {
            // The derivative path and the output-only work after it
            // (`prune`) fuse apart, so a derivative call computes no
            // output-only value inside one of its groups.
            let cut = (range.start + fx.prog.n_rhs as usize).min(range.end);
            fuse_section(&mut fx, range.start..cut);
            fuse_section(&mut fx, cut..range.end);
        } else {
            fuse_section(&mut fx, range);
        }
        section_counts[sec] = fx.sink.instrs.len() - start_len;
    }

    let mut sink = fx.sink;
    sink.stats.instrs_after = sink.instrs.len();
    let desc = |m: FxHashMap<String, usize>| -> Vec<(String, usize)> {
        let mut v: Vec<_> = m.into_iter().collect();
        v.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
        v
    };
    sink.stats.micro_hist = desc(sink.micro_hist);
    sink.stats.adj_hist = desc(sink.adj_hist);
    sink.stats.chain_hist = {
        let mut v: Vec<_> = sink.chain_hist.into_iter().collect();
        v.sort_by_key(|&(len, _)| len);
        v
    };
    prog.instrs = sink.instrs;
    prog.provenance = sink.prov;
    prog.precision = sink.prec;
    prog.n_const = section_counts[0] as u32;
    prog.n_segment = section_counts[1] as u32;
    prog.fused = sink.fused;
    prog.fuse_stats = sink.stats;
    absorb_pre_reductions(prog);
}

/// Hand a leading-axes [`Instr::Reduce`] to the fused group right after it
/// when that group runs over the reduction's source box and reads its
/// result ([`PreReduce`]): the group then folds a range of columns and
/// consumes them while they are in cache, instead of one pass over the
/// source to reduce and another to consume. Each cell still folds its terms
/// in the instruction's order, and every element of the group is computed by
/// the same micro-ops, so the program's values do not change.
///
/// Only outside conditional regions and recurrence bodies (whose skip
/// counts are positional), within one section, and with no per-instruction
/// precision.
fn absorb_pre_reductions(prog: &mut TapeProgram) {
    if !prog.precision.is_empty() {
        return;
    }
    let n = prog.instrs.len();
    let guarded = prog.guarded_positions();
    let mut absorbed: Vec<(usize, u32, PreReduce)> = Vec::new();
    for i in 0..n.saturating_sub(1) {
        if guarded[i] || guarded[i + 1] || prog.section_of(i) != prog.section_of(i + 1) {
            continue;
        }
        let Instr::Reduce {
            op,
            init,
            src,
            axes,
            reverse: false,
            src_shape,
            out,
        } = &prog.instrs[i]
        else {
            continue;
        };
        let Instr::Fused { spec } = &prog.instrs[i + 1] else {
            continue;
        };
        let k = axes.len();
        if k == 0 || k >= src_shape.len() || !axes.iter().enumerate().all(|(j, &a)| a as usize == j)
        {
            continue;
        }
        let fs = &prog.fused[*spec as usize];
        if fs.reduce.is_some()
            || fs.pre.is_some()
            || !fs.scan_fuse.is_empty()
            || fs.micro.iter().any(|m| matches!(m, MicroOp::Scan { .. }))
            || fs.shape[..] != src_shape[..]
            || !fs
                .inputs
                .iter()
                .any(|inp| matches!(inp.src, SrcRef::Slot(s) if s == *out))
        {
            continue;
        }
        let rows: usize = src_shape[..k].iter().product();
        let n_inner: usize = src_shape[k..].iter().product();
        if prog.slots[*out as usize].scalar
            || prog.slots[*out as usize].elems() != n_inner
            || !reads_own_columns(fs, *out, k, n_inner)
        {
            continue;
        }
        absorbed.push((
            i,
            *spec,
            PreReduce {
                op: *op,
                init: *init,
                src: *src,
                rows,
                n_inner,
                out: *out,
            },
        ));
    }
    if absorbed.is_empty() {
        return;
    }
    let mut drop = vec![false; n];
    for (i, spec, pre) in absorbed {
        drop[i] = true;
        prog.fused[spec as usize].pre = Some(pre);
    }
    let (nc, ns) = (prog.n_const as usize, prog.n_segment as usize);
    prog.n_const -= drop[..nc].iter().filter(|&&d| d).count() as u32;
    prog.n_segment -= drop[nc..nc + ns].iter().filter(|&&d| d).count() as u32;
    let mut k = 0usize;
    prog.instrs.retain(|_| {
        k += 1;
        !drop[k - 1]
    });
    let mut k = 0usize;
    prog.provenance.retain(|_| {
        k += 1;
        !drop[k - 1]
    });
    prog.fuse_stats.instrs_after = prog.instrs.len();
}

/// Whether every element of group `fs` that reads slot `out` (`n_inner`
/// elements, the columns of the group's box past its `lead` leading axes)
/// reads it at its own column, `out[flat % n_inner]`, or reads the zero
/// ghost. An absorbed reduction is folded and consumed one range of columns
/// at a time, so a read of any other column (a shifted read, a gather
/// through a data subscript) would see a cell another range has not folded
/// yet.
fn reads_own_columns(fs: &FusedSpec, out: SlotId, lead: usize, n_inner: usize) -> bool {
    fs.inputs
        .iter()
        .filter(|inp| matches!(inp.src, SrcRef::Slot(s) if s == out))
        .all(|inp| {
            if inp.index.is_some() {
                return false;
            }
            if let Some(g) = &inp.gather {
                return chunk_gather_own_columns(g, lead);
            }
            match inp.shifted_ix {
                // Aligned: `out` has the box's element count, one column per
                // element.
                None => true,
                Some(ix) => own_column_runs(
                    &fs.schedule.nodes,
                    ix as usize,
                    inp.elem_stride,
                    n_inner as i64,
                    0,
                    0,
                ),
            }
        })
}

/// [`reads_own_columns`] for a [`ChunkGather`], axis by axis: each of the
/// `lead` leading axes adds nothing to the source offset (the read is
/// broadcast along it), and each later axis maps every position it copies
/// to the same position at the box's own row-major stride, so the offset is
/// the element's column.
fn chunk_gather_own_columns(g: &ChunkGather, lead: usize) -> bool {
    let mut stride = 1i64;
    for a in (0..g.shape.len()).rev() {
        let ok = if a < lead {
            g.strides[a] == 0 || g.segs[a].iter().all(|&(_, len, so)| len == 1 && so == 0)
        } else {
            g.strides[a] == stride && g.segs[a].iter().all(|&(o, _, so)| so == o)
        };
        if !ok {
            return false;
        }
        stride *= g.shape[a] as i64;
    }
    true
}

/// [`reads_own_columns`] for shifted input `ix` over schedule `nodes`, every
/// run moved by `(out_d, in_d)` (see [`RunSchedule::for_each_run`]) without
/// expanding the repetitions: a run reads its own columns when its source
/// offset is its output offset's column and it stays within the columns,
/// one source element per element. A repetition that moves the output by
/// whole rows and the source not at all repeats the same columns, so its
/// body is checked once; one that moves both by the same step keeps every
/// run's source-minus-output difference, and its runs' columns move
/// linearly, so its first and last repetitions bound the rest.
fn own_column_runs(
    nodes: &[RunNode],
    ix: usize,
    stride: i64,
    n: i64,
    out_d: i64,
    in_d: i64,
) -> bool {
    let mut i = 0usize;
    while i < nodes.len() {
        match &nodes[i] {
            RunNode::Run(r) => {
                let src = r.in_off[ix];
                if src != GHOST_OFF {
                    let col = (r.out_off as i64 + out_d).rem_euclid(n);
                    let len = r.len as i64;
                    if src + in_d != col || (len > 1 && (stride != 1 || col + len > n)) {
                        return false;
                    }
                }
                i += 1;
            }
            RunNode::Repeat {
                count,
                body,
                out_step,
                in_step,
            } => {
                let body_nodes = &nodes[i + 1..i + 1 + *body as usize];
                let s = in_step[ix];
                let at = |k: i64| {
                    own_column_runs(
                        body_nodes,
                        ix,
                        stride,
                        n,
                        out_d + k * out_step,
                        in_d + k * s,
                    )
                };
                let last = *count as i64 - 1;
                let ok = if s == 0 && out_step % n == 0 {
                    at(0)
                } else if s == *out_step {
                    at(0) && at(last)
                } else {
                    false
                };
                if !ok {
                    return false;
                }
                i += 1 + *body as usize;
            }
        }
    }
    true
}

/// The rewritten program under construction.
struct Sink {
    stats: FuseStats,
    fused: Vec<FusedSpec>,
    instrs: Vec<Instr>,
    prov: Vec<u32>,
    /// Parallel to `instrs` when the source program is precision-tagged.
    prec: Vec<Precision>,
    /// Step 4b histogram accumulators (weighted by group element count).
    micro_hist: FxHashMap<String, usize>,
    adj_hist: FxHashMap<String, usize>,
    chain_hist: FxHashMap<usize, usize>,
}

impl Sink {
    fn passthrough(&mut self, prog: &TapeProgram, i: usize) {
        if matches!(prog.instrs[i], Instr::Gather { .. }) {
            self.stats.n_gathers_kept += 1;
        }
        self.instrs.push(prog.instrs[i].clone());
        self.prov.push(prog.provenance[i]);
        if let Some(p) = prec_at(prog, i) {
            self.prec.push(p);
        }
    }
}

/// Instruction `i`'s precision tag, when the program carries them.
fn prec_at(prog: &TapeProgram, i: usize) -> Option<Precision> {
    prog.precision.get(i).copied()
}

/// The state [`fuse_section`] and its helpers share for the whole pass: the
/// (read-only) source program, the global per-slot reader index, the
/// rewritten-program [`Sink`], and the superop configuration — grouped so the
/// helpers stop threading four positional arguments alongside their real
/// inputs.
struct FuseCtx<'a> {
    prog: &'a TapeProgram,
    readers: &'a [SmallVec<[u32; 4]>],
    sink: Sink,
    cfg: SuperopCfg,
    /// Per instruction: a tuple instruction of a segmented-reduction chain
    /// ([`chained_tuple_instrs`]), left as it is.
    chained: Vec<bool>,
}

/// The elementwise tuple instructions right before each CONTINUOUS
/// `SegReduce` (a gather of each tuple's source, arithmetic over the tuple
/// box): the executor runs them with the reduction as one chain, block of
/// cells by block of cells, so their values stay in cache (`exec::chain`).
/// A fused group in their place would end the chain.
pub(super) fn chained_tuple_instrs(prog: &TapeProgram) -> Vec<bool> {
    let mut chained = vec![false; prog.instrs.len()];
    let cont = prog.section_range(Cadence::Continuous);
    for seg in cont.clone() {
        let Instr::SegReduce { table, .. } = &prog.instrs[seg] else {
            continue;
        };
        let Some(&n_pairs) = prog.seg_tables[*table as usize].rows.last() else {
            continue;
        };
        let tuple = |s: &SlotId| {
            let d = &prog.slots[*s as usize];
            !d.scalar && d.shape[..] == [n_pairs as usize]
        };
        let mut pc = seg;
        while pc > cont.start {
            let ins = &prog.instrs[pc - 1];
            let ok = match ins {
                Instr::TableGather { out, .. }
                | Instr::Bin { out, .. }
                | Instr::Un { out, .. }
                | Instr::Neg { out, .. } => tuple(out),
                _ => false,
            };
            if !ok {
                break;
            }
            chained[pc - 1] = true;
            pc -= 1;
        }
    }
    chained
}

/// Finalize one group: emit either its `Fused` instruction (plus any
/// re-materialized gathers) or, for trivial groups, the original
/// instructions verbatim.
fn flush_one(g: GBuilder, fx: &mut FuseCtx) {
    let prog = fx.prog;
    let has_compute = g.micro.iter().any(|m| !matches!(m, MicroOp::Mov { .. }));
    if g.member_instrs.len() < 2 || g.micro.is_empty() || !has_compute {
        for &ix in g.member_instrs.iter().chain(&g.deferred) {
            fx.sink.passthrough(prog, ix as usize);
        }
        return;
    }
    let GBuilder {
        shape,
        member_instrs,
        members,
        mut micro,
        mut inputs,
        shifted_segs,
        scalars,
        defs,
        folded,
        folded_index,
        prov,
        prec,
        reduce,
        deferred,
        ..
    } = g;

    // Folded gathers whose value is ALSO read outside the group must stay
    // materialized: re-emit the Gather and rewire the input to an aligned
    // read of its output slot.
    let mut kept_shifted: Vec<bool> = vec![true; shifted_segs.len()];
    for &(orig_ix, slot, input_ix) in &folded {
        let external = fx.readers[slot as usize]
            .iter()
            .any(|r| !members.contains(r));
        if external {
            fx.sink.passthrough(prog, orig_ix as usize);
            // A twin fold already rewired this shared input to its own
            // materialized gather, which holds the same values.
            let Some(old) = inputs[input_ix as usize].shifted_ix else {
                continue;
            };
            kept_shifted[old as usize] = false;
            inputs[input_ix as usize] = FusedInput {
                src: SrcRef::Slot(slot),
                shifted_ix: None,
                src_shape: shape.clone(),
                elem_stride: 1,
                load_reg: GroupIx::MAX,
                index: None,
                gather: None,
            };
        } else {
            fx.sink.stats.n_gathers_folded += 1;
        }
    }
    // The same for a folded data-subscript gather.
    for &(orig_ix, slot, input_ix) in &folded_index {
        let external = fx.readers[slot as usize]
            .iter()
            .any(|r| !members.contains(r));
        if external {
            fx.sink.passthrough(prog, orig_ix as usize);
            inputs[input_ix as usize] = FusedInput {
                src: SrcRef::Slot(slot),
                shifted_ix: None,
                src_shape: shape.clone(),
                elem_stride: 1,
                load_reg: GroupIx::MAX,
                index: None,
                gather: None,
            };
        } else {
            fx.sink.stats.n_gathers_folded += 1;
        }
    }
    // Compact the shifted-input table.
    let mut new_segs: Vec<ShiftedGeom> = Vec::new();
    let mut new_ix: Vec<GroupIx> = vec![GroupIx::MAX; kept_shifted.len()];
    for (i, segs) in shifted_segs.into_iter().enumerate() {
        if kept_shifted[i] {
            new_ix[i] = new_segs.len() as GroupIx;
            new_segs.push(segs);
        }
    }
    for inp in inputs.iter_mut() {
        if let Some(s) = inp.shifted_ix {
            inp.shifted_ix = Some(new_ix[s as usize]);
        }
    }

    // Live-outs. An absorbed reduction's register is live to the end like
    // one (so neither the superop peephole nor the register allocator can
    // retire it); it rides LAST in `outputs` until allocation has renamed it.
    let mut outputs: SmallVec<[(GroupIx, SlotId); 2]> = SmallVec::new();
    for &(slot, ssa) in &defs {
        let external = fx.readers[slot as usize]
            .iter()
            .any(|r| !members.contains(r));
        if external {
            outputs.push((ssa, slot));
        }
    }
    if let Some(r) = &reduce {
        outputs.push((r.ssa, r.out));
    }

    // Step 4b histograms (pre-superop, while the program is SSA): single-use
    // producer→consumer adjacency and maximal arith chain lengths, weighted by
    // the group's element count (= per-RHS element-op cost).
    let n_elems: usize = shape.iter().product::<usize>().max(1);
    {
        let uses = ssa_uses(&micro, &outputs);
        for i in 0..micro.len().saturating_sub(1) {
            if uses[i] == 1 && reads_reg(&micro[i + 1], i as GroupIx) {
                let key = format!("{}>{}", mop_label(&micro[i]), mop_label(&micro[i + 1]));
                *fx.sink.adj_hist.entry(key).or_insert(0) += n_elems;
            }
        }
        let arith = |op: &MicroOp| matches!(op, MicroOp::Bin { op, .. } if bin2_arith(*op));
        let mut i = 0usize;
        while i < micro.len() {
            if arith(&micro[i]) {
                let mut len = 1usize;
                while i + len < micro.len()
                    && arith(&micro[i + len])
                    && uses[i + len - 1] == 1
                    && reads_reg(&micro[i + len], (i + len - 1) as GroupIx)
                {
                    len += 1;
                }
                if len >= 2 {
                    *fx.sink.chain_hist.entry(len).or_insert(0) += n_elems;
                }
                i += len;
            } else {
                i += 1;
            }
        }
    }

    merge_superops(&mut micro, &mut outputs, fx.cfg);
    let scan_fuse = scan_fusions(&mut micro, &mut outputs);
    // A chunk-read gather with short rows in a group with little else to do
    // costs more per row than reading its box once: read it whole.
    if micro.len() < LIGHT_GROUP_OPS {
        for inp in inputs.iter_mut() {
            if let Some(g) = inp.gather.as_deref_mut()
                && g.shape.last().is_some_and(|&l| l < SHORT_CHUNK_GATHER_ROW)
            {
                g.whole = true;
            }
        }
    }
    for op in &micro {
        *fx.sink.micro_hist.entry(mop_label(op)).or_insert(0) += n_elems;
    }
    let n_regs = allocate_registers(&mut micro, &mut outputs);
    if reduce.is_some() {
        fx.sink.stats.n_reduces_folded += 1;
    }
    fx.sink.stats.n_scans_folded += micro
        .iter()
        .filter(|m| matches!(m, MicroOp::Scan { .. }))
        .count();
    let reduce = reduce.map(|r| {
        let (reg, slot) = outputs.pop().expect("the reduction rides last");
        debug_assert_eq!(slot, r.out);
        FusedReduce {
            reg,
            op: r.op,
            init: r.init,
            out: r.out,
            n_inner: r.n_inner,
        }
    });
    // Bin3 executes all-pointer: its scalar operands read from per-scalar
    // splat registers and ghost reads from a trailing zero register, all
    // appended after the load registers and filled once per call.
    let n_splat_regs = if micro.iter().any(|m| matches!(m, MicroOp::Bin3 { .. })) {
        scalars.len() as GroupIx + 1
    } else {
        0
    };
    let schedule = build_schedule(&shape, &new_segs);
    let interleave = reduce.as_ref().and_then(|r| {
        const MAX_INTERLEAVED: usize = 16;
        let n_pos = n_elems / r.n_inner;
        if !(2..=MAX_INTERLEAVED).contains(&n_pos)
            || schedule.n_runs != n_pos
            || micro.iter().any(|m| matches!(m, MicroOp::Scan { .. }))
        {
            return None;
        }
        let runs = schedule.expanded();
        runs.iter()
            .enumerate()
            .all(|(k, run)| run.out_off as usize == k * r.n_inner && run.len as usize == r.n_inner)
            .then_some(runs)
    });
    // Strided shifted inputs are pre-loaded into dedicated chunk registers
    // appended after the micro register file.
    let mut n_load_regs: GroupIx = 0;
    // A zero-stride input (constant along each run) is read as a per-run
    // scalar instead, except under the all-pointer Bin3 superop, which
    // needs every operand in a register.
    for inp in inputs.iter_mut() {
        if inp.index.is_some_and(|fi| fi.preloaded())
            || inp.gather.as_deref().is_some_and(|g| !g.whole)
            || inp.shifted_ix.is_some()
                && inp.elem_stride != 1
                && (inp.elem_stride != 0 || n_splat_regs > 0)
        {
            inp.load_reg = n_regs + n_load_regs;
            n_load_regs += 1;
        }
    }

    let n_members = member_instrs.len();
    fx.sink.stats.n_groups += 1;
    fx.sink.stats.n_member_instrs += n_members;
    let bucket = match n_members {
        0..=3 => 0,
        4..=7 => 1,
        8..=15 => 2,
        16..=31 => 3,
        32..=63 => 4,
        _ => 5,
    };
    fx.sink.stats.group_size_hist[bucket] += 1;

    let direct = direct_stores(&micro, &outputs, reduce.as_ref().map(|r| r.reg));
    let spec_ix = fx.sink.fused.len() as u32;
    fx.sink.fused.push(FusedSpec {
        shape,
        inputs,
        scalars,
        micro,
        n_regs,
        n_load_regs,
        n_splat_regs,
        outputs,
        direct,
        scan_fuse,
        schedule,
        reduce,
        interleave,
        pre: None,
        n_fused_instrs: n_members as u32,
        n_folded_gathers: (folded.len() + folded_index.len()) as u32,
    });
    fx.sink.instrs.push(Instr::Fused { spec: spec_ix });
    fx.sink.prov.push(prov);
    if let Some(p) = prec {
        fx.sink.prec.push(p);
    }
    for ix in deferred {
        fx.sink.passthrough(prog, ix as usize);
    }
}

fn fuse_section(fx: &mut FuseCtx, range: std::ops::Range<usize>) {
    let mut open: Vec<GBuilder> = Vec::new();

    // Flush every open group that DEFINES one of the slots `ins` reads,
    // except (optionally) the group with shape `keep_shape` that `ins` is
    // about to join. Returns the surviving groups untouched.
    fn flush_hazards(
        ins: &Instr,
        keep: Option<(&super::super::DimU, Option<Precision>)>,
        open: &mut Vec<GBuilder>,
        fx: &mut FuseCtx,
    ) {
        let mut hazard: Vec<usize> = Vec::new();
        ins.for_each_read(&fx.prog.tables(), |s| {
            for (gi, g) in open.iter().enumerate() {
                if keep != Some((&g.shape, g.prec)) && g.defines(s) && !hazard.contains(&gi) {
                    hazard.push(gi);
                }
            }
        });
        hazard.sort_unstable();
        for &gi in hazard.iter().rev() {
            let g = open.remove(gi);
            flush_one(g, fx);
        }
    }

    /// Find the open group with `shape`, or open a new one (flushing the
    /// oldest beyond the cap).
    fn find_or_open(
        open: &mut Vec<GBuilder>,
        shape: super::super::DimU,
        prov: u32,
        prec: Option<Precision>,
        fx: &mut FuseCtx,
    ) -> usize {
        if let Some(gi) = open.iter().position(|g| g.shape == shape && g.prec == prec) {
            if open[gi].has_room() {
                return gi;
            }
            let g = open.remove(gi);
            flush_one(g, fx);
        }
        if open.len() >= MAX_OPEN_GROUPS {
            let g = open.remove(0);
            flush_one(g, fx);
        }
        open.push(GBuilder::new(shape, prov, prec));
        open.len() - 1
    }

    let prog = fx.prog;
    // Nothing reads the observed map or `dy` during a call, so an `Export`
    // or a `DyWrite` only has to run before its section ends: one of a slot
    // an open group defines is held back until that group flushes instead of
    // flushing it early (which lets the next equation over the same box join
    // the group).
    let mut i = range.start;
    while i < range.end {
        let ins = &prog.instrs[i];
        if fx.chained[i] {
            flush_hazards(ins, None, &mut open, fx);
            fx.sink.passthrough(prog, i);
            i += 1;
            continue;
        }
        let deferrable = match ins {
            Instr::Export { slot, .. } => Some(*slot),
            Instr::DyWrite { write } => Some(prog.dy_writes[*write as usize].slot),
            _ => None,
        };
        if let Some(slot) = deferrable
            && let Some(gi) = open.iter().position(|g| g.defines(slot))
        {
            open[gi].deferred.push(i as u32);
            i += 1;
            continue;
        }

        // Hard barriers: a recurrence sweep and its body (copied verbatim —
        // the body runs once per cell, in order, and its cells are not
        // independent, so nothing in it may join a group: CONFORMANCE_SPEC
        // §5.19.2), conditional regions (copied verbatim — their skip counts
        // must stay exact and phi slots are double-defined inside).
        if let Instr::Sweep { spec } = ins {
            for g in open.drain(..) {
                flush_one(g, fx);
            }
            let body = prog.sweeps[*spec as usize].body_len as usize;
            let zone_end = (i + 1 + body).min(range.end);
            for j in i..zone_end {
                fx.sink.passthrough(prog, j);
            }
            i = zone_end;
            continue;
        }
        if let Instr::JmpIfZero {
            n_true, n_false, ..
        } = ins
        {
            for g in open.drain(..) {
                flush_one(g, fx);
            }
            let zone_end = (i + 1 + *n_true as usize + *n_false as usize).min(range.end);
            for j in i..zone_end {
                fx.sink.passthrough(prog, j);
            }
            i = zone_end;
            continue;
        }

        // A copy of a value an open group of the same box computes is that
        // value (a copy joining two branches sits in a conditional region,
        // copied verbatim above).
        if let Instr::Copy {
            a: Operand::Slot(s),
            out,
        } = ins
            && let Some(gi) = open.iter().position(|g| {
                g.shape == prog.slots[*out as usize].shape
                    && g.prec == prec_at(prog, i)
                    && matches!(g.val_of.get(s), Some(MRef::Reg(_)))
            })
        {
            open[gi].alias(i as u32, *s, *out);
            i += 1;
            continue;
        }

        // Foldable gather: absorb as a shifted input of its box's group.
        if let Instr::Gather { src, plan, out } = ins {
            let plan_ref = &prog.plans[*plan as usize];
            let out_desc = &prog.slots[*out as usize];
            let fold = gather_fold(plan_ref, &out_desc.shape);
            if let Some(GatherFold::Shifted(geom)) = fold {
                // An identity read of a value an open group of the same box
                // computes is that value: alias it instead of materializing
                // the group's output to read it back.
                if let SrcRef::Slot(s) = src
                    && plan_ref.src_shape == out_desc.shape
                    && is_identity_fold(&out_desc.shape, &geom)
                    && let Some(gi) = open.iter().position(|g| {
                        g.shape == out_desc.shape
                            && g.prec == prec_at(prog, i)
                            && matches!(g.val_of.get(s), Some(MRef::Reg(_)))
                    })
                {
                    open[gi].alias(i as u32, *s, *out);
                    i += 1;
                    continue;
                }
                // A gather whose SOURCE is defined by an open group forces
                // that group to materialize first.
                flush_hazards(ins, None, &mut open, fx);
                let shape = out_desc.shape.clone();
                let gi = find_or_open(&mut open, shape, prog.provenance[i], prec_at(prog, i), fx);
                open[gi].add_folded_gather(i as u32, *src, plan_ref, geom, *out);
                i += 1;
                continue;
            }
            if let Some(GatherFold::Chunk(strides)) = fold {
                flush_hazards(ins, None, &mut open, fx);
                let shape = out_desc.shape.clone();
                let gi = find_or_open(&mut open, shape, prog.provenance[i], prec_at(prog, i), fx);
                let mut segs = plan_ref.segs.clone();
                // `ChunkGather::fill` walks each axis's segments in output
                // order.
                for axis in segs.iter_mut() {
                    axis.sort_unstable();
                }
                let g = ChunkGather {
                    segs,
                    strides,
                    shape: plan_ref.shape.clone(),
                    whole: plan_ref
                        .shape
                        .last()
                        .is_none_or(|&l| l < MIN_CHUNK_GATHER_ROW),
                };
                open[gi].add_chunk_gather(i as u32, *src, g, &plan_ref.src_shape, *out);
                i += 1;
                continue;
            }
            // Unfoldable: pass through (with the usual read hazards).
            flush_hazards(ins, None, &mut open, fx);
            fx.sink.passthrough(prog, i);
            i += 1;
            continue;
        }

        // A data-subscript gather whose other source axes are fixed or read
        // along the box's innermost axis: a pre-loaded input of its box's
        // group, when its subscript is an array defined outside that group.
        if let Instr::IndexGather {
            src,
            idx,
            spec,
            out,
        } = ins
        {
            let spec_ref = &prog.index_gathers[*spec as usize];
            let out_desc = &prog.slots[*out as usize];
            if let Some(fold) = folded_index_of(spec_ref)
                && !out_desc.shape.is_empty()
            {
                let shape = out_desc.shape.clone();
                let prec = prec_at(prog, i);
                // The group over the box may keep computing while it supplies
                // the subscript — but not when it also computes the source,
                // which the gather reads at positions other than its own.
                let src_in_box_group = matches!(src, SrcRef::Slot(s)
                    if open.iter().any(|g| g.shape == shape && g.prec == prec && g.defines(*s)));
                if src_in_box_group {
                    flush_hazards(ins, None, &mut open, fx);
                } else {
                    flush_hazards(ins, Some((&shape, prec)), &mut open, fx);
                }
                let gi = find_or_open(&mut open, shape, prog.provenance[i], prec, fx);
                let src_shape = spec_ref.src_shape.clone();
                if open[gi]
                    .add_folded_index_gather(i as u32, *src, idx, fold, &src_shape, *out, prog)
                {
                    i += 1;
                    continue;
                }
            }
            flush_hazards(ins, None, &mut open, fx);
            fx.sink.passthrough(prog, i);
            i += 1;
            continue;
        }

        // Elementwise array-out instruction?
        let elem_shape: Option<super::super::DimU> = match ins {
            Instr::Bin { out, .. }
            | Instr::Un { out, .. }
            | Instr::Neg { out, .. }
            | Instr::Select { out, .. }
            | Instr::Fill { out, .. }
            | Instr::Copy { out, .. } => {
                let d = &prog.slots[*out as usize];
                if !d.scalar && !d.shape.is_empty() && !d.shape.contains(&0) {
                    Some(d.shape.clone())
                } else {
                    None
                }
            }
            _ => None,
        };
        if let Some(shape) = elem_shape {
            let prec = prec_at(prog, i);
            flush_hazards(ins, Some((&shape, prec)), &mut open, fx);
            let gi = find_or_open(&mut open, shape.clone(), prog.provenance[i], prec, fx);
            if open[gi].try_add(i as u32, ins, prog) {
                i += 1;
                continue;
            }
            // Unfusable operand (an observed read): the target group may
            // still hold slots this instruction reads — flush it too.
            flush_hazards(ins, None, &mut open, fx);
            fx.sink.passthrough(prog, i);
            i += 1;
            continue;
        }

        // A scan along any axis joins the group of its box.
        if let Instr::Scan { src_shape, .. } = ins
            && !src_shape.is_empty()
            && !src_shape.contains(&0)
        {
            let prec = prec_at(prog, i);
            flush_hazards(ins, Some((src_shape, prec)), &mut open, fx);
            let gi = find_or_open(&mut open, src_shape.clone(), prog.provenance[i], prec, fx);
            if open[gi].has_room() && open[gi].try_add_scan(i as u32, ins, prog) {
                i += 1;
                continue;
            }
            flush_hazards(ins, None, &mut open, fx);
            fx.sink.passthrough(prog, i);
            i += 1;
            continue;
        }

        // A reduction of an open group's value folds inside that group.
        if let Instr::Reduce {
            src: SrcRef::Slot(src),
            ..
        } = ins
            && let Some(gi) = open.iter().position(|g| g.defines(*src))
            && open[gi].prec == prec_at(prog, i)
            && open[gi].try_absorb_reduce(i as u32, ins, fx.readers)
        {
            let g = open.remove(gi);
            flush_one(g, fx);
            i += 1;
            continue;
        }

        // Everything else passes through, flushing any group it reads from.
        flush_hazards(ins, None, &mut open, fx);
        fx.sink.passthrough(prog, i);
        i += 1;
    }
    for g in open.drain(..) {
        flush_one(g, fx);
    }
}

/// The [`FoldedIndex`] of an [`Instr::IndexGather`] (its `by` still unset),
/// when every source axis but the data axis is fixed or read along the
/// output box's innermost axis; `None` otherwise.
fn folded_index_of(spec: &IndexGatherSpec) -> Option<FoldedIndex> {
    let m = spec.src_shape.len();
    let nd = spec.shape.len();
    let mut strides = vec![1usize; m];
    for d in (0..m.saturating_sub(1)).rev() {
        strides[d] = strides[d + 1] * spec.src_shape[d + 1];
    }
    let mut fold = FoldedIndex {
        by: 0,
        n: 0,
        s_data: 1,
        base: 0,
        s_inner: 0,
        period: 1,
        src_len: spec.src_shape.iter().product(),
    };
    let mut inner_seen = false;
    for (d, ax) in spec.axes.iter().enumerate() {
        match *ax {
            GatherAxis::Data => {
                fold.n = spec.src_shape[d];
                fold.s_data = strides[d];
            }
            GatherAxis::Fixed(i) => fold.base += strides[d] * i,
            GatherAxis::Affine { axis, off } => {
                if nd == 0 || axis as usize != nd - 1 || inner_seen || off < 0 {
                    return None;
                }
                inner_seen = true;
                fold.base += strides[d] * off as usize;
                fold.s_inner = strides[d];
                fold.period = spec.shape[nd - 1].max(1);
            }
        }
    }
    Some(fold)
}

#[cfg(test)]
mod run_schedule_tests {
    use super::*;

    /// The schedule as it was first written: every (tile, row) piece
    /// materialized, then sorted and coalesced. The built schedule, expanded,
    /// must reproduce it run for run.
    fn reference_build_runs(shape: &[usize], shifted: &[ShiftedGeom]) -> Vec<FusedRun> {
        let n_elems: usize = shape.iter().product::<usize>().max(1);
        if shifted.is_empty() {
            return vec![FusedRun {
                out_off: 0,
                len: n_elems as u32,
                in_off: SmallVec::new(),
            }];
        }
        let nd = shape.len();
        let strides = rm_strides(shape);
        let n_shift = shifted.len();
        // Per-input element stride (coalescing contiguity is offset + len*stride).
        let elem_stride: SmallVec<[i64; 2]> =
            shifted.iter().map(ShiftedGeom::elem_stride).collect();
        let seg_strides: SmallVec<[Option<&[i64]>; 2]> = shifted
            .iter()
            .map(|g| match g {
                ShiftedGeom::Segs { strides, .. } => Some(&strides[..]),
                ShiftedGeom::Linear { .. } => None,
            })
            .collect();
        // Per axis: merged interval list
        // [(start, len, per-input Option<source start position on this axis>)]
        // (`None` for Linear inputs, which impose no cuts).
        #[allow(clippy::type_complexity)]
        let mut axis_ivals: Vec<Vec<(usize, usize, SmallVec<[Option<i64>; 2]>)>> =
            Vec::with_capacity(nd);
        for d in 0..nd {
            let mut cuts: BTreeSet<usize> = BTreeSet::new();
            cuts.insert(0);
            cuts.insert(shape[d]);
            for g in shifted {
                if let ShiftedGeom::Segs { segs, .. } = g {
                    for &(o, l, _) in &segs[d] {
                        cuts.insert(o);
                        cuts.insert(o + l);
                    }
                }
            }
            let cuts: Vec<usize> = cuts.into_iter().collect();
            let mut ivals = Vec::with_capacity(cuts.len() - 1);
            for w in cuts.windows(2) {
                let (a, b) = (w[0], w[1]);
                let mut src_start: SmallVec<[Option<i64>; 2]> = SmallVec::new();
                for g in shifted {
                    let mut pos = None;
                    if let ShiftedGeom::Segs { segs, .. } = g {
                        for &(o, l, so) in &segs[d] {
                            if a >= o && b <= o + l {
                                pos = Some(so as i64 + (a as i64 - o as i64));
                                break;
                            }
                        }
                    }
                    src_start.push(pos);
                }
                ivals.push((a, b - a, src_start));
            }
            axis_ivals.push(ivals);
        }

        // Cartesian product over per-axis intervals (tiles); each tile emits its
        // rows (inner-axis contiguous pieces) as runs.
        let mut runs: Vec<FusedRun> = Vec::new();
        let mut pick: SmallVec<[usize; 4]> = SmallVec::from_elem(0usize, nd);
        loop {
            // Per Segs-input per-axis source base of this tile (ghost if any axis
            // is uncovered).
            let mut ghost: SmallVec<[bool; 2]> = SmallVec::from_elem(false, n_shift);
            let mut src_base: SmallVec<[i64; 2]> = shifted
                .iter()
                .map(|g| match g {
                    ShiftedGeom::Segs { base, .. } => *base,
                    ShiftedGeom::Linear { .. } => 0,
                })
                .collect();
            for d in 0..nd {
                let (_, _, ref starts) = axis_ivals[d][pick[d]];
                for j in 0..n_shift {
                    if let Some(sstr) = &seg_strides[j] {
                        match starts[j] {
                            Some(p) => src_base[j] += p * sstr[d],
                            None => ghost[j] = true,
                        }
                    }
                }
            }
            // Iterate the tile's rows: odometer over axes 0..nd-1 inside the
            // tile, inner axis (nd-1) is one contiguous piece per row.
            let inner = &axis_ivals[nd - 1][pick[nd - 1]];
            let (inner_start, inner_len) = (inner.0, inner.1);
            let mut row: SmallVec<[usize; 4]> = SmallVec::from_elem(0usize, nd.saturating_sub(1));
            loop {
                let mut off = inner_start as i64 * strides[nd - 1];
                for d in 0..nd - 1 {
                    let (st, _, _) = axis_ivals[d][pick[d]];
                    off += (st + row[d]) as i64 * strides[d];
                }
                let in_off: SmallVec<[i64; 2]> = (0..n_shift)
                    .map(|j| match &shifted[j] {
                        ShiftedGeom::Linear { a, b } => a * off + b,
                        ShiftedGeom::Segs { .. } => {
                            if ghost[j] {
                                GHOST_OFF
                            } else {
                                let sstr = seg_strides[j].expect("segs have strides");
                                let mut so = src_base[j];
                                for d in 0..nd - 1 {
                                    so += row[d] as i64 * sstr[d];
                                }
                                so
                            }
                        }
                    })
                    .collect();
                runs.push(FusedRun {
                    out_off: off as u32,
                    len: inner_len as u32,
                    in_off,
                });
                // advance row odometer (innermost leading axis fastest)
                let mut d = nd - 1;
                let mut done = true;
                while d > 0 {
                    d -= 1;
                    row[d] += 1;
                    if row[d] < axis_ivals[d][pick[d]].1 {
                        done = false;
                        break;
                    }
                    row[d] = 0;
                }
                if nd == 1 || done {
                    break;
                }
            }
            // advance tile odometer
            let mut d = nd;
            let mut done = true;
            while d > 0 {
                d -= 1;
                pick[d] += 1;
                if pick[d] < axis_ivals[d].len() {
                    done = false;
                    break;
                }
                pick[d] = 0;
            }
            if done {
                break;
            }
        }

        // Sort ascending and coalesce adjacent runs.
        runs.sort_by_key(|r| r.out_off);
        let mut out: Vec<FusedRun> = Vec::with_capacity(runs.len());
        for r in runs {
            if let Some(last) = out.last_mut() {
                let contiguous =
                    last.out_off + last.len == r.out_off
                        && last.in_off.iter().zip(r.in_off.iter()).enumerate().all(
                            |(j, (&a, &b))| {
                                (a == GHOST_OFF && b == GHOST_OFF)
                                    || (a != GHOST_OFF
                                        && b != GHOST_OFF
                                        && a + last.len as i64 * elem_stride[j] == b)
                            },
                        );
                if contiguous {
                    last.len += r.len;
                    continue;
                }
            }
            out.push(r);
        }
        debug_assert_eq!(
            out.iter().map(|r| r.len as usize).sum::<usize>(),
            n_elems,
            "run schedule must tile the box"
        );
        out
    }

    /// A small deterministic generator, so the cases are reproducible.
    struct Lcg(u64);
    impl Lcg {
        fn next(&mut self, n: u64) -> u64 {
            self.0 = self
                .0
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            (self.0 >> 33) % n
        }
    }

    /// One axis of a shifted read: a constant shift into a source of the same
    /// extent (the uncovered end is ghost), a periodic wrap (two segments), or
    /// an interior-subset read of a larger source.
    fn axis_segs(rng: &mut Lcg, n: usize) -> (SmallVec<[(usize, usize, usize); 2]>, usize) {
        let mut segs = SmallVec::new();
        match rng.next(4) {
            0 => {
                let s = rng.next(5) as i64 - 2;
                let lo = (-s).max(0) as usize;
                let hi = (n as i64 - s).min(n as i64).max(lo as i64) as usize;
                if hi > lo {
                    segs.push((lo, hi - lo, (lo as i64 + s) as usize));
                }
                (segs, n)
            }
            1 if n > 1 => {
                let s = 1 + rng.next(n as u64 - 1) as usize;
                segs.push((0, n - s, s));
                segs.push((n - s, s, 0));
                (segs, n)
            }
            2 => {
                let pad = rng.next(3) as usize;
                segs.push((0, n, pad));
                (segs, n + 2 * pad)
            }
            _ => {
                segs.push((0, n, 0));
                (segs, n)
            }
        }
    }

    /// A regular sequence of rows is one repetition, not one run per row: a
    /// contraction term read by the contracted index alone (`e[j]` over a
    /// `j x i` box, constant along each row and one source element per row)
    /// and a shift along the inner axis of a deep box both keep a schedule
    /// whose size does not depend on the extents, while it still executes
    /// (and counts, for the fold test) one run per row.
    #[test]
    fn a_regular_row_pattern_is_one_repetition() {
        for n in [100usize, 3162] {
            let geom = ShiftedGeom::Segs {
                segs: SmallVec::from_vec(vec![
                    SmallVec::from_slice(&[(0, n, 0)]),
                    SmallVec::from_slice(&[(0, n, 0)]),
                ]),
                strides: SmallVec::from_slice(&[1, 0]),
                base: 0,
            };
            let s = build_schedule(&[n, n], std::slice::from_ref(&geom));
            assert_eq!(s.n_runs, n);
            assert_eq!(s.nodes.len(), 2, "{:?}", s.nodes);
            assert!(fold_runs_acceptable(&[n, n], &geom));
        }
        for n in [16usize, 40] {
            let shape = [n, n, n];
            let geom = ShiftedGeom::Segs {
                segs: SmallVec::from_vec(vec![
                    SmallVec::from_slice(&[(0, n, 0)]),
                    SmallVec::from_slice(&[(0, n, 0)]),
                    SmallVec::from_slice(&[(0, n - 1, 1)]),
                ]),
                strides: rm_strides(&shape),
                base: 0,
            };
            let s = build_schedule(&shape, std::slice::from_ref(&geom));
            assert_eq!(
                s.n_runs,
                reference_build_runs(&shape, std::slice::from_ref(&geom)).len()
            );
            assert!(s.nodes.len() <= 8, "{} nodes at n = {n}", s.nodes.len());
        }
    }

    #[test]
    fn the_schedule_expands_to_the_sorted_and_coalesced_one() {
        let mut rng = Lcg(0x5eed);
        for case in 0..3000 {
            let nd = 1 + rng.next(4) as usize;
            let shape: Vec<usize> = (0..nd).map(|_| 1 + rng.next(7) as usize).collect();
            let n_in = 1 + rng.next(4) as usize;
            let shifted: Vec<ShiftedGeom> = (0..n_in)
                .map(|_| {
                    if rng.next(6) == 0 {
                        ShiftedGeom::Linear {
                            a: 1 + rng.next(3) as i64,
                            b: rng.next(5) as i64,
                        }
                    } else {
                        // A same-rank source (row-major strides over its own
                        // extents), with now and then an axis broadcast
                        // (stride 0) and a fixed-axis base offset.
                        let mut segs: AxisSegs = SmallVec::new();
                        let mut src: Vec<usize> = Vec::new();
                        for &n in &shape {
                            let (s, m) = axis_segs(&mut rng, n);
                            segs.push(s);
                            src.push(m);
                        }
                        let mut strides: SmallVec<[i64; 4]> = rm_strides(&src);
                        for st in strides.iter_mut() {
                            if rng.next(5) == 0 {
                                *st = 0;
                            }
                        }
                        ShiftedGeom::Segs {
                            segs,
                            strides,
                            base: rng.next(4) as i64,
                        }
                    }
                })
                .collect();
            let want = reference_build_runs(&shape, &shifted);
            let sched = build_schedule(&shape, &shifted);
            let got = sched.expanded();
            let key = |r: &FusedRun| (r.out_off, r.len, r.in_off.to_vec());
            assert_eq!(
                got.iter().map(key).collect::<Vec<_>>(),
                want.iter().map(key).collect::<Vec<_>>(),
                "case {case}: shape {shape:?}"
            );
            assert_eq!(sched.n_runs, want.len(), "case {case}: run count");
            // The fold test, per input, against the rule it states.
            for g in &shifted {
                let runs = reference_build_runs(&shape, std::slice::from_ref(g)).len();
                let n_elems: usize = shape.iter().product::<usize>().max(1);
                assert_eq!(
                    fold_runs_acceptable(&shape, g),
                    runs <= 8 || n_elems / runs >= 64,
                    "case {case}: fold test, shape {shape:?}"
                );
            }
        }
    }
}