zyx 0.17.0

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0
use std::collections::BTreeSet;

use crate::{
    Map, Set, ZyxError,
    backend::{Buffer, LaunchArg, Pool, ProgramId},
    dtype::Constant,
    graph::{Graph, Node, OpId},
    kernel::BOp,
    runtime::Runtime,
    shape::Dim,
};

/// One dim of an allocation spec: an expression tree over compile-time
/// constants and leaf-class values (variables bound between plan runs).
/// Computation stays symbolic so a plan compiled once serves any variable
/// values; evaluation happens at execution time.
#[derive(Debug, Clone)]
pub enum PlanDim {
    Const(Dim),
    Leaf(OpId),
    Binary { x: Box<PlanDim>, y: Box<PlanDim>, bop: BOp },
    Cast { x: Box<PlanDim>, dtype: crate::dtype::DType },
}

impl PlanDim {
    /// Evaluate the dim expression against the leaf classes' scalar values.
    /// Fails loudly on an unbound leaf — a missing value is a bug, never a
    /// default.
    fn eval(&self, class_vars: &Map<OpId, Constant>) -> Dim {
        match self {
            PlanDim::Const(c) => *c,
            PlanDim::Leaf(cid) => class_vars
                .get(cid)
                .and_then(|c| c.as_dim())
                .unwrap_or_else(|| panic!("dynamic dim class {cid:?} is unbound at execution time")),
            PlanDim::Binary { x, y, bop } => {
                Constant::binary(Constant::idx(x.eval(class_vars)), Constant::idx(y.eval(class_vars)), *bop)
                    .as_dim()
                    .unwrap_or_else(|| panic!("dim binary op {bop:?} did not produce a dim"))
            }
            PlanDim::Cast { x, dtype } => Constant::idx(x.eval(class_vars))
                .cast(*dtype)
                .as_dim()
                .unwrap_or_else(|| panic!("dim cast to {dtype:?} did not produce a dim")),
        }
    }
}

#[derive(Debug, Clone)]
pub enum ExecNode {
    Allocate {
        class: OpId,
        pool: Pool,
        dtype_size: Dim,
        /// One dim expression per shape axis; the buffer is sized by their
        /// product, evaluated at execution time.
        dims: Vec<PlanDim>,
    },
    Copy {
        dst_class: OpId,
        src_class: OpId,
    },
    Deallocate {
        class: OpId,
    },
    Launch {
        program_id: ProgramId,
        load_classes: Box<[OpId]>,
        store_classes: Box<[OpId]>,
    },
    // Binds class_buf[class] = class_buf[to]: an After output aliases the
    // buffer of its base leaf class (in-place assign write). Preplanned by
    // ExecPlan::new so execute_plan only resolves buffers, never decides.
    Alias {
        class: OpId,
        to: OpId,
    },
}

#[derive(Debug, Clone)]
pub struct ExecPlan {
    pub nodes: Vec<ExecNode>,
    pub leaf_classes: Vec<OpId>,
    // Pool each leaf class lived in when the plan was compiled. Leaf pools
    // must not vary across plan reuse, or the preplanned Alias/Allocate/Copy
    // binding would be wrong — debug-asserted in execute_plan.
    pub leaf_pools: Map<OpId, Pool>,
}

impl ExecPlan {
    #[must_use]
    pub fn new(graph: &Graph, nodes: &[OpId], output_set: &BTreeSet<OpId>, leaf_pools: &Map<OpId, Pool>) -> Self {
        let mut rc: Map<OpId, u32> = Map::default();
        for &nid in nodes {
            match &graph.nodes[nid].node {
                Node::Kernel { inputs, .. } => {
                    for &ic in &**inputs {
                        rc.entry(ic).and_modify(|c| *c += 1).or_insert(1);
                    }
                }
                Node::ToDevice { x, .. } => {
                    rc.entry(*x).and_modify(|c| *c += 1).or_insert(1);
                }
                _ => unreachable!(),
            }
        }

        let mut plan_nodes = Vec::new();
        let mut allocated: Set<OpId> = Set::default();

        // Allocation spec of a class: dtype byte size and one `PlanDim` per
        // shape axis. Dim expressions over leaf classes stay symbolic — their
        // values live in leaf buffers set between plan runs — so execution
        // evaluates the tree and multiplies. Expression trees over Const and
        // leaf dims must terminate the walk; anything else is unreachable.
        fn alloc_spec(graph: &Graph, class: OpId) -> (Dim, Vec<PlanDim>) {
            fn dim_expr(graph: &Graph, dim: OpId) -> PlanDim {
                match graph.nodes[dim].node {
                    Node::Const { value: c, .. } => {
                        PlanDim::Const(c.as_dim().unwrap_or_else(|| panic!("dim class {dim:?} is not a constant")))
                    }
                    Node::Leaf { .. } => PlanDim::Leaf(dim),
                    Node::Binary { x, y, bop } => {
                        PlanDim::Binary { x: Box::new(dim_expr(graph, x)), y: Box::new(dim_expr(graph, y)), bop }
                    }
                    Node::Cast { x, dtype } => PlanDim::Cast { x: Box::new(dim_expr(graph, x)), dtype },
                    ref op => unreachable!("alloc dim class {dim:?} must be a dim over Const/leaf leaves, got {op:?}"),
                }
            }
            let dtype_size = Dim::from(graph.dtype(class).bit_size() / 8);
            let dims = graph.shape(class).iter().map(|&d| dim_expr(graph, d)).collect();
            (dtype_size, dims)
        }

        // After output classes alias the buffer of x's base leaf class: the
        // assign writes the new buffer version in-place into that leaf buffer,
        // so an After class (x's value after the assign) shares the leaf's
        // buffer. They must not be allocated or deallocated — the leaf's buffer
        // is owned by the realized tensor.
        let mut aliases: Vec<(OpId, OpId, Dim, Vec<PlanDim>)> = Vec::new();
        let mut alias_classes: Set<OpId> = Set::default();
        for (cid, nd) in graph.nodes.iter().filter(|(id, nd)| nd.class_of == *id) {
            if let Node::After { x, .. } = nd.node {
                let base = graph.base_leaf(x);
                let (dtype_size, dims) = alloc_spec(graph, cid);
                aliases.push((cid, base, dtype_size, dims));
                alias_classes.insert(cid);
            }
        }

        // Pool of the kernel that stores each alias class — precomputed so the
        // binding below is decided at plan time, not execution time.
        let mut store_pool: Map<OpId, Pool> = Map::default();
        for &nid in nodes {
            if let Node::Kernel { outputs, program_id, .. } = &graph.nodes[nid].node {
                let pool = program_id.dev.pool();
                for &oc in &**outputs {
                    store_pool.insert(oc, pool);
                }
            }
        }

        // Bind aliases before any kernel runs. A leaf in the same pool as its
        // assign kernel binds straight to the leaf buffer. A cross-pool leaf
        // needs one kernel-pool copy of itself shared by every alias of that
        // leaf — chained assigns must write the same physical buffer or the
        // intermediate writes are lost. Mirrors eager assign's store-to-target
        // pool handling.
        let mut leaf_copy: Map<OpId, OpId> = Map::default();
        for &(class, to, dtype_size, ref dims) in &aliases {
            match store_pool.get(&class) {
                Some(pool) if leaf_pools[&to] != *pool => {
                    let owner = *leaf_copy.entry(to).or_insert_with(|| {
                        plan_nodes.push(ExecNode::Allocate { class, pool: *pool, dtype_size, dims: dims.clone() });
                        plan_nodes.push(ExecNode::Copy { dst_class: class, src_class: to });
                        class
                    });
                    if owner != class {
                        plan_nodes.push(ExecNode::Alias { class, to: owner });
                    }
                }
                _ => plan_nodes.push(ExecNode::Alias { class, to }),
            }
        }

        for &nid in nodes {
            match &graph.nodes[nid].node {
                Node::Kernel { inputs, outputs, program_id, .. } => {
                    let pool = program_id.dev.pool();
                    for &oc in &**outputs {
                        if !allocated.insert(oc) {
                            continue;
                        }
                        // Realized leaves and after aliases already have buffers
                        // (leaf buffers via leaf_map, aliases share x's leaf
                        // buffer) — never allocate fresh buffers for them.
                        if !graph.leaf_map.contains_key(&oc) && !alias_classes.contains(&oc) {
                            let (dtype_size, dims) = alloc_spec(graph, oc);
                            plan_nodes.push(ExecNode::Allocate { class: oc, pool, dtype_size, dims });
                        }
                    }
                    plan_nodes.push(ExecNode::Launch {
                        program_id: *program_id,
                        load_classes: inputs.clone(),
                        store_classes: outputs.clone(),
                    });
                    for &ic in &**inputs {
                        let c = rc.get_mut(&ic).unwrap();
                        *c -= 1;
                        if *c == 0
                            && !graph.leaf_map.contains_key(&ic)
                            && !output_set.contains(&ic)
                            && !alias_classes.contains(&ic)
                        {
                            plan_nodes.push(ExecNode::Deallocate { class: ic });
                        }
                    }
                }
                &Node::ToDevice { x, device, .. } => {
                    // Pool is always derived from the device, never the reverse.
                    let pool = device.pool();
                    let class_of = graph.nodes[nid].class_of;
                    if allocated.insert(class_of) && !graph.leaf_map.contains_key(&class_of) && !alias_classes.contains(&class_of)
                    {
                        let (dtype_size, dims) = alloc_spec(graph, class_of);
                        plan_nodes.push(ExecNode::Allocate { class: class_of, pool, dtype_size, dims });
                    }
                    plan_nodes.push(ExecNode::Copy { dst_class: class_of, src_class: x });
                    let c = rc.get_mut(&x).unwrap();
                    *c -= 1;
                    if *c == 0 && !graph.leaf_map.contains_key(&x) && !output_set.contains(&x) && !alias_classes.contains(&x) {
                        plan_nodes.push(ExecNode::Deallocate { class: x });
                    }
                }
                _ => unreachable!(),
            }
        }

        // Deallocate kernel outputs that are neither consumed by any node nor
        // requested outputs (e.g. the extra stores of a multi-output kernel).
        let allocated: Vec<OpId> = allocated.iter().copied().collect();
        for c in allocated {
            if !graph.leaf_map.contains_key(&c) && !output_set.contains(&c) && !alias_classes.contains(&c) && !rc.contains_key(&c)
            {
                plan_nodes.push(ExecNode::Deallocate { class: c });
            }
        }

        Self { nodes: plan_nodes, leaf_classes: graph.leaf_classes.clone(), leaf_pools: leaf_pools.clone() }
    }

    #[allow(unused)]
    pub fn debug(&self) {
        let line = "─".repeat(60);
        println!("\n{}", line);
        println!("  ExecPlan");
        println!("{}", line);
        for node in &self.nodes {
            match node {
                ExecNode::Allocate { class, pool, dtype_size, dims } => {
                    println!("  Allocate class={class:?} pool={pool:?} dtype_size={dtype_size} dims={dims:?}");
                }
                ExecNode::Copy { dst_class, src_class } => {
                    println!("  Copy dst={dst_class:?} src={src_class:?}");
                }
                ExecNode::Deallocate { class } => {
                    println!("  Deallocate class={class:?}");
                }
                ExecNode::Launch { program_id, load_classes, store_classes } => {
                    println!("  Launch prog={program_id:?} loads={load_classes:?} stores={store_classes:?}");
                }
                ExecNode::Alias { class, to } => {
                    println!("  Alias class={class:?} -> to={to:?}");
                }
            }
        }
        println!("{}\n", line);
    }
}

impl Runtime {
    pub fn execute_plan(
        &mut self,
        cache_key: u64,
        class_buf: &mut Map<OpId, Buffer>,
        class_vars: &Map<OpId, Constant>,
    ) -> Result<(), ZyxError> {
        let plan = self.plan_cache.get(&cache_key).unwrap();

        #[cfg(debug_assertions)]
        {
            for (&cid, &pool) in &plan.leaf_pools {
                debug_assert_eq!(
                    class_buf[&cid].pool, pool,
                    "leaf class {cid:?} moved pools since the plan was compiled — preplanned \
                     Alias/Allocate/Copy binding would be wrong"
                );
            }
        }

        for node in &plan.nodes {
            match node {
                ExecNode::Allocate { class, pool, dtype_size, dims } => {
                    // Evaluate the dim expressions against the leaf classes'
                    // scalar values (variables bound between plan runs), then
                    // size the buffer: one element per dim-product element,
                    // plus one extra trash element.
                    let mut elements: Dim = 1;
                    for dim in dims {
                        let v = dim.eval(class_vars);
                        debug_assert!(v > 0, "dim of class {class:?} evaluated to non-positive value {v}");
                        elements *= v;
                    }
                    debug_assert!(elements > 0, "allocation for class {class:?} would be empty ({elements} elements)");
                    let bytes = (elements + 1) * dtype_size;
                    let buf = pool.allocate(bytes)?;
                    let buf_id = Buffer { pool: *pool, buffer_id: buf };
                    class_buf.insert(*class, buf_id);
                }
                ExecNode::Launch { program_id, load_classes, store_classes } => {
                    let mut args = Vec::new();
                    let mut kernel_bufs = BTreeSet::new();
                    for c in load_classes.iter().chain(store_classes.iter()) {
                        if let Some(&value) = class_vars.get(c) {
                            // Variable leaf: bound from variable_map, no buffer.
                            args.push(LaunchArg::Variable(value));
                            continue;
                        }
                        let Some(buf) = class_buf.get(c) else {
                            panic!(
                                "DEBUG launch: class {c:?} (program {program_id:?}) has no allocated buffer; load_classes={load_classes:?}, store_classes={store_classes:?}"
                            );
                        };
                        args.push(LaunchArg::Buffer(buf.buffer_id));
                        kernel_bufs.insert(*buf);
                    }
                    if crate::debug_mask().dev() {
                        println!("launching kernel {program_id:?}");
                    }
                    program_id.dev.launch(program_id.program_id, &args)?;
                }
                ExecNode::Copy { dst_class, src_class } => {
                    let src = class_buf[src_class];
                    let dst = class_buf[dst_class];
                    debug_assert_ne!(src.pool, dst.pool);
                    // Cross-pool transfer. Event bookkeeping (barrier events
                    // on the source, deferred foreign release) is handled
                    // inside the receiving pool's worker.
                    dst.pool.pool_to_pool(src.pool, src.buffer_id, dst.buffer_id)?;
                }
                ExecNode::Deallocate { class } => {
                    let buf = class_buf.remove(class).unwrap();
                    buf.pool.release(buf.buffer_id);
                }
                ExecNode::Alias { class, to } => {
                    let buf = class_buf[to];
                    class_buf.insert(*class, buf);
                }
            }
        }

        Ok(())
    }
}