oxicuda_graph/optimizer/

fusion.rs

//! Operator fusion pass — merges chains of fusible element-wise kernels.
//!
//! # What is fused
//!
//! Two consecutive kernel nodes `a → b` are **fusion candidates** when:
//!
//! 1. Both are marked `fusible` in their `NodeKind::KernelLaunch`.
//! 2. `a` dominates `b` in the dominator tree (no side path can bypass `b`).
//! 3. There is no non-fusible node on any path from `a` to `b`.
//! 4. `a` has exactly one output buffer shared with `b`'s inputs (single
//!    producer–consumer chain — avoids creating broadcast copies).
//! 5. Both have the same launch configuration (same total thread count),
//!    or `b` uses a configuration that is a submultiple of `a`'s grid.
//!
//! When a group is fusible, the pass produces a `FusionGroup` describing
//! which original nodes should be merged. The graph itself is not modified
//! by this pass — the `Executor` is responsible for lowering the fusion
//! groups into combined PTX.
//!
//! # Algorithm
//!
//! 1. Run topological analysis to get ASAP order and level info.
//! 2. Run dominance analysis.
//! 3. Traverse nodes in topological order; greedily extend chains of
//!    fusible nodes that satisfy the rules above.
//! 4. Return the list of `FusionGroup`s.

use std::collections::{HashMap, HashSet};

use crate::analysis::{dominance_analyse, topo_analyse};
use crate::error::{GraphError, GraphResult};
use crate::graph::ComputeGraph;
use crate::node::{KernelConfig, NodeId, NodeKind};

// ---------------------------------------------------------------------------
// FusionGroup
// ---------------------------------------------------------------------------

/// A group of nodes that can be merged into a single fused kernel.
///
/// The nodes are listed in topological execution order. The fused kernel
/// will be generated by the PTX codegen layer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FusionGroup {
    /// Group identifier (sequential, 0-based).
    pub id: usize,
    /// Original node IDs in topological order (oldest first).
    pub members: Vec<NodeId>,
    /// Combined launch configuration (uses the first member's config).
    pub config: KernelConfig,
    /// Human-readable tag for debugging.
    pub tag: String,
}

impl FusionGroup {
    /// Returns the number of nodes in this group.
    #[must_use]
    pub fn size(&self) -> usize {
        self.members.len()
    }

    /// Returns `true` if this group contains just one node (trivial — no fusion).
    #[must_use]
    pub fn is_trivial(&self) -> bool {
        self.members.len() == 1
    }
}

// ---------------------------------------------------------------------------
// FusionPlan
// ---------------------------------------------------------------------------

/// The complete fusion plan produced by the fusion pass.
#[derive(Debug, Clone)]
pub struct FusionPlan {
    /// All fusion groups (including trivial single-node groups).
    pub groups: Vec<FusionGroup>,
    /// Map from NodeId to the FusionGroup index it belongs to.
    pub node_to_group: HashMap<NodeId, usize>,
}

impl FusionPlan {
    /// Returns the number of non-trivial fusion groups (size ≥ 2).
    pub fn fusion_count(&self) -> usize {
        self.groups.iter().filter(|g| !g.is_trivial()).count()
    }

    /// Returns the total number of nodes saved by fusion.
    ///
    /// Each non-trivial group of size `k` saves `k-1` kernel launches.
    pub fn nodes_saved(&self) -> usize {
        self.groups
            .iter()
            .filter(|g| !g.is_trivial())
            .map(|g| g.size() - 1)
            .sum()
    }

    /// Returns the fusion group that owns `node`.
    pub fn group_of(&self, node: NodeId) -> Option<&FusionGroup> {
        self.node_to_group
            .get(&node)
            .and_then(|&idx| self.groups.get(idx))
    }
}

// ---------------------------------------------------------------------------
// Fusion eligibility checks
// ---------------------------------------------------------------------------

/// Returns `true` if two kernel configs are compatible for fusion.
///
/// Two configs are compatible when they have the same total number of threads
/// (same grid×block volume), so the fused kernel can use an identical launch.
fn configs_compatible(a: &KernelConfig, b: &KernelConfig) -> bool {
    a.total_threads() == b.total_threads()
}

/// Returns `true` if node `a` is immediately before `b` in the topological
/// order AND there are no intervening non-fusible nodes on the single path.
fn only_fusible_between(
    graph: &ComputeGraph,
    a: NodeId,
    b: NodeId,
    topo_pos: &HashMap<NodeId, usize>,
) -> bool {
    let pos_a = topo_pos[&a];
    let pos_b = topo_pos[&b];
    if pos_b <= pos_a + 1 {
        return true; // adjacent
    }
    // BFS from a; if we can reach b without going through a non-fusible
    // compute node, the path is clean.
    let mut visited = HashSet::new();
    let mut stack = vec![a];
    while let Some(cur) = stack.pop() {
        if cur == b {
            continue;
        }
        for &s in graph.successors(cur).unwrap_or(&[]) {
            if visited.insert(s) {
                if s == b {
                    continue;
                }
                let node = graph.node(s).ok();
                let is_fusible = node.map(|n| n.kind.is_fusible()).unwrap_or(false);
                let is_barrier = node
                    .map(|n| matches!(n.kind, NodeKind::Barrier))
                    .unwrap_or(false);
                let spos = topo_pos.get(&s).copied().unwrap_or(usize::MAX);
                if spos < pos_b && (is_fusible || is_barrier) {
                    stack.push(s);
                } else if spos < pos_b && !is_fusible && !is_barrier {
                    return false; // non-fusible node between a and b
                }
            }
        }
    }
    true
}

// ---------------------------------------------------------------------------
// analyse — entry point
// ---------------------------------------------------------------------------

/// Runs the fusion analysis pass on `graph`.
///
/// Returns a [`FusionPlan`] describing which nodes can be merged.
///
/// # Errors
///
/// Returns [`GraphError::EmptyGraph`] if the graph has no nodes.
pub fn analyse(graph: &ComputeGraph) -> GraphResult<FusionPlan> {
    if graph.is_empty() {
        return Err(GraphError::EmptyGraph);
    }

    let topo = topo_analyse(graph)?;
    let dt = dominance_analyse(graph)?;

    let topo_pos: HashMap<NodeId, usize> = topo
        .order
        .iter()
        .enumerate()
        .map(|(p, &id)| (id, p))
        .collect();

    // Assigned group index for each node (None = unassigned).
    let mut assigned: HashMap<NodeId, usize> = HashMap::new();
    let mut groups: Vec<FusionGroup> = Vec::new();

    // Traverse in topological order.
    for &node_id in &topo.order {
        if assigned.contains_key(&node_id) {
            continue;
        }

        let node = graph.node(node_id)?;

        // Only kernel nodes participate in fusion.
        let (is_fusible, base_config) = match &node.kind {
            NodeKind::KernelLaunch {
                fusible, config, ..
            } => (*fusible, *config),
            _ => {
                // Non-kernel node: assign to a trivial group.
                let gid = groups.len();
                groups.push(FusionGroup {
                    id: gid,
                    members: vec![node_id],
                    config: KernelConfig::linear(1, 1, 0),
                    tag: format!("non_kernel_{}", node.kind.tag()),
                });
                assigned.insert(node_id, gid);
                continue;
            }
        };

        if !is_fusible {
            let gid = groups.len();
            groups.push(FusionGroup {
                id: gid,
                members: vec![node_id],
                config: base_config,
                tag: format!("non_fusible_{}", node.display_name()),
            });
            assigned.insert(node_id, gid);
            continue;
        }

        // Start a new fusion group with this node.
        let gid = groups.len();
        let mut members = vec![node_id];
        assigned.insert(node_id, gid);

        // Greedily extend: try to add direct successors that are fusible.
        let mut frontier = graph.successors(node_id)?.to_vec();
        while let Some(succ_id) = frontier.first().copied() {
            frontier.remove(0);
            if assigned.contains_key(&succ_id) {
                continue;
            }
            let succ = graph.node(succ_id)?;
            let (succ_fusible, succ_config) = match &succ.kind {
                NodeKind::KernelLaunch {
                    fusible, config, ..
                } => (*fusible, *config),
                _ => continue,
            };
            if !succ_fusible {
                continue;
            }
            // Check fusion rules:
            // 1. Compatible launch config.
            if !configs_compatible(&base_config, &succ_config) {
                continue;
            }
            // 2. Dominator: the last member must dominate succ.
            let last_member = *members.last().unwrap();
            if !dt.dominates(last_member, succ_id) {
                continue;
            }
            // 3. No non-fusible nodes between them.
            if !only_fusible_between(graph, last_member, succ_id, &topo_pos) {
                continue;
            }
            // Accept this node into the group.
            members.push(succ_id);
            assigned.insert(succ_id, gid);
            // Continue extending from succ.
            for &next in graph.successors(succ_id)? {
                if !assigned.contains_key(&next) {
                    frontier.push(next);
                }
            }
        }

        let tag = if members.len() > 1 {
            format!(
                "fused_{}..{}",
                graph.node(members[0])?.display_name(),
                graph.node(*members.last().unwrap())?.display_name()
            )
        } else {
            format!("solo_{}", graph.node(node_id)?.display_name())
        };

        groups.push(FusionGroup {
            id: gid,
            members,
            config: base_config,
            tag,
        });
    }

    // Build node_to_group map.
    let node_to_group: HashMap<NodeId, usize> = assigned;

    Ok(FusionPlan {
        groups,
        node_to_group,
    })
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::builder::GraphBuilder;
    use crate::node::MemcpyDir;

    fn fusible_kernel(b: &mut GraphBuilder, name: &str) -> NodeId {
        b.add_kernel(name, 4, 256, 0).fusible(true).finish()
    }

    fn non_fusible_kernel(b: &mut GraphBuilder, name: &str) -> NodeId {
        b.add_kernel(name, 4, 256, 0).fusible(false).finish()
    }

    #[test]
    fn fusion_empty_graph() {
        let g = ComputeGraph::new();
        assert!(matches!(analyse(&g), Err(GraphError::EmptyGraph)));
    }

    #[test]
    fn fusion_single_fusible_kernel_trivial_group() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k = fusible_kernel(&mut b, "add");
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        assert_eq!(plan.groups.len(), 1);
        assert!(plan.groups[0].is_trivial());
        assert_eq!(plan.group_of(k).unwrap().members, vec![k]);
    }

    #[test]
    fn fusion_chain_of_fusible_kernels_merged() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = fusible_kernel(&mut b, "k0");
        let k1 = fusible_kernel(&mut b, "k1");
        let k2 = fusible_kernel(&mut b, "k2");
        b.chain(&[k0, k1, k2]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // All three are fusible and in a chain → one non-trivial group.
        assert_eq!(plan.fusion_count(), 1);
        let group = plan.group_of(k0).unwrap();
        assert_eq!(group.size(), 3);
        assert!(group.members.contains(&k0));
        assert!(group.members.contains(&k1));
        assert!(group.members.contains(&k2));
    }

    #[test]
    fn fusion_non_fusible_breaks_chain() {
        // k0 (fusible) → k1 (non-fusible) → k2 (fusible)
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = fusible_kernel(&mut b, "k0");
        let k1 = non_fusible_kernel(&mut b, "k1");
        let k2 = fusible_kernel(&mut b, "k2");
        b.chain(&[k0, k1, k2]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // k0 and k2 should be in different groups.
        let g0 = plan.group_of(k0).unwrap().id;
        let g2 = plan.group_of(k2).unwrap().id;
        assert_ne!(g0, g2);
    }

    #[test]
    fn fusion_memcpy_not_fused() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let upload = b.add_memcpy("up", MemcpyDir::HostToDevice, 1024);
        let k = fusible_kernel(&mut b, "k");
        let download = b.add_memcpy("dn", MemcpyDir::DeviceToHost, 1024);
        b.chain(&[upload, k, download]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // Upload and download are non-kernel nodes → trivial groups.
        let gup = plan.group_of(upload).unwrap();
        let gdn = plan.group_of(download).unwrap();
        assert!(gup.is_trivial());
        assert!(gdn.is_trivial());
    }

    #[test]
    fn fusion_incompatible_configs_not_fused() {
        // k0: 4 blocks × 256 threads = 1024 threads
        // k1: 8 blocks × 256 threads = 2048 threads (incompatible)
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = b.add_kernel("k0", 4, 256, 0).fusible(true).finish();
        let k1 = b.add_kernel("k1", 8, 256, 0).fusible(true).finish();
        b.dep(k0, k1);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        let gk0 = plan.group_of(k0).unwrap().id;
        let gk1 = plan.group_of(k1).unwrap().id;
        assert_ne!(gk0, gk1);
    }

    #[test]
    fn fusion_nodes_saved_count() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = fusible_kernel(&mut b, "k0");
        let k1 = fusible_kernel(&mut b, "k1");
        let k2 = fusible_kernel(&mut b, "k2");
        b.chain(&[k0, k1, k2]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // Group of 3 saves 2 kernel launches.
        assert_eq!(plan.nodes_saved(), 2);
    }

    #[test]
    fn fusion_plan_covers_all_nodes() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = fusible_kernel(&mut b, "k0");
        let k1 = non_fusible_kernel(&mut b, "k1");
        let upload = b.add_memcpy("up", MemcpyDir::HostToDevice, 512);
        b.chain(&[upload, k0, k1]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // Every node must appear in exactly one group.
        let total: usize = plan.groups.iter().map(|g| g.size()).sum();
        assert_eq!(total, 3);
        // Every node must be in node_to_group.
        assert!(plan.node_to_group.contains_key(&k0));
        assert!(plan.node_to_group.contains_key(&k1));
        assert!(plan.node_to_group.contains_key(&upload));
    }

    #[test]
    fn fusion_parallel_branches_not_fused() {
        // src → k0 and src → k1 (independent branches, cannot fuse across them).
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let src = b.add_barrier("src");
        let k0 = fusible_kernel(&mut b, "k0");
        let k1 = fusible_kernel(&mut b, "k1");
        b.fan_out(src, &[k0, k1]);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        // k0 and k1 are in separate groups (not dominator-related to each other).
        let gk0 = plan.group_of(k0).unwrap().id;
        let gk1 = plan.group_of(k1).unwrap().id;
        assert_ne!(gk0, gk1);
    }

    #[test]
    fn fusion_group_tag_contains_names() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k0 = fusible_kernel(&mut b, "relu");
        let k1 = fusible_kernel(&mut b, "scale");
        b.dep(k0, k1);
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        let group = plan.group_of(k0).unwrap();
        assert!(!group.tag.is_empty());
    }

    #[test]
    fn fusion_empty_fusible_graph_one_group() {
        let mut b = GraphBuilder::new().with_auto_infer_edges(false);
        let k = fusible_kernel(&mut b, "solo");
        let g = b.build().unwrap();
        let plan = analyse(&g).unwrap();
        assert_eq!(plan.fusion_count(), 0); // trivial, not fused
        assert_eq!(plan.nodes_saved(), 0);
        assert_eq!(plan.group_of(k).unwrap().size(), 1);
    }
}