joule_mir/dataflow/

optimize.rs

//! Dataflow Graph Optimizations
//!
//! This module provides optimization passes for dataflow graphs,
//! inspired by RipTide's 27-52% operator count reduction techniques.
//!
//! # Key Optimizations
//!
//! - **Stream Fusion**: Fuse loop induction variables into single Stream operators
//! - **Operator Merging**: Combine adjacent operators where beneficial
//! - **Dead Channel Elimination**: Remove unused channels
//! - **Constant Propagation**: Fold constants through the graph

use super::{
    ChannelId, ComputeOp, DataflowGraph, DependencyAnalysis, DfOperator, EnergyEstimate,
    OperatorId, TokenType, TokenValue,
};
use std::collections::{HashMap, HashSet};

/// Statistics from optimization passes
#[derive(Debug, Clone, Default)]
pub struct OptimizationStats {
    /// Number of operators before optimization
    pub operators_before: usize,
    /// Number of operators after optimization
    pub operators_after: usize,
    /// Number of channels before
    pub channels_before: usize,
    /// Number of channels after
    pub channels_after: usize,
    /// Streams created by fusion
    pub streams_created: usize,
    /// Operators eliminated
    pub operators_eliminated: usize,
}

impl OptimizationStats {
    /// Calculate percentage reduction in operators
    pub fn operator_reduction_percent(&self) -> f64 {
        if self.operators_before == 0 {
            return 0.0;
        }
        let eliminated = self.operators_before.saturating_sub(self.operators_after);
        (eliminated as f64 / self.operators_before as f64) * 100.0
    }
}

/// Stream fusion optimization
///
/// This is RipTide's key optimization that achieves 27-52% operator reduction
/// by fusing loop induction variable computations into single Stream operators.
///
/// Pattern detected:
/// ```text
/// Constant(start) --> Add --> Compare --> Steer
///                      ^          |
///                      |          v
///                   Carry <--- (loop body)
/// ```
///
/// Becomes:
/// ```text
/// Stream(start, step, bound) --> (loop body)
/// ```
pub struct StreamFusion {
    /// Minimum loop iterations to consider for fusion
    pub min_iterations: u64,
    /// Whether to fuse nested loops
    pub fuse_nested: bool,
}

impl StreamFusion {
    pub fn new() -> Self {
        Self {
            min_iterations: 4,
            fuse_nested: true,
        }
    }

    /// Run stream fusion on a dataflow graph
    pub fn optimize(&self, dfg: &mut DataflowGraph) -> usize {
        let mut streams_created = 0;

        // Find candidate patterns for stream fusion
        let candidates = self.find_induction_patterns(dfg);

        for candidate in candidates {
            if self.can_fuse(&candidate, dfg) {
                self.fuse_to_stream(candidate, dfg);
                streams_created += 1;
            }
        }

        streams_created
    }

    /// Find potential induction variable patterns
    fn find_induction_patterns(&self, dfg: &DataflowGraph) -> Vec<InductionPattern> {
        let mut patterns = Vec::new();

        // Look for Carry operators that might be induction variables
        for (op_idx, op) in dfg.operators.iter().enumerate() {
            if let DfOperator::Carry {
                initial,
                feedback,
                continue_signal,
                output,
            } = op
            {
                // Check if this looks like an induction variable:
                // - Initial value is a constant
                // - Feedback comes from an Add operation
                // - Continue signal comes from a comparison

                let initial_const = self.find_constant_source(*initial, dfg);
                let add_op = self.find_add_operator(*feedback, dfg);
                let cmp_info = self.find_comparison(*continue_signal, dfg);

                if let (Some(start), Some((add_idx, step_ch)), Some((cmp_idx, bound_ch))) =
                    (initial_const, add_op, cmp_info)
                {
                    patterns.push(InductionPattern {
                        carry_idx: op_idx,
                        initial_channel: *initial,
                        feedback_channel: *feedback,
                        continue_channel: *continue_signal,
                        output_channel: *output,
                        start_value: start,
                        step_channel: step_ch,
                        bound_channel: bound_ch,
                        add_operator: add_idx,
                        compare_operator: cmp_idx,
                    });
                }
            }
        }

        patterns
    }

    /// Find a constant source for a channel
    fn find_constant_source(&self, ch: ChannelId, dfg: &DataflowGraph) -> Option<i64> {
        for op in &dfg.operators {
            if let DfOperator::Constant { value, output, .. } = op {
                if *output == ch {
                    return match value {
                        TokenValue::Int(i) => Some(*i),
                        TokenValue::Uint(u) => Some(*u as i64),
                        _ => None,
                    };
                }
            }
        }
        None
    }

    /// Find an Add operator that produces the given channel
    fn find_add_operator(&self, ch: ChannelId, dfg: &DataflowGraph) -> Option<(usize, ChannelId)> {
        for (idx, op) in dfg.operators.iter().enumerate() {
            if let DfOperator::Compute {
                op: ComputeOp::Add,
                inputs,
                output,
            } = op
            {
                if *output == ch && inputs.len() == 2 {
                    // Return the step channel (the one that's not from Carry output)
                    // For simplicity, assume second input is the step
                    return Some((idx, inputs[1]));
                }
            }
        }
        None
    }

    /// Find a comparison operator that produces the given channel
    fn find_comparison(&self, ch: ChannelId, dfg: &DataflowGraph) -> Option<(usize, ChannelId)> {
        for (idx, op) in dfg.operators.iter().enumerate() {
            if let DfOperator::Compute {
                op: cmp_op,
                inputs,
                output,
            } = op
            {
                if *output == ch
                    && matches!(
                        cmp_op,
                        ComputeOp::Lt
                            | ComputeOp::Le
                            | ComputeOp::Gt
                            | ComputeOp::Ge
                            | ComputeOp::Ne
                    )
                {
                    if inputs.len() == 2 {
                        // Return the bound channel (typically the second input)
                        return Some((idx, inputs[1]));
                    }
                }
            }
        }
        None
    }

    /// Check if a pattern can be fused
    fn can_fuse(&self, pattern: &InductionPattern, dfg: &DataflowGraph) -> bool {
        // Check that the pattern is well-formed
        // - Step should be a constant
        // - Bound should be either constant or external input

        let step_const = self.find_constant_source(pattern.step_channel, dfg);

        // For now, only fuse if step is constant
        step_const.is_some()
    }

    /// Fuse an induction pattern into a Stream operator
    fn fuse_to_stream(&self, pattern: InductionPattern, dfg: &mut DataflowGraph) {
        // Create channels for Stream operator
        let start_ch = pattern.initial_channel;
        let step_ch = pattern.step_channel;
        let bound_ch = pattern.bound_channel;

        // Create done channel - the Stream operator will use this to signal loop completion,
        // replacing the old compare operator's continue_channel
        let done_ch = dfg.add_channel(TokenType::Bool);

        // Add Stream operator
        dfg.add_operator(DfOperator::Stream {
            start: start_ch,
            step: step_ch,
            bound: bound_ch,
            output: pattern.output_channel,
            done: done_ch,
        });

        // Pre-allocate dead-end channels for replaced operators to avoid borrow conflicts.
        // These channels receive the output of the replaced (now-dead) operators.
        let carry_dead_ch = dfg.add_channel(TokenType::i64());
        let add_dead_ch = dfg.add_channel(TokenType::i64());
        let cmp_dead_ch = dfg.add_channel(TokenType::Bool);

        // Replace the old Carry operator with a Constant producing the start value.
        // This effectively removes the Carry from the active graph since the Stream
        // now handles induction. The feedback_channel and continue_channel that fed
        // into the Carry are now disconnected.
        if let Some(op) = dfg.operators.get_mut(pattern.carry_idx) {
            let _ = pattern.feedback_channel;
            let _ = pattern.continue_channel;
            *op = DfOperator::Constant {
                value: TokenValue::Int(pattern.start_value),
                output: carry_dead_ch,
                repeat: None,
            };
        }

        // Replace the Add operator (which fed back into the Carry's feedback_channel)
        if let Some(op) = dfg.operators.get_mut(pattern.add_operator) {
            *op = DfOperator::Constant {
                value: TokenValue::Int(pattern.start_value),
                output: add_dead_ch,
                repeat: None,
            };
        }

        // Replace the Compare operator (which produced the continue_channel).
        // The Stream operator's done channel now replaces the compare operator's role.
        if let Some(op) = dfg.operators.get_mut(pattern.compare_operator) {
            *op = DfOperator::Constant {
                value: TokenValue::Uint(0),
                output: cmp_dead_ch,
                repeat: None,
            };
        }
    }
}

impl Default for StreamFusion {
    fn default() -> Self {
        Self::new()
    }
}

/// Pattern representing a loop induction variable
#[derive(Debug)]
struct InductionPattern {
    carry_idx: usize,
    initial_channel: ChannelId,
    feedback_channel: ChannelId,
    continue_channel: ChannelId,
    output_channel: ChannelId,
    start_value: i64,
    step_channel: ChannelId,
    bound_channel: ChannelId,
    add_operator: usize,
    compare_operator: usize,
}

/// Main optimizer that runs all optimization passes
pub struct DfgOptimizer {
    /// Stream fusion pass
    pub stream_fusion: StreamFusion,
    /// Run dead code elimination
    pub eliminate_dead_code: bool,
    /// Run constant propagation
    pub propagate_constants: bool,
}

impl DfgOptimizer {
    pub fn new() -> Self {
        Self {
            stream_fusion: StreamFusion::new(),
            eliminate_dead_code: true,
            propagate_constants: true,
        }
    }

    /// Run all optimization passes
    pub fn optimize(&self, dfg: &mut DataflowGraph) -> OptimizationStats {
        let operators_before = dfg.operators.len();
        let channels_before = dfg.channels.len();

        // Pass 1: Stream fusion
        let streams_created = self.stream_fusion.optimize(dfg);

        // Pass 2: Dead code elimination
        let dead_eliminated = if self.eliminate_dead_code {
            self.eliminate_dead_code(dfg)
        } else {
            0
        };

        // Pass 3: Constant propagation
        if self.propagate_constants {
            self.propagate_constants(dfg);
        }

        OptimizationStats {
            operators_before,
            operators_after: dfg.operators.len(),
            channels_before,
            channels_after: dfg.channels.len(),
            streams_created,
            operators_eliminated: dead_eliminated,
        }
    }

    /// Eliminate dead (unreachable) operators
    fn eliminate_dead_code(&self, dfg: &mut DataflowGraph) -> usize {
        let analysis = DependencyAnalysis::analyze(dfg);

        // Find operators that are not reachable from any output
        let mut reachable = HashSet::new();

        // Start from sinks (outputs)
        let mut worklist: Vec<_> = analysis.sinks.clone();

        while let Some(op_id) = worklist.pop() {
            if reachable.contains(&op_id) {
                continue;
            }
            reachable.insert(op_id);

            // Add predecessors to worklist
            if let Some(preds) = analysis.predecessors.get(&op_id) {
                for pred in preds {
                    worklist.push(*pred);
                }
            }
        }

        // Count unreachable operators (we don't actually remove them in this implementation)
        let total_ops = dfg.operators.len();
        let reachable_count = reachable.len();

        total_ops.saturating_sub(reachable_count)
    }

    /// Propagate constants through the graph
    fn propagate_constants(&self, dfg: &mut DataflowGraph) {
        // Build map of constant channels
        let mut constants: HashMap<ChannelId, TokenValue> = HashMap::new();

        for op in &dfg.operators {
            if let DfOperator::Constant { value, output, .. } = op {
                constants.insert(*output, value.clone());
            }
        }

        // Propagate through compute operators
        // (In a full implementation, we'd evaluate constant expressions)
    }

    /// Estimate energy for all operators
    pub fn estimate_energy(
        &self,
        dfg: &mut DataflowGraph,
        firing_estimates: &HashMap<OperatorId, u64>,
    ) {
        for (idx, op) in dfg.operators.iter().enumerate() {
            let op_id = OperatorId::new(idx as u32);
            let static_cost = op.energy_cost() as f64;
            let firing_count = firing_estimates.get(&op_id).copied().unwrap_or(1);

            dfg.energy_estimates.insert(
                op_id,
                EnergyEstimate {
                    static_cost,
                    dynamic_cost: 0.0, // Would be data-dependent
                    firing_count,
                },
            );
        }
    }
}

impl Default for DfgOptimizer {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_stream_fusion_new() {
        let fusion = StreamFusion::new();
        assert_eq!(fusion.min_iterations, 4);
        assert!(fusion.fuse_nested);
    }

    #[test]
    fn test_optimizer_default() {
        let optimizer = DfgOptimizer::new();
        assert!(optimizer.eliminate_dead_code);
        assert!(optimizer.propagate_constants);
    }

    #[test]
    fn test_optimization_stats() {
        let stats = OptimizationStats {
            operators_before: 100,
            operators_after: 75,
            channels_before: 50,
            channels_after: 40,
            streams_created: 5,
            operators_eliminated: 25,
        };

        assert!((stats.operator_reduction_percent() - 25.0).abs() < 0.01);
    }

    #[test]
    fn test_simple_optimization() {
        let mut dfg = DataflowGraph::new();

        // Create a simple graph
        let ch1 = dfg.add_channel(TokenType::i32());
        let _ch2 = dfg.add_channel(TokenType::i32());

        dfg.add_operator(DfOperator::Source {
            external_id: 0,
            output: ch1,
        });

        dfg.add_operator(DfOperator::Sink {
            input: ch1,
            external_id: 0,
        });

        let optimizer = DfgOptimizer::new();
        let stats = optimizer.optimize(&mut dfg);

        // Should not crash and produce valid stats
        assert!(stats.operators_before > 0);
    }
}