zyx 0.17.0

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0

//! Kernel Intermediate Representation for building custom compute kernels.
//!
//! This module provides the IR builder API for constructing custom kernels
//! that can be compiled and executed on any backend (CPU, CUDA, Vulkan, etc.).
//!
//! # Quick start — WMMA matrix multiply
//!
//! Tensor-core matmul using the m16n8k8 WMMA instruction with one warp per 16×8 tile.
//! Requires CUDA with tensor cores (compute capability ≥ 7.0).
//!
//! ```rust
//! use zyx::kernel::{Dev, Kernel, MMADType, MMADims, MMALayout, MemLayout, MemScope, ParamKind};
//! use zyx::DType;
//!
//! let (m, n, k) = (1024, 1024, 1024);
//! let mut kernel = Kernel::new(Dev::Auto);
//!
//! let a_shape = kernel.add_shape(&[m, k]);
//! let b_shape = kernel.add_shape(&[k, n]);
//! let c_shape = kernel.add_shape(&[m, n]);
//! let a_buf = kernel.param(DType::F16);
//! let b_buf = kernel.param(DType::F16);
//! let c_buf = kernel.param_mut(DType::F32);
//!
//! let glen_x = kernel.const_idx(m / 16);
//! let glen_y = kernel.const_idx(n / 8);
//! let gidx = kernel.group_range(0, glen_x);
//! let gidy = kernel.group_range(1, glen_y);
//! let wid = kernel.local_range(0, 32);
//!
//! let [c0, c1, c2, c4, c8, c16] = kernel.const_idxs([0u32, 1, 2, 4, 8, 16]);
//! let n_const = kernel.const_idx(n);
//! let k_const = kernel.const_idx(k);
//!
//! let row_in_tile = kernel.div(wid, c4);
//! let sub_col = kernel.mod_(wid, c4);
//! let col_in_tile = kernel.mul(sub_col, c2);
//!
//! let a_row = kernel.mad(gidx, c16, row_in_tile);
//! let b_col = kernel.mad(gidy, c8, row_in_tile);
//! let tile_base_col = kernel.mul(gidy, c8);
//!
//! let acc = kernel.storage(DType::F32, MemScope::Register, 4);
//! let zf = kernel.const_val(0.0f32);
//! let zero_acc = kernel.stack(&[zf, zf, zf, zf]);
//! kernel.store_vector(acc, zero_acc, c0, 4);
//!
//! let k_div_8 = kernel.const_idx(k / 8);
//! kernel.loop_over(k_div_8, |kernel, k_loop| {
//! let k_off = kernel.mul(k_loop, c8);
//!
//! let a_base = kernel.mad(a_row, k_const, k_off);
//! let a_base = kernel.add(a_base, col_in_tile);
//! let a_load_0 = kernel.load(a_buf, a_base);
//! let a_base_p1 = kernel.add(a_base, c1);
//! let a_load_1 = kernel.load(a_buf, a_base_p1);
//! let a_base2 = kernel.mad(c8, k_const, a_base);
//! let a_load_2 = kernel.load(a_buf, a_base2);
//! let a_base2_p1 = kernel.add(a_base2, c1);
//! let a_load_3 = kernel.load(a_buf, a_base2_p1);
//! let a_frag = kernel.stack(&[a_load_0, a_load_1, a_load_2, a_load_3]);
//!
//! let b_row = kernel.add(k_off, col_in_tile);
//! let b_base = kernel.mad(b_row, n_const, b_col);
//! let b_load_0 = kernel.load(b_buf, b_base);
//! let b_base_n = kernel.add(b_base, n_const);
//! let b_load_1 = kernel.load(b_buf, b_base_n);
//! let b_frag = kernel.stack(&[b_load_0, b_load_1]);
//!
//! let acc_old = kernel.load_vector(acc, c0, 4);
//! let acc_new = kernel.wmma(
//!     MMADims::m16n8k8, MMALayout::row_col, MMADType::f16_f16_f16_f32,
//!     a_frag, b_frag, acc_old,
//! );
//! kernel.store_vector(acc, acc_new, c0, 4);
//! });
//!
//! let acc_final = kernel.load_vector(acc, c0, 4);
//! let [co, c1v, c2v, c3v] = kernel.devectorize(acc_final);
//!
//! let c_col = kernel.add(tile_base_col, col_in_tile);
//! let c_base = kernel.mad(a_row, n_const, c_col);
//! kernel.store(c_buf, co, c_base);
//! let c_base_p1 = kernel.add(c_base, c1);
//! kernel.store(c_buf, c1v, c_base_p1);
//! let c_base2 = kernel.mad(c8, n_const, c_base);
//! kernel.store(c_buf, c2v, c_base2);
//! let c_base2_p1 = kernel.add(c_base2, c1);
//! kernel.store(c_buf, c3v, c_base2_p1);
//!
//! // kernel.compile()?;  // requires CUDA with tensor cores
//! ```

pub use crate::backend::{Dev, DeviceInfo};
pub use custom::{Acc, CompiledKernel, LocalPartition, Partition};
pub use ops::{BOp, MMADType, MMADims, MMALayout, OpId, ParamKind, TileDim};
pub(crate) use ops::{MoveOp, Op, OpLinked, RangeKind, UOp};

use crate::{DType, Map, Set, dtype::Constant, shape::Dim, slab::Slab};
use nanoserde::{DeBin, SerBin};
use std::collections::BTreeMap;
use std::sync::Arc;
use std::{hash::BuildHasherDefault, hash::Hash};

mod algebraic;
pub mod autotune;
mod coarsen;
mod cost;
mod custom;
mod debug;
mod fold_constants;
mod fold_loops;
mod fuse;
mod instr_sched;
mod licm;
mod linearize;
mod local_reduce;
mod merge_loops;
mod mma;
mod ops;
mod pad_range;
mod predict_cost;
mod split_loops;
mod tenstorrent;
mod transforms;
mod unroll_loops;
mod vectorize;
mod verify;

// TODO later make this dynamic u32 or u64 depending on max range
/// Type used for indexing into arrays within kernels.
pub(crate) const IDX_T: DType = DType::I64;

/// Kernel builder for constructing custom compute kernels.
///
/// This struct represents a kernel in the intermediate representation (IR)
/// that can be compiled and executed on any backend (CPU, CUDA, Vulkan, etc.).
///
/// The kernel IR supports:
/// - Element-wise operations (add, mul, sin, exp, etc.)
/// - Reductions (sum, max, etc.)
/// - Memory operations (load, store)
/// - Control flow (loops, conditionals)
/// - Tensor transformations (reshape, permute, expand, pad)
///
/// # Example
///
/// Build a kernel that computes `sin(x) + cos(x)` element-wise:
///
/// ```
/// use zyx::kernel::{Kernel, MemLayout, Dev, ParamKind};
/// use zyx::DType;
///
/// let mut kernel = Kernel::new(Dev::Auto);
/// let n = 256;
/// let shape = kernel.add_shape(&[n]);
/// let len = kernel.const_idx(n);
/// let inp = kernel.param(DType::F32);
/// let gidx = kernel.group_range(0, len);
/// let loaded = kernel.load(inp, gidx);
/// let s = kernel.sin(loaded);
/// let c = kernel.cos(loaded);
/// let result = kernel.add(s, c);
/// let out = kernel.param_mut(DType::F32);
/// kernel.store(out, result, gidx);
/// ```
///
/// # Compile
///
/// Build a kernel using fused multiply-add and compile it:
///
/// ```
/// use zyx::kernel::{Kernel, MemLayout, Dev, ParamKind};
/// use zyx::{DType, Tensor, ZyxError};
///
/// let mut kernel = Kernel::new(Dev::Auto);
/// let n = 4;
/// let shape = kernel.add_shape(&[n]);
/// let inp = kernel.param(DType::F32);
/// let len = kernel.const_idx(n);
/// let gidx = kernel.group_range(0, len);
/// let loaded = kernel.load(inp, gidx);
/// let result = kernel.mad(loaded, loaded, loaded); // x*x + x
/// let out = kernel.param_mut(DType::F32);
/// kernel.store(out, result, gidx);
///
/// let compiled = kernel.compile()?;
/// let x = Tensor::from([1.0f32, 2.0, 3.0, 4.0]);
/// let result = compiled.forward(&[&x], vec![[n]])?;
/// let data: Vec<f32> = result.into_iter().next().unwrap().try_into()?;
/// assert_eq!(data, vec![2.0, 6.0, 12.0, 20.0]);
/// # Ok::<_, ZyxError>(())
/// ```
#[derive(Debug, Clone)]
pub struct Kernel {
    /// Operation slab containing the kernel IR.
    pub(crate) ops: Slab<OpId, OpLinked>,
    /// Head of the operation linked list.
    pub(crate) head: OpId,
    /// Tail of the operation linked list.
    pub(crate) tail: OpId,
    /// Target device for compilation.
    pub(crate) dev: Dev,
    /// Snapshot of the device's [`DeviceInfo`] taken when a device is bound —
    /// the warp size and limits live here so kernel passes never need to lock
    /// the RT (it is not reentrant). Shared via `Arc` so kernel clones
    /// (autotune creates thousands) don't re-allocate. `None` only for
    /// placeholder kernels (`Dev::Auto`) before their device is bound.
    /// Hardware properties are fixed, so the snapshot cannot go stale.
    pub(crate) dev_info: Option<Arc<DeviceInfo>>,
    /// Memoized [`Self::shape_ids`] results. Only valid pre-linearization —
    /// [`linearize`](crate::kernel::linearize) clears it so kernels stay
    /// lightweight during autotuning.
    pub(crate) shape_cache: Map<OpId, Vec<OpId>>,
}

/// Scope for memory.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, SerBin, DeBin)]
pub enum MemScope {
    /// Global memory scope (DRAM).
    Global,
    /// Local memory scope (SRAM).
    Local,
    /// Register scope (registers).
    Register,
    /// Circular buffer, SRAM (tenstorrent)
    Circular,
}

/// Memory layout for kernel operations.
///
/// Specifies how data is laid out in memory for efficient access.
///
/// The layout also determines how much a **single** store / load op
/// writes (reads): it is the *indexing granularity* of the op, not the rank
/// or scalarity of the value being moved. `Scalar` means the op writes one
/// single element per op; `Vector` writes a full vector with a single op;
/// `Tile` writes a large tile with a single op. A multi-element tensor can
/// perfectly well be stored through repeated `Scalar`-layout store ops (one
/// per element) — that says nothing about the value's shape.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, SerBin, DeBin)]
pub enum MemLayout {
    /// Scalar layout: one element per memory location
    Scalar,
    /// Vector layout: vector of size `x`
    Vector(u16),
    /// Tile layout: tile of `x` × `y` elements with stride
    Tile {
        /// Width of the tile
        x: u16,
        /// Height of the tile
        y: u16,
        /// Stride between tiles
        stride: u32,
    },
}

impl PartialEq for Kernel {
    fn eq(&self, other: &Self) -> bool {
        self.ops == other.ops && self.head == other.head && self.dev == other.dev
    }
}

impl Eq for Kernel {}

impl SerBin for Kernel {
    fn ser_bin(&self, output: &mut Vec<u8>) {
        self.ops.ser_bin(output);
        self.head.ser_bin(output);
        self.tail.ser_bin(output);
    }
}

/*impl DeBin for Kernel {
    fn de_bin(offset: &mut usize, bytes: &[u8]) -> Result<Self, nanoserde::DeBinErr> {
        let ops = Slab::<OpId, OpNode>::de_bin(offset, bytes)?;
        let start = OpId::de_bin(offset, bytes)?;
        let end = OpId::de_bin(offset, bytes)?;
        Ok(Self { head: start, tail: end, ops, device_id: Dev::Auto })
    }
}*/

impl Hash for Kernel {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.ops.hash(state);
        self.dev.hash(state);
    }
}

// Custom kernel machinery
impl Kernel {
    /// Creates an empty kernel bound to the given device.
    ///
    /// `Dev::Auto` (placeholder kernels) gets no [`DeviceInfo`] — it is
    /// bound together with the real device before compilation
    /// (`kernel.device_id = dev` sites must set both).
    ///
    /// Takes the device info as an argument and never locks the RT: this
    /// constructor runs while callers already hold the RT lock (it is not
    /// reentrant). Pass `None` for late-bound placeholders (`Dev::Auto`).
    pub(crate) fn from_device_id(dev: Dev, dev_info: Option<Arc<DeviceInfo>>) -> Self {
        Self { ops: Slab::new(), head: OpId::NULL, tail: OpId::NULL, dev, dev_info, shape_cache: Map::default() }
    }

    /// The device info snapshot bound to this kernel.
    ///
    /// # Panics
    ///
    /// If the kernel has no device bound (placeholder kernels with
    /// `Dev::Auto` that were never compiled).
    pub(crate) fn dev_info(&self) -> &DeviceInfo {
        self.dev_info.as_ref().expect("kernel has no device bound (Dev::Auto placeholder)")
    }

    /// Compute dtypes and reference counts for all operations.
    pub(crate) fn compute_dtypes_and_rcs(&self) -> (Map<OpId, (DType, MemLayout)>, Map<OpId, u32>) {
        let mut rcs: Map<OpId, u32> = Map::with_capacity_and_hasher(self.ops.len().into(), BuildHasherDefault::new());
        let mut dtypes: Map<OpId, (DType, MemLayout)> = Map::with_capacity_and_hasher(100, BuildHasherDefault::new());

        let mut op_id = self.head;
        for _ in 0..10_000 {
            if op_id.is_null() {
                break;
            }
            match self.ops[op_id].op {
                Op::Move { .. } | Op::Reduce { .. } | Op::ReduceTile { .. } => {
                    unreachable!()
                }
                Op::Const(x) => {
                    dtypes.insert(op_id, (x.dtype(), MemLayout::Scalar));
                }
                Op::Param { dtype, .. } => {
                    dtypes.insert(op_id, (dtype, MemLayout::Scalar));
                }
                Op::Storage { dtype, .. } => {
                    dtypes.insert(op_id, (dtype, MemLayout::Scalar));
                }
                Op::Load { src, index, layout } => {
                    dtypes.insert(op_id, (dtypes[&src].0, layout));
                    *rcs.entry(index).or_insert(0) += 1;
                }
                Op::Store { dst, src: x, index, layout } => {
                    debug_assert_eq!(dtypes[&x].1, layout);
                    dtypes.insert(op_id, dtypes[&x]);
                    *rcs.entry(dst).or_insert(0) += 1;
                    *rcs.entry(x).or_insert(0) += 1;
                    *rcs.entry(index).or_insert(0) += 1;
                }
                Op::Cast { x, dtype } => {
                    dtypes.insert(op_id, (dtype, dtypes[&x].1));
                    *rcs.entry(x).or_insert(0) += 1;
                }
                Op::Bitcast { x, dtype } => {
                    dtypes.insert(op_id, (dtype, dtypes[&x].1));
                    *rcs.entry(x).or_insert(0) += 1;
                }
                Op::Unary { x, .. } => {
                    dtypes.insert(op_id, dtypes[&x]);
                    *rcs.entry(x).or_insert(0) += 1;
                }
                Op::BroadcastTile { x, .. } => {
                    dtypes.insert(op_id, dtypes[&x]);
                    *rcs.entry(x).or_insert(0) += 1;
                }
                Op::Binary { x, y, bop } => {
                    let dtype = if bop.returns_bool() {
                        (DType::Bool, dtypes[&x].1)
                    } else {
                        dtypes[&x]
                    };
                    dtypes.insert(op_id, dtype);
                    *rcs.entry(x).or_insert(0) += 1;
                    *rcs.entry(y).or_insert(0) += 1;
                }
                Op::Asm { ref ops, .. } => {
                    // Result takes ops[0]'s dtype/layout; the template's
                    // `{i}` placeholders substitute the operand expressions.
                    let dtype = dtypes[&ops[0]];
                    dtypes.insert(op_id, dtype);
                    for &x in ops.iter() {
                        *rcs.entry(x).or_insert(0) += 1;
                    }
                }
                Op::Stack { ref ops } => {
                    let dtype = dtypes[&ops[0]];
                    dtypes.insert(op_id, (dtype.0, MemLayout::Vector(ops.len().try_into().unwrap())));
                    for &x in ops.iter() {
                        *rcs.entry(x).or_insert(0) += 1;
                    }
                }
                Op::Index { vec, idx: _ } => {
                    let dtype = dtypes[&vec];
                    dtypes.insert(op_id, (dtype.0, MemLayout::Scalar));
                    *rcs.entry(vec).or_insert(0) += 1;
                }
                Op::Wmma { dims: _, layout: _, dtype, a, b, c } => {
                    let out_dtype = match dtype {
                        MMADType::f16_f16_f16_f32 => DType::F32,
                        MMADType::f16_f16_f16_f16 => DType::F16,
                        MMADType::s8_s8_s32_s32
                        | MMADType::s4_s4_s32_s32
                        | MMADType::b1_b1_s32_xor_popc
                        | MMADType::b1_b1_s32_and_popc => DType::I32,
                    };
                    dtypes.insert(op_id, (out_dtype, MemLayout::Vector(4)));
                    *rcs.entry(a).or_insert(0) += 1;
                    *rcs.entry(b).or_insert(0) += 1;
                    *rcs.entry(c).or_insert(0) += 1;
                }
                Op::MatmulTile { x, y, acc } => {
                    dtypes.insert(op_id, dtypes[&acc]);
                    *rcs.entry(x).or_insert(0) += 1;
                    *rcs.entry(y).or_insert(0) += 1;
                    *rcs.entry(acc).or_insert(0) += 1;
                }
                Op::TransposeTile { x } => {
                    dtypes.insert(op_id, dtypes[&x]);
                    *rcs.entry(x).or_insert(0) += 1;
                }
                Op::Mad { x, y, z } => {
                    dtypes.insert(op_id, dtypes[&x]);
                    *rcs.entry(x).or_insert(0) += 1;
                    *rcs.entry(y).or_insert(0) += 1;
                    *rcs.entry(z).or_insert(0) += 1;
                }
                Op::Range { kind, .. } => {
                    // Group length is a consumed operand (e.g. a Binary
                    // expression) and needs a refcount for codegen register
                    // allocation. Same for the local id a Warp view reads.
                    if let RangeKind::Group(len) = kind {
                        *rcs.entry(len).or_insert(0) += 1;
                    }
                    if let RangeKind::Warp(local_id) = kind {
                        *rcs.entry(local_id).or_insert(0) += 1;
                    }
                    dtypes.insert(op_id, (IDX_T, MemLayout::Scalar));
                }
                Op::Loop { len, .. } => {
                    // Loop length is a consumed operand (e.g. a Binary
                    // expression) and needs a refcount for codegen register
                    // allocation.
                    *rcs.entry(len).or_insert(0) += 1;
                    dtypes.insert(op_id, (IDX_T, MemLayout::Scalar));
                }
                Op::If { condition } => {
                    *rcs.entry(condition).or_insert(0) += 1;
                }
                Op::Barrier | Op::EndIf | Op::EndLoop => {}
            }
            op_id = self.next_op(op_id);
        }
        if !op_id.is_null() {
            panic!("compute_dtypes_and_rcs did not finish in 10000 steps");
        }
        (dtypes, rcs)
    }

    /// Resolve the layout of an operation's result by walking the IR.
    pub(crate) fn layout(&self, mut op_id: OpId) -> MemLayout {
        for _ in 0..10000 {
            match self.ops[op_id].op {
                Op::Const(_) | Op::Param { .. } | Op::Storage { .. } => return MemLayout::Scalar,
                Op::Range { .. } => return MemLayout::Scalar,
                Op::Cast { x, .. } | Op::Bitcast { x, .. } => op_id = x,
                Op::Load { layout, .. } => return layout,
                Op::Store { src: x, .. } => op_id = x,
                Op::Unary { x, .. } => op_id = x,
                Op::Binary { x, .. } => op_id = x,
                Op::Mad { x, .. } => op_id = x,
                Op::Wmma { dims, .. } => match dims {
                    MMADims::m8n8k16 => return MemLayout::Vector(2),
                    MMADims::m16n8k8 => return MemLayout::Vector(4),
                    MMADims::m16n8k16 => return MemLayout::Vector(4),
                    MMADims::m32n8k16 => return MemLayout::Vector(8),
                    MMADims::m8n32k16 => return MemLayout::Vector(8),
                    MMADims::m8n8k32 => return MemLayout::Vector(2),
                    MMADims::m8n8k128 => return MemLayout::Vector(2),
                },
                Op::MatmulTile { acc, .. } => op_id = acc,
                Op::TransposeTile { x } => op_id = x,
                Op::Stack { ref ops } => {
                    return MemLayout::Vector(ops.len().try_into().unwrap());
                }
                Op::Asm { ref ops, .. } => op_id = ops[0],
                // Index extracts a single lane: a scalar, not the vec layout.
                Op::Index { .. } => return MemLayout::Scalar,
                Op::Move { x, .. } => op_id = x,
                Op::Reduce { x, .. } => op_id = x,
                Op::ReduceTile { acc, .. } => op_id = acc,
                Op::BroadcastTile { x, .. } => op_id = x,
                Op::EndLoop | Op::Loop { .. } => return MemLayout::Scalar,
                Op::Barrier | Op::If { .. } | Op::EndIf => todo!(),
            }
        }
        panic!("layout not found for too long time");
    }

    /// Resolve the dtype of an operation's result by walking the IR.
    pub(crate) fn dtype(&self, mut op_id: OpId) -> DType {
        //println!("getting dtype of id: {op_id:?}'");
        for _ in 0..10000 {
            match self.ops[op_id].op {
                Op::Const(c) => return c.dtype(),
                Op::Param { dtype, .. } => return dtype,
                Op::Storage { dtype, .. } => return dtype,
                Op::Cast { dtype, .. } => return dtype,
                Op::Bitcast { dtype, .. } => return dtype,
                Op::Range { .. } => return IDX_T,
                Op::Load { src, .. } => op_id = src,
                Op::Unary { x, .. } => op_id = x,
                Op::Binary { x, bop, .. } => {
                    if bop.returns_bool() {
                        return DType::Bool;
                    }
                    op_id = x;
                }
                Op::Mad { x, .. } => op_id = x,
                Op::Wmma { dtype, .. } => match dtype {
                    MMADType::f16_f16_f16_f32 => return DType::F32,
                    MMADType::f16_f16_f16_f16 => return DType::F16,
                    MMADType::s8_s8_s32_s32
                    | MMADType::s4_s4_s32_s32
                    | MMADType::b1_b1_s32_xor_popc
                    | MMADType::b1_b1_s32_and_popc => return DType::I32,
                },
                Op::MatmulTile { acc, .. } => op_id = acc,
                Op::TransposeTile { x } => op_id = x,
                Op::Stack { ref ops } => op_id = ops[0],
                Op::Asm { ref ops, .. } => op_id = ops[0],
                Op::Index { vec, .. } => op_id = vec,
                Op::Store { src: x, .. } => op_id = x,
                Op::Move { x, .. } => op_id = x,
                Op::Reduce { x, .. } => op_id = x,
                Op::ReduceTile { acc, .. } => op_id = acc,
                Op::BroadcastTile { x, .. } => op_id = x,
                Op::EndLoop | Op::Loop { .. } => return IDX_T,
                Op::Barrier | Op::If { .. } | Op::EndIf => todo!(),
            }
        }
        panic!("dtype not found for too long time");
    }

    #[track_caller]
    pub(crate) fn at(&self, op_id: OpId) -> &Op {
        &self.ops[op_id].op
    }

    pub(crate) fn prev_op(&self, op_id: OpId) -> OpId {
        self.ops[op_id].prev
    }

    pub(crate) fn next_op(&self, op_id: OpId) -> OpId {
        self.ops[op_id].next
    }

    /*pub fn ops_mut(&mut self) -> impl Iterator<Item = &mut Op> {
        self.ops.values_mut().map(|op_node| &mut op_node.op)
    }*/

    pub(crate) fn insert_before(&mut self, before_id: OpId, op: Op) -> OpId {
        debug_assert!(!before_id.is_null());
        debug_assert!(!self.ops.is_empty());

        let prev = self.ops[before_id].prev;
        let op_node = OpLinked { prev, next: before_id, op };
        let op_id = self.ops.push(op_node);
        self.ops[before_id].prev = op_id;
        if prev.is_null() {
            self.head = op_id;
        } else {
            self.ops[prev].next = op_id;
        }
        op_id
    }

    pub(crate) fn insert_const_idx_before(&mut self, before_id: OpId, val: impl crate::scalar::Scalar) -> OpId {
        self.insert_before(before_id, Op::Const(Constant::idx(val)))
    }

    pub(crate) fn insert_after(&mut self, after_id: OpId, op: Op) -> OpId {
        debug_assert!(!after_id.is_null());
        debug_assert!(!self.ops.is_empty());

        let next = self.ops[after_id].next;
        let op_node = OpLinked { prev: after_id, next, op };
        let op_id = self.ops.push(op_node);
        self.ops[after_id].next = op_id;
        if next.is_null() {
            self.tail = op_id;
        } else {
            self.ops[next].prev = op_id;
        }
        op_id
    }

    pub(crate) fn move_op_after(&mut self, op_id: OpId, after_id: OpId) {
        debug_assert!(!op_id.is_null());
        debug_assert!(!after_id.is_null());
        debug_assert!(!self.ops.is_empty());

        if op_id == after_id {
            return;
        }

        //println!("moving op={op_id}, after={after_id}");

        // Remove
        let OpLinked { prev, next, .. } = self.ops[op_id];
        if prev.is_null() {
            self.head = next;
        } else {
            self.ops[prev].next = next;
        }
        if next.is_null() {
            self.tail = prev;
        } else {
            self.ops[next].prev = prev;
        }

        // Insert
        self.ops[op_id].prev = after_id;
        let next = self.ops[after_id].next;
        self.ops[op_id].next = next;
        self.ops[after_id].next = op_id;
        if next.is_null() {
            self.tail = op_id;
        } else {
            self.ops[next].prev = op_id;
        }
    }

    /// Move an operation before another operation.
    ///
    /// Moves `op_id` to appear immediately before `before_id` in the operation chain.
    pub(crate) fn move_op_before(&mut self, op_id: OpId, before_id: OpId) {
        debug_assert!(!op_id.is_null());
        debug_assert!(!before_id.is_null());
        debug_assert!(!self.ops.is_empty());

        if op_id == before_id {
            return;
        }

        //println!("moving op={op_id}, before={before_id}");

        // Remove
        let OpLinked { prev, next, .. } = self.ops[op_id];
        if prev.is_null() {
            self.head = next;
        } else {
            self.ops[prev].next = next;
        }
        if next.is_null() {
            self.tail = prev;
        } else {
            self.ops[next].prev = prev;
        }

        // Insert
        self.ops[op_id].next = before_id;
        let prev = self.ops[before_id].prev;
        self.ops[op_id].prev = prev;
        self.ops[before_id].prev = op_id;
        if prev.is_null() {
            self.head = op_id;
        } else {
            self.ops[prev].next = op_id;
        }
    }

    /// Remove an operation from the kernel.
    ///
    /// Removes the operation with `op_id` from the kernel IR.
    pub(crate) fn remove_op(&mut self, op_id: OpId) {
        debug_assert!(!op_id.is_null());
        debug_assert!(!self.ops.is_empty());

        let OpLinked { prev, next, .. } = self.ops[op_id];
        if prev.is_null() {
            self.head = next;
        } else {
            self.ops[prev].next = next;
        }
        if next.is_null() {
            self.tail = prev;
        } else {
            self.ops[next].prev = prev;
        }

        self.ops.remove(op_id);
    }

    /// Remove the transitive dependency chain of `x` that is not needed
    /// by any store or any op in `keep_alive`.
    /// Removes ops from `x` backwards that are no longer needed.
    /// Returns a filtered version of `loads` with entries for removed LoadViews removed.
    /// The i-th entry in `loads` corresponds to the i-th LoadView in op order.
    /// Generic over the load payload (`TensorId` in runtime kernels,
    /// `ClassId` in kernelizer jit kernels) — it is never inspected.
    pub(crate) fn remove_unused_chain<T: Copy>(&mut self, x: OpId, keep_alive: &[OpId], loads: &[T]) -> Vec<T> {
        let mut chain: Set<OpId> = Set::default();
        let mut stack = vec![x];
        for _ in 0..30_000 {
            let Some(op) = stack.pop() else { break };
            if chain.insert(op) {
                stack.extend(self.ops[op].op.parameters().filter(|&p| !p.is_null()));
            }
        }
        if !stack.is_empty() {
            panic!("remove_unused_chain did not finish in 10000 steps");
        }

        let mut live: Set<OpId> = Set::default();
        stack.extend_from_slice(keep_alive);
        let mut op_id = self.head;
        for _ in 0..30_000 {
            if op_id.is_null() {
                break;
            }
            if matches!(self.ops[op_id].op, Op::Store { .. }) {
                stack.push(op_id);
            }
            op_id = self.next_op(op_id);
        }
        if !op_id.is_null() {
            panic!("remove_unused_chain did not finish in 10000 steps");
        }
        for _ in 0..30_000 {
            let Some(op) = stack.pop() else { break };
            if live.insert(op) {
                stack.extend(self.ops[op].op.parameters().filter(|&p| !p.is_null()));
            }
        }
        if !stack.is_empty() {
            panic!("remove_unused_chain did not finish in 10000 steps");
        }

        // Collect input (ro) Param OpIds in op order (before removal). These are
        // the Params that correspond 1:1 with `loads`: Global and Variable Params
        // are loads; GlobalMut Params are stores (mapped to `stores`, not `loads`).
        let param_ops: Vec<OpId> = {
            let mut ops = Vec::new();
            let mut id = self.head;
            for _ in 0..30_000 {
                if id.is_null() {
                    break;
                }
                if let Op::Param { kind, .. } = &self.ops[id].op {
                    if *kind != ParamKind::GlobalMut {
                        ops.push(id);
                    }
                }
                id = self.next_op(id);
            }
            if !id.is_null() {
                panic!("remove_unused_chain did not finish in 10000 steps");
            }
            ops
        };

        let to_remove: Set<OpId> = chain.difference(&live).copied().collect();
        let mut op_id = self.head;
        for _ in 0..30_000 {
            if op_id.is_null() {
                break;
            }
            let next = self.next_op(op_id);
            if to_remove.contains(&op_id) {
                self.remove_op(op_id);
            }
            op_id = next;
        }
        if !op_id.is_null() {
            panic!("remove_unused_chain did not finish in 10000 steps");
        }

        // Keep only loads whose corresponding input Param was not removed.
        // `loads[i]` maps to the i-th Global/Variable Param in op order.
        param_ops.iter().enumerate().filter(|&(_, &lv_id)| !to_remove.contains(&lv_id)).map(|(i, _)| loads[i]).collect()
    }

    /// Iterate over all operations in the kernel.
    ///
    /// Returns an iterator over all operations without any ordering guarantees.
    pub(crate) fn iter_unordered(&self) -> impl Iterator<Item = (OpId, &Op)> {
        self.ops.iter().map(|(id, node)| (id, &node.op))
    }

    pub(crate) fn name(&self) -> String {
        let mut parts: Vec<&str> = Vec::new();
        let mut op_id = self.head;
        for _ in 0..10_000 {
            if op_id.is_null() {
                break;
            }
            match self.at(op_id) {
                Op::Unary { uop, .. } => parts.push(match uop {
                    UOp::Neg => "neg",
                    UOp::Not => "not",
                    UOp::BitNot => "bitnot",
                    UOp::Exp => "exp",
                    UOp::Exp2 => "exp2",
                    UOp::Log2 => "log2",
                    UOp::Reciprocal => "reciprocal",
                    UOp::Sqrt => "sqrt",
                    UOp::Rsqrt => "rsqrt",
                    UOp::Sin => "sin",
                    UOp::Cos => "cos",
                    UOp::Floor => "floor",
                    UOp::Trunc => "trunc",
                    UOp::Abs => "abs",
                }),
                Op::Binary { bop, .. } => parts.push(match bop {
                    BOp::Add => "add",
                    BOp::Sub => "sub",
                    BOp::Mul => "mul",
                    BOp::Div => "div",
                    BOp::Pow => "pow",
                    BOp::Mod => "mod",
                    BOp::Cmplt => "cmplt",
                    BOp::Cmpgt => "cmpgt",
                    BOp::Cmpge => "cmpge",
                    BOp::Max => "max",
                    BOp::Or => "or",
                    BOp::And => "and",
                    BOp::BitXor => "bitxor",
                    BOp::BitOr => "bitor",
                    BOp::BitAnd => "bitand",
                    BOp::BitShiftLeft => "shl",
                    BOp::BitShiftRight => "shr",
                    BOp::NotEq => "neq",
                    BOp::Eq => "eq",
                }),
                Op::Reduce { rop, .. } => parts.push(match rop {
                    BOp::Add => "sum",
                    BOp::Max => "max",
                    BOp::Mul => "prod",
                    _ => "reduce",
                }),
                Op::ReduceTile { rop, .. } => parts.push(match rop {
                    BOp::Add => "reduce_tile_sum",
                    BOp::Max => "reduce_tile_max",
                    BOp::Mul => "reduce_tile_prod",
                    _ => "reduce_tile",
                }),
                Op::Mad { .. } => parts.push("mad"),
                Op::Wmma { .. } => parts.push("wmma"),
                Op::Cast { .. } => parts.push("cast"),
                Op::Bitcast { .. } => parts.push("bitcast"),
                _ => {}
            }
            op_id = self.next_op(op_id);
        }
        if !op_id.is_null() {
            panic!("name did not finish in 10000 steps");
        }
        parts.dedup();
        if parts.is_empty() {
            return "copy".into();
        }
        parts.join("_")
    }

    /// Check if the kernel contains any store operations.
    pub(crate) fn contains_stores(&self) -> bool {
        self.ops.values().any(|x| matches!(x.op, Op::Store { .. }))
    }

    /// Estimated flops, global bytes read, global bytes written for a launch
    /// of this kernel. Iteration counts from enclosing `Loop`/`Range`
    /// (Groups/Locals) multiply per-op costs; variable trip counts that
    /// cannot be resolved at compile time are treated as 1.
    pub fn flop_mem_rw(&self) -> (u64, u64, u64) {
        let mut flops: u64 = 0;
        let mut read: u64 = 0;
        let mut write: u64 = 0;
        let mut loop_stack: Vec<u64> = Vec::new();
        let mut group_stack: Vec<u64> = Vec::new();
        let mut local_prod: u64 = 1;

        let mut op_id = self.head;
        for _ in 0..20_000 {
            if op_id.is_null() {
                break;
            }
            // Multiplier for this op: all enclosing loops * all parallel threads.
            let loop_prod: u64 = loop_stack.iter().product::<u64>().max(1);
            let group_prod: u64 = group_stack.iter().product::<u64>().max(1);
            let mult: u64 = loop_prod.saturating_mul(group_prod).saturating_mul(local_prod);

            match self.at(op_id) {
                Op::Loop { len } => {
                    let v = self.resolve_const(*len).and_then(|c| c.as_dim()).unwrap_or(1).max(1) as u64;
                    // The loop body executes v times; push for following ops.
                    // Count is applied via mult on body ops, so push now.
                    loop_stack.push(v);
                }
                Op::EndLoop => {
                    loop_stack.pop();
                }
                Op::Range { kind, .. } => match kind {
                    RangeKind::Group(l) => {
                        let v = self.resolve_const(*l).and_then(|c| c.as_dim()).unwrap_or(1).max(1) as u64;
                        group_stack.push(v);
                    }
                    RangeKind::Local(n) => {
                        local_prod = local_prod.saturating_mul(*n as u64);
                    }
                    RangeKind::Warp(_) => {}
                },
                Op::Unary { .. } => {
                    if self.dtype(op_id).is_float() {
                        flops = flops.saturating_add(mult)
                    }
                }
                Op::Binary { .. } => {
                    if self.dtype(op_id).is_float() {
                        flops = flops.saturating_add(mult)
                    }
                }
                Op::Mad { .. } => {
                    if self.dtype(op_id).is_float() {
                        flops = flops.saturating_add(2 * mult)
                    }
                }
                Op::Wmma { dims, .. } => {
                    let (m, n, k) = dims.decompose_mnk();
                    // One wmma per warp, not per thread.
                    let ws = u64::from(self.dev_info().warp_size);
                    let warps = (group_prod * local_prod / ws.max(1)).max(1);
                    let loop_prod: u64 = loop_stack.iter().product::<u64>().max(1);
                    flops = flops.saturating_add(2 * m * n * k * warps * loop_prod);
                }
                Op::ReduceTile { .. } => flops = flops.saturating_add(mult),
                Op::Load { src, layout, .. } => {
                    if let Op::Param { kind: ParamKind::Global, dtype, .. } = &self.ops[*src].op {
                        let bytes = (dtype.bit_size() as u64 / 8) * layout.n_elements() as u64;
                        read = read.saturating_add(bytes.saturating_mul(mult));
                    }
                }
                Op::Store { dst, layout, .. } => {
                    if let Op::Param { kind: ParamKind::GlobalMut, dtype, .. } = &self.ops[*dst].op {
                        let bytes = (dtype.bit_size() as u64 / 8) * layout.n_elements() as u64;
                        write = write.saturating_add(bytes.saturating_mul(mult));
                    }
                }
                _ => {}
            }
            op_id = self.next_op(op_id);
        }
        (flops, read, write)
    }

    /// Check if the kernel is a reduction kernel.
    pub(crate) fn is_reduce(&self) -> bool {
        self.ops.values().any(|x| matches!(x.op, Op::Reduce { .. } | Op::ReduceTile { .. }))
    }

    /// Tensor shape of the value produced by `op_id`, as per-dimension op
    /// ids. Symbolic walk over existing ops with a per-op memo (shared
    /// sub-DAGs are re-entered by many parents; re-walking them is
    /// exponential without the memo). A `Stack` accessed *directly* (as a
    /// value) produces `[len] + first_element_shape`: its length is emitted
    /// as a const dim. A `Stack` referenced as a *shape descriptor* (through
    /// `Param { shape }`, `Reshape { shape }`, `Expand { shape }`) yields
    /// its element ops directly. Movement ops apply their transformation
    /// (`Permute` reorders, `Pad`/`Narrow` substitute their stored `len`);
    /// a scalar value (`Const`, scalar `Param`) has an empty shape.
    #[must_use]
    pub(crate) fn shape_ids(&mut self, op_id: OpId) -> Vec<OpId> {
        if op_id.is_null() {
            return Vec::new();
        }
        if let Some(cached) = self.shape_cache.get(&op_id) {
            return cached.clone();
        }
        // Unpacks a shape descriptor (a `Stack` of dims or a bare `Const`
        // dim, as stored in `Param { shape }`, `Reshape { shape }`,
        // `Expand { shape }`) into its dimension op ids.
        //
        // Each returned entry is the op computing that dim's length. A
        // dynamic (runtime-known only) dim length is [`OpId::NULL`]: rank is
        // always statically known, individual lengths may not be. An empty
        // vec means a scalar (rank 0).
        //
        // A Param as a shape descriptor is a shape *tensor*: scalar
        // (`shape == NULL`) → one dynamic dim; 1d vector → one dynamic dim
        // per element, loaded element-wise from the param. Higher rank is
        // invalid. Shape tensors must be [`IDX_T`].
        fn descriptor(k: &mut Kernel, id: OpId) -> Vec<OpId> {
            if id.is_null() {
                return Vec::new();
            }
            let (shape, dtype) = match k.ops[id].op {
                Op::Stack { ref ops } => return ops.to_vec(),
                // A bare Const is a single concrete dim length.
                Op::Const(_) => return vec![id],
                // A scalar dim *expression* (over consts, params or loads)
                // computes its own single dim length.
                Op::Unary { .. } | Op::Binary { .. } | Op::Load { .. } => return vec![id],
                Op::Param { shape, dtype, .. } => (shape, dtype),
                ref op => todo!("shape_ids: invalid shape descriptor {op:?}"),
            };
            debug_assert!(dtype == IDX_T, "shape tensor must be {IDX_T:?}, got {dtype:?}");
            // A scalar shape tensor (shape == NULL) IS its single dim: the
            // param op itself is the length (a variable is its value — no
            // load). Only a 1-d shape tensor needs element-wise loads.
            if shape.is_null() {
                return vec![id];
            }
            let rank = k.shape(shape).len();
            debug_assert!(rank <= 1, "shape_ids: param shape descriptor must be 0d or 1d, got rank {rank}");
            let mut dims = Vec::with_capacity(rank);
            for i in 0..rank {
                dims.push(k.push_back(Op::Index { vec: id, idx: i }));
            }
            dims
        }
        let root = op_id;
        let mut stack = vec![op_id];
        let mut visited: Map<OpId, Vec<OpId>> = Map::default();
        for _ in 0..10_000 {
            let Some(op_id) = stack.pop() else {
                let shape = visited.remove(&op_id).unwrap();
                self.shape_cache.insert(root, shape.clone());
                return shape;
            };
            if visited.contains_key(&op_id) {
                continue;
            }
            let node_op = self.ops[op_id].op.clone();
            match node_op {
                Op::Const(_) => {
                    visited.insert(op_id, vec![]);
                }
                Op::Param { shape, .. } => {
                    visited.insert(op_id, descriptor(self, shape));
                }
                // Direct access of a stack op: it is a rank-1+ tensor whose
                // leading dim is the element count (emitted as a const).
                Op::Stack { ref ops } => {
                    if ops.is_empty() {
                        visited.insert(op_id, vec![]);
                    } else if let Some(mut dims) = visited.get(&ops[0]).cloned() {
                        dims.insert(0, self.const_idx(ops.len() as u32));
                        visited.insert(op_id, dims);
                    } else {
                        stack.push(op_id);
                        stack.push(ops[0]);
                    }
                }
                Op::Move { x, ref mop } => match mop.as_ref() {
                    MoveOp::Reshape { shape, .. } | MoveOp::Expand { shape } => {
                        visited.insert(op_id, descriptor(self, *shape));
                    }
                    MoveOp::Permute { axes } => match visited.get(&x) {
                        Some(dims) => {
                            visited.insert(op_id, crate::shape::permute(dims, axes));
                        }
                        None => {
                            stack.push(op_id);
                            stack.push(x);
                        }
                    },
                    MoveOp::Flip { .. } => match visited.get(&x) {
                        Some(dims) => {
                            visited.insert(op_id, dims.clone());
                        }
                        None => {
                            stack.push(op_id);
                            stack.push(x);
                        }
                    },
                    &MoveOp::Narrow { axis, len, .. } | &MoveOp::Pad { axis, len, .. } => match visited.get(&x).cloned() {
                        Some(mut dims) => {
                            if dims.is_empty() {
                                dims = vec![len];
                            } else {
                                dims[axis as usize] = len;
                            }
                            visited.insert(op_id, dims);
                        }
                        None => {
                            stack.push(op_id);
                            stack.push(x);
                        }
                    },
                },
                // Scalars broadcast implicitly: if `x` is a scalar (empty
                // shape), the binary takes `y`'s shape.
                Op::Binary { x, y, .. } => {
                    let x_done = visited.contains_key(&x);
                    if x_done && !visited[&x].is_empty() {
                        visited.insert(op_id, visited[&x].clone());
                    } else if visited.contains_key(&y) {
                        visited.insert(op_id, visited[&y].clone());
                    } else {
                        stack.push(op_id);
                        if !x_done {
                            stack.push(x);
                        }
                        stack.push(y);
                    }
                }
                Op::Reduce { x, .. } => match visited.get(&x).cloned() {
                    Some(mut dims) => {
                        dims.truncate(dims.len().saturating_sub(1));
                        visited.insert(op_id, dims);
                    }
                    None => {
                        stack.push(op_id);
                        stack.push(x);
                    }
                },
                Op::Load { src: x, .. }
                | Op::Store { src: x, .. }
                | Op::Cast { x, .. }
                | Op::Bitcast { x, .. }
                | Op::Unary { x, .. }
                | Op::Mad { x, .. }
                | Op::MatmulTile { x, .. }
                | Op::ReduceTile { x, .. }
                | Op::Index { vec: x, .. }
                | Op::TransposeTile { x } => match visited.get(&x) {
                    Some(dims) => {
                        visited.insert(op_id, dims.clone());
                    }
                    None => {
                        stack.push(op_id);
                        stack.push(x);
                    }
                },
                Op::Asm { ref ops, .. } => match visited.get(&ops[0]) {
                    Some(dims) => {
                        visited.insert(op_id, dims.clone());
                    }
                    None => {
                        stack.push(op_id);
                        stack.push(ops[0]);
                    }
                },
                // Group/loop indices and scalar storages are scalar values.
                Op::Range { .. } | Op::Loop { .. } => {
                    visited.insert(op_id, vec![]);
                }
                Op::Storage { len, .. } if len == 1 => {
                    visited.insert(op_id, vec![]);
                }
                ref op => todo!("shape_ids of {op:?}"),
            }
        }
        panic!("shape_ids did not resolve in 10000 steps");
    }

    /// Numeric tensor shape of the value produced by `op_id`: each dim
    /// resolved to its const value, `-1` where resolution fails (symbolic).
    /// Immutable — unlike [`Self::shape_ids`] this never emits ops; a
    /// directly accessed `Stack` contributes its element count as a plain
    /// number instead of a const dim op.
    pub(crate) fn shape(&self, op_id: OpId) -> Vec<Dim> {
        // Unpacks a shape descriptor into its dimension op ids. A descriptor
        // is a `Stack` of dims, or a single dim op — a `Const`, a scalar
        // `Param`, or a dim *expression* (Const/Unary/Binary/Cast/Load tree,
        // folded per-dim by `resolve_const`, `-1` where it contains a
        // variable). Anything else resolves to `-1` via `resolve_const` and
        // fails loudly downstream where the dim is consumed.
        fn descriptor(k: &Kernel, id: OpId) -> Vec<OpId> {
            if id.is_null() {
                return Vec::new();
            }
            match k.ops[id].op {
                Op::Stack { ref ops } => ops.to_vec(),
                _ => vec![id],
            }
        }
        if op_id.is_null() {
            return Vec::new();
        }
        let mut stack = vec![op_id];
        let mut visited: Map<OpId, Vec<Dim>> = Map::default();
        for _ in 0..10_000 {
            let Some(id) = stack.pop() else {
                return visited[&op_id].clone();
            };
            if visited.contains_key(&id) {
                continue;
            }
            match self.ops[id].op {
                Op::Const(_) => {
                    visited.insert(id, vec![]);
                }
                Op::Param { shape, .. } => {
                    visited.insert(
                        id,
                        descriptor(self, shape)
                            .iter()
                            .map(|&d| self.resolve_const(d).and_then(crate::dtype::Constant::as_dim).unwrap_or(-1))
                            .collect(),
                    );
                }
                // Direct access of a stack op: `[len] + first_element_shape`.
                Op::Stack { ref ops } => {
                    // Shape of a stacked value: `[len] + first element's shape`.
                    if ops.is_empty() {
                        visited.insert(id, vec![]);
                    } else if let Some(mut dims) = visited.get(&ops[0]).cloned() {
                        dims.insert(0, ops.len() as Dim);
                        visited.insert(id, dims);
                    } else {
                        stack.push(id);
                        stack.push(ops[0]);
                    }
                }
                Op::Move { x, ref mop } => match mop.as_ref() {
                    MoveOp::Reshape { shape, .. } | MoveOp::Expand { shape } => {
                        visited.insert(
                            id,
                            descriptor(self, *shape)
                                .iter()
                                .map(|&d| self.resolve_const(d).and_then(crate::dtype::Constant::as_dim).unwrap_or(-1))
                                .collect(),
                        );
                    }
                    MoveOp::Permute { axes } => match visited.get(&x) {
                        Some(dims) => {
                            visited.insert(id, crate::shape::permute(dims, axes));
                        }
                        None => {
                            stack.push(id);
                            stack.push(x);
                        }
                    },
                    MoveOp::Flip { .. } => match visited.get(&x) {
                        Some(dims) => {
                            visited.insert(id, dims.clone());
                        }
                        None => {
                            stack.push(id);
                            stack.push(x);
                        }
                    },
                    &MoveOp::Narrow { axis, len, .. } | &MoveOp::Pad { axis, len, .. } => match visited.get(&x).cloned() {
                        Some(mut dims) => {
                            if dims.is_empty() {
                                dims = vec![self.resolve_const(len).and_then(crate::dtype::Constant::as_dim).unwrap_or(-1)];
                            } else {
                                dims[axis as usize] =
                                    self.resolve_const(len).and_then(crate::dtype::Constant::as_dim).unwrap_or(-1);
                            }
                            visited.insert(id, dims);
                        }
                        None => {
                            stack.push(id);
                            stack.push(x);
                        }
                    },
                },
                // Scalars broadcast implicitly: if `x` is a scalar (empty
                // shape), the binary takes `y`'s shape.
                Op::Binary { x, y, .. } => {
                    let x_done = visited.contains_key(&x);
                    if x_done && !visited[&x].is_empty() {
                        visited.insert(id, visited[&x].clone());
                    } else if visited.contains_key(&y) {
                        visited.insert(id, visited[&y].clone());
                    } else {
                        stack.push(id);
                        if !x_done {
                            stack.push(x);
                        }
                        stack.push(y);
                    }
                }
                Op::Reduce { x, .. } => match visited.get(&x).cloned() {
                    Some(mut dims) => {
                        dims.truncate(dims.len().saturating_sub(1));
                        visited.insert(id, dims);
                    }
                    None => {
                        stack.push(id);
                        stack.push(x);
                    }
                },
                Op::Load { src: x, .. }
                | Op::Store { src: x, .. }
                | Op::Cast { x, .. }
                | Op::Bitcast { x, .. }
                | Op::Unary { x, .. }
                | Op::Mad { x, .. }
                | Op::MatmulTile { x, .. }
                | Op::ReduceTile { x, .. }
                | Op::Index { vec: x, .. }
                | Op::TransposeTile { x } => match visited.get(&x) {
                    Some(dims) => {
                        visited.insert(id, dims.clone());
                    }
                    None => {
                        stack.push(id);
                        stack.push(x);
                    }
                },
                Op::Asm { ref ops, .. } => match visited.get(&ops[0]) {
                    Some(dims) => {
                        visited.insert(id, dims.clone());
                    }
                    None => {
                        stack.push(id);
                        stack.push(ops[0]);
                    }
                },
                // Group/loop indices and scalar storages are scalar values.
                Op::Range { .. } | Op::Loop { .. } => {
                    visited.insert(id, vec![]);
                }
                Op::Storage { len, .. } if len == 1 => {
                    visited.insert(id, vec![]);
                }
                ref op => todo!("shape of {op:?}"),
            }
        }
        panic!("shape did not resolve in 10000 steps");
    }

    /// Builds a single shape-descriptor op for the value produced by
    /// `op_id`, used to size a store/param buffer: `NULL` for scalars, the
    /// bare dim for rank-1, a `Stack` otherwise.
    #[must_use]
    pub(crate) fn stack_shape_dims(&mut self, op_id: OpId) -> OpId {
        match self.shape_ids(op_id).as_slice() {
            [] => OpId::NULL,
            [dim] => *dim,
            dims => self.stack(dims),
        }
    }

    /// Get index loop ids, dimensions and strides.
    ///
    /// Returns `loop_id` -> (dimension, stride) where NULL means unknown stride.
    pub(crate) fn get_strides(&self, index: OpId) -> Map<OpId, (Dim, Dim)> {
        //println!("Get index {index}");

        let index_len_of = |op: &Op| -> Dim {
            match op {
                Op::Range { kind, .. } => match kind {
                    // Dynamic dims are `-1`; autotune substitutes 42 (see `alloc_buffers`).
                    RangeKind::Group(len) => self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap_or(42),
                    RangeKind::Local(len) => i64::from(*len),
                    RangeKind::Warp(_) => i64::from(self.dev_info().warp_size),
                },
                _ => unreachable!(),
            }
        };

        let mut params = vec![(index, 1i64)];
        let mut indices = Map::default();

        for _ in 0..10_000 {
            let Some((param, scale)) = params.pop() else { break };
            match self.ops[param].op {
                Op::Binary { x, y, bop } => {
                    if bop == BOp::Add {
                        if let Op::Loop { len, .. } = self.ops[x].op {
                            indices.insert(x, (self.resolve_const(len).and_then(crate::dtype::Constant::as_dim).unwrap(), 1));
                            params.push((y, scale));
                        } else if let Op::Range { .. } = self.ops[x].op {
                            indices.insert(x, (index_len_of(&self.ops[x].op), 1));
                            params.push((y, scale));
                        } else if let Op::Loop { len, .. } = self.ops[y].op {
                            indices.insert(y, (self.resolve_const(len).and_then(crate::dtype::Constant::as_dim).unwrap(), 1));
                            params.push((x, scale));
                        } else if let Op::Range { .. } = self.ops[y].op {
                            indices.insert(y, (index_len_of(&self.ops[y].op), 1));
                            params.push((x, scale));
                        } else {
                            params.push((x, scale));
                            params.push((y, scale));
                        }
                    }
                    if bop == BOp::Mul {
                        match (&self.ops[x].op, &self.ops[y].op) {
                            (Op::Loop { len, .. }, Op::Const(c)) => {
                                indices.insert(
                                    x,
                                    (
                                        self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                        c.as_dim().unwrap() * scale,
                                    ),
                                );
                            }
                            (Op::Const(c), Op::Loop { len, .. }) => {
                                indices.insert(
                                    y,
                                    (
                                        self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                        c.as_dim().unwrap() * scale,
                                    ),
                                );
                            }
                            (Op::Range { .. }, Op::Const(c)) => {
                                indices.insert(x, (index_len_of(&self.ops[x].op), c.as_dim().unwrap() * scale));
                            }
                            (Op::Const(c), Op::Range { .. }) => {
                                indices.insert(y, (index_len_of(&self.ops[y].op), c.as_dim().unwrap() * scale));
                            }
                            _ => {}
                        }
                    }
                    if bop == BOp::BitShiftLeft {
                        match (&self.ops[x].op, &self.ops[y].op) {
                            (Op::Loop { len, .. }, Op::Const(c)) => {
                                indices.insert(
                                    x,
                                    (
                                        self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                        (1i64 << c.as_dim().unwrap()) * scale,
                                    ),
                                );
                            }
                            (Op::Range { .. }, Op::Const(c)) => {
                                indices.insert(x, (index_len_of(&self.ops[x].op), (1i64 << c.as_dim().unwrap()) * scale));
                            }
                            (Op::Const(c), Op::Range { .. }) => {
                                indices.insert(y, (index_len_of(&self.ops[y].op), (1i64 << c.as_dim().unwrap()) * scale));
                            }
                            (Op::Const(c), Op::Loop { len, .. }) => {
                                indices.insert(
                                    y,
                                    (
                                        self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                        (1i64 << c.as_dim().unwrap()) * scale,
                                    ),
                                );
                            }
                            _ => {
                                if let Op::Const(c) = self.ops[y].op {
                                    params.push((x, scale * (1i64 << c.as_dim().unwrap())));
                                }
                            }
                        }
                    }
                }
                Op::Mad { x, y, z } => {
                    match &self.ops[z].op {
                        Op::Loop { len, .. } => {
                            indices.insert(z, (self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(), 1));
                        }
                        Op::Range { .. } => {
                            indices.insert(z, (index_len_of(&self.ops[z].op), 1));
                        }
                        _ => {
                            params.push((z, scale));
                        }
                    }
                    match (&self.ops[x].op, &self.ops[y].op) {
                        (Op::Loop { len, .. }, Op::Const(c)) => {
                            indices.insert(
                                x,
                                (
                                    self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                    c.as_dim().unwrap() * scale,
                                ),
                            );
                        }
                        (Op::Range { .. }, Op::Const(c)) => {
                            indices.insert(x, (index_len_of(&self.ops[x].op), c.as_dim().unwrap() * scale));
                        }
                        (Op::Const(c), Op::Loop { len, .. }) => {
                            indices.insert(
                                y,
                                (
                                    self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim).unwrap(),
                                    c.as_dim().unwrap() * scale,
                                ),
                            );
                        }
                        (Op::Const(c), Op::Range { .. }) => {
                            indices.insert(y, (index_len_of(&self.ops[y].op), c.as_dim().unwrap() * scale));
                        }
                        _ => {}
                    }
                }
                Op::Const(c) => {
                    indices
                        .entry(OpId::NULL)
                        .and_modify(|(_, v)| *v += c.as_dim().unwrap() * scale)
                        .or_insert((0, c.as_dim().unwrap() * scale));
                }
                _ => {}
            }
        }
        if !params.is_empty() {
            panic!("get_strides did not finish in 10000 steps");
        }

        indices
    }

    /// Fully resolve a kernel-shape value (an `OpId`) down to a constant dim.
    ///
    /// Follows `Cast`, `Unary`, `Binary`, `Loop`, and `Index` operands,
    /// evaluating constant expressions with `Constant::unary`/`Constant::binary`.
    /// Returns `Some(dim)` when the value depends only on constants, and `None`
    /// when it is genuinely dynamic (e.g. involves a `Param` variable). Callers
    /// must bail (not optimize) on `None`.
    pub(crate) fn resolve_const(&self, op_id: OpId) -> Option<Constant> {
        // Collect the constant-expression subgraph reachable through `Cast`,
        // `Unary`, `Binary`, `Loop`, and `Index` operands in postorder
        // (children before parents), then evaluate it bottom-up
        // (iteratively, no recursion). Postorder matters: reversed preorder
        // evaluates a parent before a shared child first reached through a
        // later branch, poisoning the parent to None (e.g. an Add whose
        // addend is also used deeper in the other addend's subtree).
        let mut seen: Set<OpId> = Set::default();
        let mut order: Vec<OpId> = Vec::new();
        let mut stack = vec![(op_id, false)];
        for _ in 0..10_000 {
            let Some((id, emit)) = stack.pop() else { break };
            if id.is_null() {
                continue;
            }
            if emit {
                order.push(id);
                continue;
            }
            if !seen.insert(id) {
                continue;
            }
            stack.push((id, true));
            match &self.ops[id].op {
                Op::Cast { x, .. } => stack.push((*x, false)),
                Op::Bitcast { x, .. } => stack.push((*x, false)),
                Op::Unary { x, .. } => stack.push((*x, false)),
                Op::Binary { x, y, .. } => {
                    stack.push((*x, false));
                    stack.push((*y, false));
                }
                Op::Stack { ops } => stack.extend(ops.iter().copied().map(|o| (o, false))),
                Op::Loop { len } => stack.push((*len, false)),
                &Op::Range { kind, .. } => match kind {
                    RangeKind::Group(len) => stack.push((len, false)),
                    RangeKind::Local(_) | RangeKind::Warp(_) => {}
                },
                Op::Mad { x, y, z } => {
                    stack.push((*x, false));
                    stack.push((*y, false));
                    stack.push((*z, false));
                }
                _ => {}
            }
        }
        if !stack.is_empty() {
            panic!("resolve_const did not finish in 10000 steps");
        }

        let mut values: Map<OpId, Option<Constant>> = Map::default();
        for &id in &order {
            let v = match &self.ops[id].op {
                Op::Const(c) => Some(*c),
                // A cast preserves the integer value of a length.
                Op::Cast { x, dtype } => values.get(x).copied().flatten().map(|v| v.cast(*dtype)),
                // A bitcast reinterprets the bits (equal widths asserted upstream).
                Op::Bitcast { x, dtype } => values.get(x).copied().flatten().map(|v| v.bitcast(*dtype)),
                Op::Unary { x, uop } => values.get(x).copied().flatten().map(|v| v.unary(*uop)),
                Op::Binary { x, y, bop } => {
                    values.get(x).copied().flatten().zip(values.get(y).copied().flatten()).map(|(a, b)| {
                        // Shape expressions may mix dtypes (e.g. a U32 group
                        // index with an I64 pad constant); evaluate in their
                        // least upper dtype.
                        let dt = a.dtype().least_upper_dtype(b.dtype());
                        Constant::binary(a.cast(dt), b.cast(dt), *bop)
                    })
                }
                // Fused multiply add: x * y + z, evaluated in the operands'
                // least upper dtype (same rule as Binary).
                Op::Mad { x, y, z } => values
                    .get(x)
                    .copied()
                    .flatten()
                    .zip(values.get(y).copied().flatten())
                    .zip(values.get(z).copied().flatten())
                    .map(|((a, b), c)| {
                        let dt = a.dtype().least_upper_dtype(b.dtype()).least_upper_dtype(c.dtype());
                        let prod = Constant::binary(a.cast(dt), b.cast(dt), BOp::Mul);
                        Constant::binary(prod, c.cast(dt), BOp::Add)
                    }),
                Op::Loop { len } => values.get(len).copied().flatten(),
                &Op::Range { kind, .. } => match kind {
                    RangeKind::Group(len) => values.get(&len).copied().flatten(),
                    RangeKind::Local(len) => Some(Constant::idx(len)),
                    // The warp op's value is the lane id — runtime data.
                    RangeKind::Warp(_) => None,
                },
                // Param is dynamic
                Op::Param { .. } => None,
                // A loaded value (e.g. a lowered variable: Binary{mask, Load})
                // is runtime data — never compile-time constant, but a
                // legitimate dynamic length (same as Param).
                Op::Load { .. } => None,
                ref op => todo!("{op:?}"),
            };
            values.insert(id, v);
        }
        values[&op_id]
    }

    /// Remap slab indices from x to y
    fn remap(&mut self, x: OpId, y: OpId) {
        for op_node in self.ops.values_mut() {
            for param in op_node.op.parameters_mut() {
                if *param == x {
                    *param = y;
                }
            }
        }
    }

    /// Add an operation to the kernel.
    pub(crate) fn push_back(&mut self, op: Op) -> OpId {
        let op_node = OpLinked { prev: self.tail, next: OpId::NULL, op };
        let op_id = self.ops.push(op_node);
        if self.head.is_null() {
            self.head = op_id;
        } else {
            self.ops[self.tail].next = op_id;
        }
        self.tail = op_id;
        op_id
    }

    /// Duplicates the dependency chain of `root_op` into a fresh kernel.
    ///
    /// # The duplication design (runtime vs kernelizer)
    ///
    /// Kernels can share their ops' results in two ways, and the two callers
    /// of this method use two different ownership models:
    ///
    /// - **Runtime** (eager tensors): ops never *consume* — they only create
    ///   new tensors. When a kernel must be split, the original kernel keeps
    ///   the output tensor in its `outputs` and keeps ALL of its ops; the
    ///   fresh kernel receives a **clone** of the root chain and *recomputes*
    ///   the same values from the shared input loads. Nothing is removed,
    ///   no refcount changes hands — the only bookkeeping is that every load
    ///   the fresh kernel reads gains one reference (retain per `new_loads`
    ///   occurrence). The fresh kernel's `outputs` stays empty: it produces
    ///   values only for the ops built on it afterwards.
    /// - **Kernelizer** (graph classes): ops DO consume edges (`rcs`,
    ///   `visited`, `outputs`). The split keeps the same observable
    ///   accounting as a move: one `outputs` occurrence of the child class
    ///   leaves the original kernel and is pushed onto the fresh kernel,
    ///   where the caller's `consume` immediately removes it. The original
    ///   kernel keeps its output registration for its remaining consumers;
    ///   if the split consumer was the last one, the original's orphaned
    ///   chain is cleaned up separately via [`Kernel::remove_unused_chain`].
    ///
    /// In both models this method is a **pure duplication**: `self` is left
    /// completely untouched (ops AND loads), and the fresh kernel contains
    /// exactly the ops `root_op` transitively depends on, cloned in topo
    /// order with remapped OpIds.
    ///
    /// Returns the fresh kernel, the remapped `root_op` OpId inside it, and
    /// `new_loads` — the caller's `loads` entries (parallel to the input
    /// `Param` ops, linked-list order) for the Params the fresh kernel
    /// cloned, in the fresh kernel's Param order. Every `new_loads` entry is
    /// an additional reader of the underlying tensor: the runtime caller
    /// retains one reference per occurrence.
    pub(crate) fn duplicate_subkernel<T: Copy>(&mut self, root_op: OpId, loads: &[T]) -> (Self, OpId, Vec<T>) {
        // Walk 1: from root_op
        let mut root_required = Set::default();
        let mut stack = vec![root_op];
        for _ in 0..10_000 {
            let Some(op) = stack.pop() else { break };
            if root_required.insert(op) {
                stack.extend(self.at(op).parameters());
            }
        }
        if !stack.is_empty() {
            panic!("duplicate_subkernel did not finish in 10000 steps");
        }

        // Partition loads: the fresh kernel receives the caller's `loads`
        // entries for every input Param it cloned. `self` keeps ALL of its
        // loads — it is untouched. `loads` is parallel to the kernel Param
        // ops. Kernels before linearize contain only Params — Storage
        // (accumulators etc.) is resolved later, so none may appear here.
        let mut new_loads: Vec<T> = Vec::new();
        let mut load_idx = 0;
        let mut oid = self.head;
        for _ in 0..10_000 {
            if oid.is_null() {
                break;
            }
            match self.at(oid) {
                // Global and Variable Params are loads; GlobalMut is a store and
                // is not part of the loads list.
                Op::Param { kind: ParamKind::Global | ParamKind::Variable, .. } => {
                    if root_required.contains(&oid) {
                        new_loads.push(loads[load_idx]);
                    }
                    load_idx += 1;
                }
                Op::Param { kind: ParamKind::GlobalMut, .. } => {}
                Op::Storage { .. } => {
                    panic!("duplicate_subkernel: unexpected Op::Storage (pre-linearize kernels contain only Params)")
                }
                _ => {}
            }
            oid = self.next_op(oid);
        }
        if !oid.is_null() {
            panic!("duplicate_subkernel did not finish in 10000 steps");
        }

        // Build new kernel by cloning root's ops (in topo order) with remapped OpIds
        let mut new_kernel = Kernel::from_device_id(self.dev, self.dev_info.clone());
        let mut remap: Map<OpId, OpId> =
            Map::with_capacity_and_hasher(root_required.len(), core::hash::BuildHasherDefault::default());
        let mut new_root_op = OpId::NULL;
        let mut old_id = self.head;
        for _ in 0..10_000 {
            if old_id.is_null() {
                break;
            }
            if root_required.contains(&old_id) {
                let mut op = self.at(old_id).clone();
                op.remap_params(&remap);
                let new_id = new_kernel.push_back(op);
                if old_id == root_op {
                    new_root_op = new_id;
                }
                remap.insert(old_id, new_id);
            }
            old_id = self.next_op(old_id);
        }
        if !old_id.is_null() {
            panic!("duplicate_subkernel did not finish in 10000 steps");
        }

        (new_kernel, new_root_op, new_loads)
    }

    /// Get all group indices used in the kernel.
    pub(crate) fn get_group_indices(&self) -> std::collections::BTreeMap<u32, OpId> {
        let mut indices = std::collections::BTreeMap::new();
        for (op_id, op_node) in self.ops.iter() {
            if let Op::Range { axis, kind: RangeKind::Group(_), .. } = op_node.op {
                indices.insert(axis, op_id);
            }
        }
        indices
    }

    /// Renumber indices to be in order.
    pub(crate) fn renumber_indices(&mut self) {
        let mut group_indices = BTreeMap::default();
        let mut local_indices = BTreeMap::default();
        for (op_id, op_node) in self.ops.iter() {
            match op_node.op {
                Op::Range { axis, kind: RangeKind::Group(_), .. } => group_indices.insert(axis, op_id),
                Op::Range { axis, kind: RangeKind::Local(_), .. } => local_indices.insert(axis, op_id),
                _ => None,
            };
        }
        let mut ax = 0;
        for &idx_id in group_indices.values() {
            let Op::Range { axis, kind: RangeKind::Group(_), .. } = &mut self.ops[idx_id].op else {
                unreachable!()
            };
            *axis = ax;
            ax += 1;
        }
        for &idx_id in local_indices.values() {
            let Op::Range { axis, kind: RangeKind::Local(_), .. } = &mut self.ops[idx_id].op else {
                unreachable!()
            };
            *axis = ax;
            ax += 1;
        }
    }

    pub(crate) fn is_preceded_by_reduce(&self, x: OpId) -> bool {
        //if self.ops.values().filter(|node| matches!(node.op, Op::Reduce { .. })).count() > 1 { return true; }
        let mut params = vec![x];
        let mut found = false;
        let mut seen: Set<OpId> = Set::default();
        for _ in 0..10_000 {
            let Some(param) = params.pop() else { break };
            if !seen.insert(param) {
                continue;
            }
            if let &Op::Reduce { x, .. } = self.at(param) {
                params = vec![x];
                found = true;
                break;
            }
            params.extend(self.ops[param].op.parameters());
        }
        if !found && !params.is_empty() {
            panic!("is_preceded_by_reduce did not finish in 10000 steps");
        }
        //if params.is_empty() { return false; }
        //println!("Found reduce at {params:?}");
        // If there is a load (non constant reduce) or multiple reduces, return true
        let mut seen: Set<OpId> = Set::default();
        for _ in 0..10_000 {
            let Some(param) = params.pop() else { break };
            if !seen.insert(param) {
                continue;
            }
            if matches!(self.ops[param].op, Op::Param { .. } | Op::Reduce { .. }) {
                return true;
            }
            params.extend(self.ops[param].op.parameters());
        }
        if !params.is_empty() {
            panic!("is_preceded_by_reduce did not finish in 10000 steps");
        }
        false
    }

    pub(crate) fn is_preceded_by_compute(&self, x: OpId) -> bool {
        let mut params = vec![x];
        let mut seen: Set<OpId> = Set::default();
        let (mut has_compute, mut has_param) = (false, false);
        for _ in 0..10_000 {
            let Some(param) = params.pop() else { break };
            if !seen.insert(param) {
                continue;
            }
            match &self.ops[param].op {
                Op::Binary { .. } | Op::Unary { .. } | Op::Reduce { .. } => {
                    has_compute = true;
                    params.extend(self.ops[param].op.parameters());
                }
                Op::Param { .. } => has_param = true,
                Op::Const(_) => {}
                _ => params.extend(self.ops[param].op.parameters()),
            }
        }
        if !params.is_empty() {
            panic!("is_preceded_by_compute did not finish in 10000 steps");
        }
        has_compute && has_param
    }
}