joule_mir/

ndarray_fusion.rs

//! NDArray expression fusion pass.
//!
//! Identifies chains of element-wise ndarray operations that produce single-use
//! temporaries, and fuses them into a single loop. For example:
//!
//! ```text
//! // Before fusion:
//! let t1 = a.add(&b);    // allocates + loops
//! let t2 = t1.mul(&c);   // allocates + loops
//! let result = t2.sub(&d); // allocates + loops
//!
//! // After fusion (conceptually):
//! let result = fused_loop(a, b, c, d, |a, b, c, d| (a + b) * c - d);
//! ```
//!
//! This eliminates intermediate allocations and reduces memory traffic from
//! 3 full passes over data to 1 pass. The energy savings are substantial:
//! each eliminated pass saves ~5 pJ per cacheline of data.
//!
//! The fusion pass works at the MIR level by:
//! 1. Identifying `Terminator::Call` sequences targeting ndarray element-wise functions
//! 2. Building a `FusedExpr` tree from the chain
//! 3. Replacing the chain with a single fused call
//!
//! The C codegen then emits the fused call as one loop with the composed expression.

use std::collections::HashMap;

/// A fused expression tree representing composed element-wise operations.
#[derive(Debug, Clone)]
pub enum FusedExpr {
    /// Reference to an input source array (index into sources list)
    Source(usize),
    /// Binary operation on two sub-expressions
    BinaryOp {
        op: FusedBinOp,
        left: Box<FusedExpr>,
        right: Box<FusedExpr>,
    },
    /// Unary operation on a sub-expression
    UnaryOp {
        op: FusedUnaryOp,
        operand: Box<FusedExpr>,
    },
    /// Scalar constant
    Scalar(f64),
}

/// Binary operations that can be fused
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FusedBinOp {
    Add,
    Sub,
    Mul,
    Div,
}

/// Unary operations that can be fused
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FusedUnaryOp {
    Neg,
    Abs,
}

impl FusedBinOp {
    /// Get the C operator symbol for this operation
    pub fn c_op(&self) -> &'static str {
        match self {
            Self::Add => "+",
            Self::Sub => "-",
            Self::Mul => "*",
            Self::Div => "/",
        }
    }

    /// Parse from ndarray function name suffix
    pub fn from_suffix(s: &str) -> Option<Self> {
        match s {
            "add" => Some(Self::Add),
            "sub" => Some(Self::Sub),
            "mul" => Some(Self::Mul),
            "div" => Some(Self::Div),
            _ => None,
        }
    }
}

/// A fusion opportunity: a chain of element-wise ops that can be fused.
#[derive(Debug, Clone)]
pub struct FusionChain {
    /// The fused expression tree
    pub expr: FusedExpr,
    /// Source array locals (in order of Source(0), Source(1), ...)
    pub sources: Vec<crate::Local>,
    /// The result local where the fused output is stored
    pub result_local: crate::Local,
    /// The element type (C identifier, e.g. "double")
    pub elem_c: String,
    /// Rank of the ndarray
    pub rank: u32,
    /// Intermediate locals that can be eliminated after fusion
    pub eliminated_locals: Vec<crate::Local>,
}

/// Analyze a function's MIR to find fusion opportunities.
///
/// Returns a list of fusion chains found. Each chain represents a sequence
/// of element-wise operations that can be replaced by a single fused loop.
pub fn find_fusion_chains(func: &crate::FunctionMIR) -> Vec<FusionChain> {
    // Map from local → the ndarray call that defines it
    let mut definitions: HashMap<crate::Local, NdarrayCallInfo> = HashMap::new();
    // Map from local → number of uses
    let mut use_counts: HashMap<crate::Local, usize> = HashMap::new();

    // Pass 1: collect ndarray element-wise call definitions and count uses
    for bb in &func.basic_blocks {
        // Count uses in statements
        for stmt in &bb.statements {
            if let crate::Statement::Assign { rvalue, .. } = stmt {
                count_rvalue_uses(rvalue, &mut use_counts);
            }
        }

        // Check terminator for ndarray calls
        if let crate::Terminator::Call {
            func_name: Some(ref name),
            ref args,
            ref destination,
            ..
        } = bb.terminator
        {
            let name_str = name.as_str();
            let dest_local = destination.local;
            if let Some(info) = parse_ndarray_elementwise_call(&name_str, args, dest_local) {
                definitions.insert(dest_local, info);
            }
            // Count args as uses
            for arg in args {
                if let crate::Operand::Move(place) | crate::Operand::Copy(place) = arg {
                    *use_counts.entry(place.local).or_insert(0) += 1;
                }
            }
        }
    }

    // Pass 2: build chains by following single-use temporaries backward
    let mut chains = Vec::new();
    let mut visited: std::collections::HashSet<crate::Local> = std::collections::HashSet::new();

    for (&local, info) in &definitions {
        if visited.contains(&local) {
            continue;
        }
        // Only start chains from results that are NOT single-use intermediates
        // (i.e., they're used more than once, or they're the final result)
        if use_counts.get(&local).copied().unwrap_or(0) <= 1
            && definitions.values().any(|d| d.args.contains(&local))
        {
            continue; // This is an intermediate — will be caught by another chain
        }

        let mut sources = Vec::new();
        let mut eliminated = Vec::new();
        let expr = build_fused_expr(
            local,
            &definitions,
            &use_counts,
            &mut sources,
            &mut eliminated,
            &mut visited,
        );

        // Only worth fusing if we eliminated at least one intermediate
        if !eliminated.is_empty() {
            chains.push(FusionChain {
                expr,
                sources,
                result_local: local,
                elem_c: info.elem_c.clone(),
                rank: info.rank,
                eliminated_locals: eliminated,
            });
        }
    }

    chains
}

/// Info about an ndarray element-wise call
#[derive(Debug, Clone)]
struct NdarrayCallInfo {
    op: FusedBinOp,
    args: Vec<crate::Local>,
    elem_c: String,
    rank: u32,
}

/// Parse a Terminator::Call to see if it's an ndarray element-wise operation
fn parse_ndarray_elementwise_call(
    name: &str,
    args: &[crate::Operand],
    _destination: crate::Local,
) -> Option<NdarrayCallInfo> {
    // Pattern: JouleNDArray_{elem}_{rank}_{op}
    let prefix = "JouleNDArray_";
    if !name.starts_with(prefix) {
        return None;
    }
    let rest = &name[prefix.len()..];

    // Find the last underscore to get the operation name
    let last_sep = rest.rfind('_')?;
    let op_name = &rest[last_sep + 1..];
    let type_part = &rest[..last_sep];

    let op = FusedBinOp::from_suffix(op_name)?;

    // Parse elem and rank from type_part (e.g., "double_2")
    let rank_sep = type_part.rfind('_')?;
    let elem_c = type_part[..rank_sep].to_string();
    let rank: u32 = type_part[rank_sep + 1..].parse().ok()?;

    // Must have exactly 2 args for binary op
    if args.len() != 2 {
        return None;
    }

    let arg_locals: Vec<_> = args
        .iter()
        .filter_map(|a| match a {
            crate::Operand::Move(p) | crate::Operand::Copy(p) => Some(p.local),
            _ => None,
        })
        .collect();

    if arg_locals.len() != 2 {
        return None;
    }

    Some(NdarrayCallInfo {
        op,
        args: arg_locals,
        elem_c,
        rank,
    })
}

/// Recursively build a fused expression tree from a chain of operations
fn build_fused_expr(
    local: crate::Local,
    definitions: &HashMap<crate::Local, NdarrayCallInfo>,
    use_counts: &HashMap<crate::Local, usize>,
    sources: &mut Vec<crate::Local>,
    eliminated: &mut Vec<crate::Local>,
    visited: &mut std::collections::HashSet<crate::Local>,
) -> FusedExpr {
    visited.insert(local);

    if let Some(info) = definitions.get(&local) {
        // Check if both args can be fused (single-use intermediates)
        let left_local = info.args[0];
        let right_local = info.args[1];

        let left = if definitions.contains_key(&left_local)
            && use_counts.get(&left_local).copied().unwrap_or(0) == 1
            && !visited.contains(&left_local)
        {
            eliminated.push(left_local);
            build_fused_expr(left_local, definitions, use_counts, sources, eliminated, visited)
        } else {
            make_source(left_local, sources)
        };

        let right = if definitions.contains_key(&right_local)
            && use_counts.get(&right_local).copied().unwrap_or(0) == 1
            && !visited.contains(&right_local)
        {
            eliminated.push(right_local);
            build_fused_expr(right_local, definitions, use_counts, sources, eliminated, visited)
        } else {
            make_source(right_local, sources)
        };

        FusedExpr::BinaryOp {
            op: info.op,
            left: Box::new(left),
            right: Box::new(right),
        }
    } else {
        make_source(local, sources)
    }
}

fn make_source(local: crate::Local, sources: &mut Vec<crate::Local>) -> FusedExpr {
    if let Some(idx) = sources.iter().position(|&l| l == local) {
        FusedExpr::Source(idx)
    } else {
        let idx = sources.len();
        sources.push(local);
        FusedExpr::Source(idx)
    }
}

fn count_rvalue_uses(rvalue: &crate::Rvalue, counts: &mut HashMap<crate::Local, usize>) {
    match rvalue {
        crate::Rvalue::Use(op) | crate::Rvalue::UnaryOp { operand: op, .. } => {
            count_operand_uses(op, counts);
        }
        crate::Rvalue::BinaryOp { left, right, .. } => {
            count_operand_uses(left, counts);
            count_operand_uses(right, counts);
        }
        crate::Rvalue::Ref { place, .. } => {
            *counts.entry(place.local).or_insert(0) += 1;
        }
        crate::Rvalue::Aggregate { operands, .. } => {
            for op in operands {
                count_operand_uses(op, counts);
            }
        }
        crate::Rvalue::Cast { operand, .. } => {
            count_operand_uses(operand, counts);
        }
        crate::Rvalue::Discriminant { place } | crate::Rvalue::Len { place } => {
            *counts.entry(place.local).or_insert(0) += 1;
        }
        _ => {}
    }
}

fn count_operand_uses(op: &crate::Operand, counts: &mut HashMap<crate::Local, usize>) {
    if let crate::Operand::Move(place) | crate::Operand::Copy(place) = op {
        *counts.entry(place.local).or_insert(0) += 1;
    }
}

/// Generate the C expression string for a fused expression tree.
/// The sources are accessed as `src0->data[i]`, `src1->data[i]`, etc.
pub fn fused_expr_to_c(expr: &FusedExpr) -> String {
    match expr {
        FusedExpr::Source(idx) => format!("src{idx}->data[i]"),
        FusedExpr::BinaryOp { op, left, right } => {
            format!(
                "({} {} {})",
                fused_expr_to_c(left),
                op.c_op(),
                fused_expr_to_c(right)
            )
        }
        FusedExpr::UnaryOp {
            op: FusedUnaryOp::Neg,
            operand,
        } => {
            format!("(-{})", fused_expr_to_c(operand))
        }
        FusedExpr::UnaryOp {
            op: FusedUnaryOp::Abs,
            operand,
        } => {
            let inner = fused_expr_to_c(operand);
            format!("({inner} < 0 ? -{inner} : {inner})")
        }
        FusedExpr::Scalar(v) => format!("{v}"),
    }
}

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

    #[test]
    fn test_fused_expr_to_c() {
        // a + b * c → (src0->data[i] + (src1->data[i] * src2->data[i]))
        let expr = FusedExpr::BinaryOp {
            op: FusedBinOp::Add,
            left: Box::new(FusedExpr::Source(0)),
            right: Box::new(FusedExpr::BinaryOp {
                op: FusedBinOp::Mul,
                left: Box::new(FusedExpr::Source(1)),
                right: Box::new(FusedExpr::Source(2)),
            }),
        };
        let c = fused_expr_to_c(&expr);
        assert_eq!(
            c,
            "(src0->data[i] + (src1->data[i] * src2->data[i]))"
        );
    }

    #[test]
    fn test_fused_bin_op_from_suffix() {
        assert_eq!(FusedBinOp::from_suffix("add"), Some(FusedBinOp::Add));
        assert_eq!(FusedBinOp::from_suffix("sub"), Some(FusedBinOp::Sub));
        assert_eq!(FusedBinOp::from_suffix("mul"), Some(FusedBinOp::Mul));
        assert_eq!(FusedBinOp::from_suffix("div"), Some(FusedBinOp::Div));
        assert_eq!(FusedBinOp::from_suffix("matmul"), None);
    }

    #[test]
    fn test_neg_expr() {
        let expr = FusedExpr::UnaryOp {
            op: FusedUnaryOp::Neg,
            operand: Box::new(FusedExpr::Source(0)),
        };
        assert_eq!(fused_expr_to_c(&expr), "(-src0->data[i])");
    }
}