zyx 0.17.0

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

//! Loop unrolling optimization.
//!
//! This module provides loop unrolling optimizations for kernels,
//! which unroll loops to reduce loop overhead and enable better
//! instruction scheduling and vectorization.
//!
//! Loop unrolling is useful for:
//!
//! - Reducing loop overhead
//! - Enabling better instruction-level parallelism
//! - Improving vectorization opportunities
//! - Unrolling loops with constant lengths

use super::autotune::Optimization;
use crate::backend::DeviceInfo;
use crate::kernel::MemLayout;
#[allow(unused)]
use crate::{
    Map, Set,
    dtype::Constant,
    kernel::{BOp, Kernel, MemScope, Op, OpId, RangeKind},
    shape::Dim,
};

/// Unroll loops with a specific factor.
#[derive(Debug)]
pub struct UnrollLoops {
    /// Unroll factors to try for each loop.
    pub factors: Vec<u64>,
}

impl Optimization for UnrollLoops {
    fn nconfigs(&self) -> u64 {
        self.factors.len() as u64
    }

    fn apply(&self, kernel: &mut Kernel, config: u64) {
        let factor = self.factors[config as usize];
        if (kernel.ops.len().0 as usize) < 5000 {
            kernel.unroll_loops(factor as Dim);
        }
    }
}

/// Unroll loops with constant lengths.
#[derive(Debug)]
pub struct UnrollConstantLoops;

impl Optimization for UnrollConstantLoops {
    fn nconfigs(&self) -> u64 {
        1
    }

    fn apply(&self, kernel: &mut Kernel, _config: u64) {
        kernel.unroll_constant_loops();
    }
}

impl Kernel {
    /// Make the `UnrollLoops` optimization.
    pub fn opt_unroll(&self, _dev_info: &DeviceInfo) -> Box<dyn Optimization> {
        Box::new(UnrollLoops { factors: vec![8, 4, 16, 2] })
    }

    /// Make the `UnrollConstantLoops` optimization.
    pub fn opt_unroll_constant_loops(&self, _dev_info: &DeviceInfo) -> Box<dyn Optimization> {
        Box::new(UnrollConstantLoops)
    }

    /// Eliminate zero-length index operations.
    ///
    /// This method removes index operations with length 1,
    /// which can be optimized away.
    pub(crate) fn delete_zero_len_indices(&mut self) {
        #[cfg(feature = "time")]
        let _timer = crate::Timer::new("eliminate zero index");
        let node_ids: Vec<OpId> = self.ops.ids().collect();
        for id in node_ids.iter().copied() {
            let is_len1 = match &self.ops[id].op {
                Op::Range { kind, .. } => match kind {
                    RangeKind::Group(len) => self.resolve_const(*len).and_then(crate::dtype::Constant::as_dim) == Some(1),
                    RangeKind::Local(len) => *len == 1,
                    // A warp is a view over a local range — never a removable len-1 index.
                    RangeKind::Warp(_) => false,
                },
                _ => false,
            };
            if is_len1 {
                self.ops[id].op = Op::Const(Constant::idx(0));
            }
        }
        self.verify();
    }

    pub(crate) fn unroll_len1_loops(&mut self) {
        #[cfg(feature = "time")]
        let _timer = crate::Timer::new("eliminate zero index");
        let mut op_id = self.head;
        let mut depth = 0;
        let mut deleted = Vec::new();
        while !op_id.is_null() {
            let next = self.next_op(op_id);
            if let Op::Loop { len: len_id, .. } = self.ops[op_id].op {
                depth += 1;
                if self.resolve_const(len_id).and_then(crate::dtype::Constant::as_dim) == Some(1) {
                    self.ops[op_id].op = Op::Const(Constant::idx(0));
                    deleted.push(depth);
                }
            } else if self.ops[op_id].op == Op::EndLoop {
                if deleted.pop_if(|x| *x == depth).is_some() {
                    self.remove_op(op_id);
                }
                depth -= 1;
            }
            op_id = next;
        }
        self.verify();
    }

    pub(crate) fn unroll_loops(&mut self, unroll_dim: Dim) {
        #[cfg(feature = "time")]
        let _timer = crate::Timer::new("unroll_loops");
        let mut endloop_ids = Vec::new();
        let mut op_id = self.tail;
        while !op_id.is_null() {
            let prev = self.prev_op(op_id);
            if self.ops[op_id].op == Op::EndLoop {
                endloop_ids.push(op_id);
            }
            if let Op::Loop { len: len_id, .. } = self.ops[op_id].op {
                let Some(len) = self.resolve_const(len_id).and_then(crate::dtype::Constant::as_dim) else {
                    op_id = prev;
                    continue;
                };
                let _ = endloop_ids.pop().unwrap();
                if len as usize <= unroll_dim as usize
                    && self.ops.len().0 as usize + (self.n_ops_in_loop(op_id) * (len as usize - 1)) < 5_000
                {
                    self.unroll_loop(op_id);
                }
            }
            op_id = prev;
        }
    }

    fn n_ops_in_loop(&self, loop_id: OpId) -> usize {
        let mut op_id = self.next_op(loop_id);
        let mut n_loops = 1;
        let mut n_ops = 0;
        while !op_id.is_null() {
            match self.ops[op_id].op {
                Op::Loop { .. } => {
                    n_loops += 1;
                }
                Op::EndLoop => {
                    n_loops -= 1;
                    if n_loops == 0 {
                        return n_ops;
                    }
                }
                _ => {}
            }
            n_ops += 1;
            op_id = self.next_op(op_id);
        }
        n_ops
    }

    pub(crate) fn unroll_constant_loops(&mut self) {
        #[cfg(feature = "time")]
        let _timer = crate::Timer::new("unroll_constant_loops");
        let mut endloop_ids = Vec::new();
        let mut op_id = self.tail;
        let mut constant_loops = vec![true];
        while !op_id.is_null() {
            let prev = self.prev_op(op_id);
            match self.ops[op_id].op {
                Op::EndLoop => {
                    endloop_ids.push(op_id);
                    constant_loops.push(true);
                }
                Op::Loop { len: len_id, .. } => {
                    let Some(len) = self.resolve_const(len_id).and_then(crate::dtype::Constant::as_dim) else {
                        continue;
                    };
                    endloop_ids.pop().unwrap();
                    let is_const = constant_loops.pop().unwrap();
                    if !is_const && let Some(inner_loop) = constant_loops.last_mut() {
                        *inner_loop = false;
                    }
                    if len == 1
                        || (is_const && self.ops.len().0 as usize + (self.n_ops_in_loop(op_id) * (len as usize - 1)) < 5_000)
                    {
                        self.unroll_loop(op_id);
                    }
                }
                Op::Store { dst, .. } => {
                    let Op::Storage { scope, .. } = self.ops[dst].op else {
                        unreachable!()
                    };
                    if scope != MemScope::Register {
                        *constant_loops.last_mut().unwrap() = false;
                    }
                }
                Op::Load { src, .. } => {
                    let Op::Storage { scope, .. } = self.ops[src].op else {
                        unreachable!()
                    };
                    if scope != MemScope::Register {
                        *constant_loops.last_mut().unwrap() = false;
                    }
                }
                _ => {}
            }
            op_id = prev;
        }
    }

    /// Unroll a loop.
    ///
    /// Unrolls the loop at `loop_id` to reduce loop overhead and
    /// enable better instruction scheduling and vectorization.
    pub fn unroll_loop(&mut self, loop_id: OpId) {
        let Op::Loop { len: len_id } = self.ops[loop_id].op else {
            return;
        };
        let Some(len) = self.resolve_const(len_id).and_then(crate::dtype::Constant::as_dim) else {
            return;
        };
        //println!("UNROLL len={} limit={}", len, len > 64);
        if len == 0 || len > 64 {
            return;
        }

        let mut loop_depth = 1;
        let mut endloop_id = self.next_op(loop_id);
        while !endloop_id.is_null() {
            match self.ops[endloop_id].op {
                Op::Loop { .. } => {
                    loop_depth += 1;
                }
                Op::EndLoop => {
                    loop_depth -= 1;
                    if loop_depth == 0 {
                        break;
                    }
                }
                _ => {}
            }
            endloop_id = self.next_op(endloop_id);
        }

        self.ops[loop_id].op = Op::Const(Constant::idx(0));
        let last_loop_op = self.prev_op(endloop_id);

        for idx in 1..len {
            let mut new_ops_map = Map::default();
            let idx_op = self.insert_before(endloop_id, Op::Const(Constant::idx(idx)));
            new_ops_map.insert(loop_id, idx_op);

            let mut op_id = self.next_op(loop_id);
            loop {
                let mut op = self.ops[op_id].op.clone();
                for param in op.parameters_mut() {
                    if let Some(&new_param) = new_ops_map.get(param) {
                        *param = new_param;
                    }
                }
                let new_op_id = self.insert_before(endloop_id, op);
                new_ops_map.insert(op_id, new_op_id);

                if op_id == last_loop_op {
                    break;
                }
                op_id = self.next_op(op_id);
            }
        }
        self.remove_op(endloop_id);

        self.verify();
    }

    /// Unrolls loop:
    /// acc a
    /// for 0..32 {
    ///     z = ...
    ///     x = load a
    ///     x1 = x + z
    ///     store a <- x1
    /// }
    ///
    /// Using tree reduce unroll to:
    ///
    /// acc a
    /// for 0..8 {
    ///     z = ...
    ///     z1 = ...
    ///     z2 = ...
    ///     z3 = ...
    ///     x = load a
    ///     x1 = x + z
    ///     x2 = x1 + z1
    ///     x3 = x2 + z2
    ///     x4 = x3 + z3
    ///     store a <- x4
    /// }
    ///
    pub fn unroll_tree_reduce(&mut self, loop_id: OpId, factor: Dim) {
        #[cfg(feature = "time")]
        let _timer = crate::Timer::new("unroll_tree_reduce");
        let Op::Loop { len: len_id } = self.ops[loop_id].op else {
            return;
        };
        let Some(len) = self.resolve_const(len_id).and_then(crate::dtype::Constant::as_dim) else {
            return;
        };
        if factor < 2 || len % factor as Dim != 0 {
            return;
        }

        if self.ops.len().0 as i64 * factor > 5000 {
            return;
        }

        //println!("Unroll tree reduce for loop={loop_id}, factor={factor}");

        // Find the acc
        let acc_id;
        let mut op_id = self.prev_op(loop_id);
        loop {
            if op_id.is_null() {
                return; // not a reduce, just a loop
            }
            match self.ops[op_id].op {
                Op::Loop { .. } => return, // nested reduce or no reduce
                Op::Storage { scope: MemScope::Register, .. } => {
                    acc_id = op_id;
                    break;
                }
                _ => {}
            }
            op_id = self.prev_op(op_id);
        }

        // Find the store to acc
        let mut op_id = acc_id;
        let acc_init;
        loop {
            if let Op::Store { dst, src: x, .. } = self.ops[op_id].op
                && dst == acc_id
            {
                acc_init = x;
                break;
            }
            op_id = self.next_op(op_id);
            if op_id == loop_id {
                unreachable!();
            }
        }

        let mut has_store = false;
        let mut op_id = self.next_op(loop_id);
        let endloop_id;
        loop {
            match self.ops[op_id].op {
                Op::Loop { .. } => return, // no nested loops
                Op::Store { layout, .. } => {
                    if has_store || layout != MemLayout::Scalar {
                        return;
                    }
                    has_store = true;
                }
                Op::EndLoop => {
                    endloop_id = op_id;
                    break;
                }
                _ => {}
            }
            op_id = self.next_op(op_id);
        }

        let mut map = Map::default();

        let new_len = self.insert_const_idx_before(loop_id, len / factor);
        let new_loop = self.insert_before(loop_id, Op::Loop { len: new_len });
        let mut op_id = self.next_op(loop_id);
        let stride = self.insert_before(loop_id, Op::Const(Constant::idx(factor)));
        self.ops[loop_id].op = Op::Binary { x: new_loop, y: stride, bop: BOp::Mul };
        let mut new_ones = Vec::with_capacity(factor as usize - 1);
        for i in 1..factor {
            let offset = self.insert_before(op_id, Op::Const(Constant::idx(i)));
            let new_id = self.insert_before(op_id, Op::Binary { x: loop_id, y: offset, bop: BOp::Add });
            new_ones.push(new_id);
        }
        map.insert(loop_id, new_ones);

        while op_id != endloop_id {
            let this_id = op_id;
            op_id = self.next_op(op_id);
            if let Op::Load { src, index: _, layout: MemLayout::Scalar } = self.ops[this_id].op
                && src == acc_id
            {
                // TODO debug assert index is const zero
                map.insert(this_id, vec![acc_init; factor as usize - 1]);
            } else if let Op::Store { dst, src: x, index, layout: MemLayout::Scalar } = self.ops[this_id].op
                && dst == acc_id
            {
                let Op::Binary { bop, .. } = self.ops[x].op else {
                    unreachable!()
                };
                // TODO debug assert index is const zero
                let y = if let Some(mapping) = map.get(&x) { mapping[0] } else { x };
                let mut carry = this_id;
                self.ops[this_id].op = Op::Binary { x, y, bop };
                for i in 1..factor - 1 {
                    let x = if let Some(mapping) = map.get(&x) {
                        mapping[i as usize]
                    } else {
                        x
                    };
                    carry = self.insert_before(op_id, Op::Binary { x, y: carry, bop });
                }
                self.insert_before(op_id, Op::Store { dst, src: carry, index, layout: MemLayout::Scalar });
            } else {
                let mut new_ones = Vec::with_capacity(factor as usize - 1);
                for i in 1..factor {
                    let mut new_op = self.ops[this_id].op.clone();
                    for param in new_op.parameters_mut() {
                        if let Some(mapping) = map.get(param) {
                            *param = mapping[i as usize - 1];
                        }
                    }
                    let new_id = self.insert_before(op_id, new_op);
                    new_ones.push(new_id);
                }
                map.insert(this_id, new_ones);
            }
        }

        self.verify();
    }
}