joule_mir/dataflow/

schedule.rs

//! Energy-Aware Dataflow Scheduler
//!
//! This module provides scheduling for dataflow graphs that considers:
//! - Data dependencies (traditional dataflow)
//! - Thermal state (reduce parallelism when hot)
//! - Energy budgets (constrain total energy consumption)
//! - Available parallelism (match to hardware)
//!
//! # Scheduling Strategy
//!
//! 1. Build dependency graph from channels
//! 2. Topological sort to determine execution order
//! 3. Group operators into parallel execution steps
//! 4. Apply energy/thermal constraints to limit parallelism
//!
//! # Example
//!
//! ```ignore
//! let scheduler = EnergyAwareScheduler::new(SchedulingConfig::default());
//! let scheduled = scheduler.schedule(&dfg, ThermalState::Nominal);
//! ```

use super::{DataflowGraph, DependencyAnalysis, OperatorId, ScheduleStep, ScheduledDfg};
use std::cmp::Ordering;
use std::collections::{HashMap, VecDeque};

/// Thermal state for scheduling decisions
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ThermalConstraint {
    /// Cool - can use maximum parallelism
    Cool,
    /// Nominal - use standard parallelism
    Nominal,
    /// Elevated - reduce parallelism to 75%
    Elevated,
    /// Hot - reduce parallelism to 50%
    Hot,
    /// Critical - serialize execution
    Critical,
}

impl ThermalConstraint {
    /// Get parallelism multiplier for this thermal state
    pub fn parallelism_factor(&self) -> f64 {
        match self {
            ThermalConstraint::Cool => 1.0,
            ThermalConstraint::Nominal => 1.0,
            ThermalConstraint::Elevated => 0.75,
            ThermalConstraint::Hot => 0.5,
            ThermalConstraint::Critical => 0.0, // Forces serialization
        }
    }

    /// Get energy multiplier (hot execution uses more energy)
    pub fn energy_multiplier(&self) -> f64 {
        match self {
            ThermalConstraint::Cool => 0.9,
            ThermalConstraint::Nominal => 1.0,
            ThermalConstraint::Elevated => 1.3,
            ThermalConstraint::Hot => 1.6,
            ThermalConstraint::Critical => 2.0,
        }
    }
}

/// Configuration for the scheduler
#[derive(Debug, Clone)]
pub struct SchedulingConfig {
    /// Maximum parallelism (number of operators in one step)
    pub max_parallelism: usize,
    /// Energy budget in Joules (None = unlimited)
    pub energy_budget_j: Option<f64>,
    /// Whether to prefer energy efficiency over speed
    pub prefer_efficiency: bool,
    /// Minimum operators to run in parallel (below this, don't bother)
    pub min_parallelism: usize,
    /// Energy cost per operator (baseline, in arbitrary units)
    pub base_energy_per_op: f64,
}

impl Default for SchedulingConfig {
    fn default() -> Self {
        Self {
            max_parallelism: 8, // Typical multicore CPU
            energy_budget_j: None,
            prefer_efficiency: false,
            min_parallelism: 1,
            base_energy_per_op: 1e-6, // 1 microjoule per op baseline
        }
    }
}

impl SchedulingConfig {
    /// Create config for high-performance mode
    pub fn high_performance(cores: usize) -> Self {
        Self {
            max_parallelism: cores,
            energy_budget_j: None,
            prefer_efficiency: false,
            min_parallelism: 2,
            base_energy_per_op: 1e-6,
        }
    }

    /// Create config for energy-efficient mode
    pub fn energy_efficient(budget_j: f64) -> Self {
        Self {
            max_parallelism: 4,
            energy_budget_j: Some(budget_j),
            prefer_efficiency: true,
            min_parallelism: 1,
            base_energy_per_op: 1e-6,
        }
    }

    /// Create config for thermal-constrained mode
    pub fn thermal_constrained() -> Self {
        Self {
            max_parallelism: 2,
            energy_budget_j: None,
            prefer_efficiency: true,
            min_parallelism: 1,
            base_energy_per_op: 1e-6,
        }
    }
}

/// Energy-aware dataflow scheduler
pub struct EnergyAwareScheduler {
    config: SchedulingConfig,
}

impl EnergyAwareScheduler {
    pub fn new(config: SchedulingConfig) -> Self {
        Self { config }
    }

    /// Schedule a dataflow graph considering thermal constraints
    pub fn schedule(&self, dfg: &DataflowGraph, thermal: ThermalConstraint) -> ScheduledDfg {
        // Analyze dependencies
        let analysis = DependencyAnalysis::analyze(dfg);

        // Calculate effective parallelism
        let effective_parallelism = self.calculate_effective_parallelism(thermal);

        // Build schedule using level-based approach
        let schedule = self.build_schedule(dfg, &analysis, effective_parallelism);

        // Calculate energy estimate
        let estimated_energy = self.estimate_total_energy(dfg, &schedule, thermal);

        // Estimate duration (simplified)
        let estimated_duration = self.estimate_duration(&schedule);

        ScheduledDfg {
            dfg: dfg.clone(),
            schedule,
            estimated_energy_j: estimated_energy,
            estimated_duration_s: estimated_duration,
        }
    }

    /// Calculate effective parallelism given config and thermal state
    fn calculate_effective_parallelism(&self, thermal: ThermalConstraint) -> usize {
        let base = self.config.max_parallelism as f64;
        let thermal_factor = thermal.parallelism_factor();

        let effective = (base * thermal_factor).round() as usize;

        // Apply energy budget constraint if set
        let energy_constrained = if let Some(budget) = self.config.energy_budget_j {
            // Rough heuristic: limit parallelism based on budget
            let ops_allowed = (budget / self.config.base_energy_per_op) as usize;
            effective.min(ops_allowed.max(1))
        } else {
            effective
        };

        energy_constrained.max(self.config.min_parallelism)
    }

    /// Build the schedule using topological levels
    fn build_schedule(
        &self,
        dfg: &DataflowGraph,
        analysis: &DependencyAnalysis,
        max_parallelism: usize,
    ) -> Vec<ScheduleStep> {
        if dfg.operators.is_empty() {
            return Vec::new();
        }

        // Assign levels using topological sort
        let levels = self.assign_levels(dfg, analysis);

        // Find maximum level
        let max_level = levels.values().max().copied().unwrap_or(0);

        // Group operators by level
        let mut level_groups: Vec<Vec<OperatorId>> = vec![Vec::new(); max_level + 1];
        for (&op_id, &level) in &levels {
            level_groups[level].push(op_id);
        }

        // Create schedule steps, respecting parallelism limit
        let mut schedule = Vec::new();

        for level_ops in level_groups {
            if level_ops.is_empty() {
                continue;
            }

            // Build ScheduleEntry list for priority-based sorting within each level.
            // Each entry captures the operator's energy cost for scheduling decisions.
            let mut entries: Vec<ScheduleEntry> = level_ops
                .iter()
                .map(|&op_id| {
                    let energy = dfg
                        .operators
                        .get(op_id.0 as usize)
                        .map(|op| op.energy_cost())
                        .unwrap_or(0);
                    ScheduleEntry {
                        op_id,
                        priority: if self.config.prefer_efficiency {
                            // Invert energy cost so low-energy ops sort first (higher priority)
                            i32::MAX - energy as i32
                        } else {
                            0 // Equal priority when not preferring efficiency
                        },
                        energy_cost: energy,
                    }
                })
                .collect();

            // Sort by ScheduleEntry ordering (higher priority first, then lower energy)
            entries.sort_by(|a, b| a.cmp(b).reverse());

            let sorted_ops: Vec<OperatorId> = entries.iter().map(|e| e.op_id).collect();

            // Split into chunks of max_parallelism
            for chunk in sorted_ops.chunks(max_parallelism) {
                let step_energy: f64 = chunk
                    .iter()
                    .map(|op_id| {
                        dfg.operators
                            .get(op_id.0 as usize)
                            .map(|op| op.energy_cost() as f64 * self.config.base_energy_per_op)
                            .unwrap_or(0.0)
                    })
                    .sum();

                schedule.push(ScheduleStep {
                    operators: chunk.to_vec(),
                    max_parallelism: chunk.len(),
                    estimated_energy: step_energy,
                });
            }
        }

        schedule
    }

    /// Assign topological levels to operators
    fn assign_levels(
        &self,
        dfg: &DataflowGraph,
        analysis: &DependencyAnalysis,
    ) -> HashMap<OperatorId, usize> {
        let mut levels: HashMap<OperatorId, usize> = HashMap::new();
        let mut queue: VecDeque<OperatorId> = VecDeque::new();

        // Initialize sources at level 0
        for &src in &analysis.sources {
            levels.insert(src, 0);
            queue.push_back(src);
        }

        // BFS to assign levels
        while let Some(op) = queue.pop_front() {
            let op_level = *levels.get(&op).unwrap_or(&0);

            if let Some(succs) = analysis.successors.get(&op) {
                for &succ in succs {
                    let new_level = op_level + 1;
                    let current = levels.entry(succ).or_insert(0);
                    if new_level > *current {
                        *current = new_level;
                    }

                    // Add to queue if all predecessors have been processed
                    let all_preds_done = analysis
                        .predecessors
                        .get(&succ)
                        .map(|preds| preds.iter().all(|p| levels.contains_key(p)))
                        .unwrap_or(true);

                    if all_preds_done && !queue.contains(&succ) {
                        queue.push_back(succ);
                    }
                }
            }
        }

        // Handle any unvisited operators (disconnected from sources)
        for i in 0..dfg.operators.len() {
            let op_id = OperatorId::new(i as u32);
            levels.entry(op_id).or_insert(0);
        }

        levels
    }

    /// Estimate total energy for the schedule
    fn estimate_total_energy(
        &self,
        _dfg: &DataflowGraph,
        schedule: &[ScheduleStep],
        thermal: ThermalConstraint,
    ) -> f64 {
        let base_energy: f64 = schedule.iter().map(|s| s.estimated_energy).sum();

        // Apply thermal multiplier
        base_energy * thermal.energy_multiplier()
    }

    /// Estimate execution duration
    fn estimate_duration(&self, schedule: &[ScheduleStep]) -> f64 {
        // Simplified: assume each step takes 1 microsecond base time
        // divided by parallelism
        let base_time_per_op = 1e-6; // 1 microsecond

        schedule
            .iter()
            .map(|step| {
                let ops = step.operators.len().max(1);
                let parallelism = step.max_parallelism.max(1);
                (ops as f64 / parallelism as f64) * base_time_per_op
            })
            .sum()
    }

    /// Check if energy budget is exceeded
    pub fn check_energy_budget(&self, scheduled: &ScheduledDfg) -> Option<f64> {
        if let Some(budget) = self.config.energy_budget_j {
            if scheduled.estimated_energy_j > budget {
                return Some(scheduled.estimated_energy_j - budget);
            }
        }
        None
    }

    /// Reschedule with reduced parallelism to meet energy budget
    pub fn reschedule_for_energy(&self, dfg: &DataflowGraph, target_energy_j: f64) -> ScheduledDfg {
        // Start with critical thermal state (minimum parallelism)
        let thermal = ThermalConstraint::Critical;
        let mut scheduled = self.schedule(dfg, thermal);

        // If still over budget, there's not much more we can do
        if scheduled.estimated_energy_j > target_energy_j {
            // Return the minimum-parallelism schedule anyway
            return scheduled;
        }

        // Try increasing parallelism while staying under budget
        for state in &[
            ThermalConstraint::Hot,
            ThermalConstraint::Elevated,
            ThermalConstraint::Nominal,
            ThermalConstraint::Cool,
        ] {
            let candidate = self.schedule(dfg, *state);
            if candidate.estimated_energy_j <= target_energy_j {
                scheduled = candidate;
            } else {
                break;
            }
        }

        scheduled
    }
}

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

/// Priority queue entry for scheduling
#[derive(Debug, Clone, Eq, PartialEq)]
struct ScheduleEntry {
    op_id: OperatorId,
    priority: i32, // Higher = schedule sooner
    energy_cost: u32,
}

impl Ord for ScheduleEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        // Higher priority first, then lower energy cost
        self.priority
            .cmp(&other.priority)
            .then_with(|| other.energy_cost.cmp(&self.energy_cost))
    }
}

impl PartialOrd for ScheduleEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dataflow::{ComputeOp, DfOperator, TokenType};

    #[test]
    #[allow(clippy::float_cmp)]
    fn test_thermal_constraint_factors() {
        assert_eq!(ThermalConstraint::Cool.parallelism_factor(), 1.0);
        assert_eq!(ThermalConstraint::Hot.parallelism_factor(), 0.5);
        assert_eq!(ThermalConstraint::Critical.parallelism_factor(), 0.0);

        assert!(ThermalConstraint::Hot.energy_multiplier() > 1.0);
        assert!(ThermalConstraint::Cool.energy_multiplier() < 1.0);
    }

    #[test]
    fn test_scheduling_config() {
        let config = SchedulingConfig::default();
        assert_eq!(config.max_parallelism, 8);
        assert!(config.energy_budget_j.is_none());

        let efficient = SchedulingConfig::energy_efficient(1.0);
        assert_eq!(efficient.energy_budget_j, Some(1.0));
        assert!(efficient.prefer_efficiency);

        let perf = SchedulingConfig::high_performance(16);
        assert_eq!(perf.max_parallelism, 16);
    }

    #[test]
    fn test_scheduler_empty_graph() {
        let dfg = DataflowGraph::new();
        let scheduler = EnergyAwareScheduler::default();
        let scheduled = scheduler.schedule(&dfg, ThermalConstraint::Nominal);

        assert!(scheduled.schedule.is_empty());
    }

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

        // A -> B -> C
        let ch1 = dfg.add_channel(TokenType::i32());
        let ch2 = dfg.add_channel(TokenType::i32());
        let _ch3 = dfg.add_channel(TokenType::i32());

        dfg.add_operator(DfOperator::Source {
            external_id: 0,
            output: ch1,
        });
        dfg.add_operator(DfOperator::Compute {
            op: ComputeOp::Add,
            inputs: vec![ch1],
            output: ch2,
        });
        dfg.add_operator(DfOperator::Sink {
            input: ch2,
            external_id: 0,
        });

        let scheduler = EnergyAwareScheduler::default();
        let scheduled = scheduler.schedule(&dfg, ThermalConstraint::Nominal);

        // Linear graph should have 3 steps (one per operator)
        assert_eq!(scheduled.schedule.len(), 3);
    }

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

        // Source splits to two parallel computes, then merges
        let ch_in = dfg.add_channel(TokenType::i32());
        let ch_a = dfg.add_channel(TokenType::i32());
        let ch_b = dfg.add_channel(TokenType::i32());
        let ch_mul_out = dfg.add_channel(TokenType::i32());
        let ch_add_out = dfg.add_channel(TokenType::i32());

        dfg.add_operator(DfOperator::Source {
            external_id: 0,
            output: ch_in,
        });
        dfg.add_operator(DfOperator::Split {
            input: ch_in,
            outputs: vec![ch_a, ch_b],
        });
        dfg.add_operator(DfOperator::Compute {
            op: ComputeOp::Mul,
            inputs: vec![ch_a],
            output: ch_mul_out,
        });
        dfg.add_operator(DfOperator::Compute {
            op: ComputeOp::Add,
            inputs: vec![ch_b],
            output: ch_add_out,
        });

        let scheduler = EnergyAwareScheduler::new(SchedulingConfig {
            max_parallelism: 4,
            ..Default::default()
        });
        let scheduled = scheduler.schedule(&dfg, ThermalConstraint::Nominal);

        // Should have fewer steps than operators due to parallelism
        assert!(scheduled.schedule.len() <= dfg.operators.len());
    }

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

        // Create several independent operators
        for i in 0..10 {
            let ch = dfg.add_channel(TokenType::i32());
            dfg.add_operator(DfOperator::Source {
                external_id: i,
                output: ch,
            });
        }

        let scheduler = EnergyAwareScheduler::new(SchedulingConfig {
            max_parallelism: 8,
            ..Default::default()
        });

        let cool = scheduler.schedule(&dfg, ThermalConstraint::Cool);
        let hot = scheduler.schedule(&dfg, ThermalConstraint::Hot);
        let critical = scheduler.schedule(&dfg, ThermalConstraint::Critical);

        // Hotter thermal states should have more steps (less parallelism)
        assert!(hot.schedule.len() >= cool.schedule.len());
        assert!(critical.schedule.len() >= hot.schedule.len());
    }

    #[test]
    fn test_energy_budget_check() {
        let mut dfg = DataflowGraph::new();
        let ch = dfg.add_channel(TokenType::i32());
        dfg.add_operator(DfOperator::Source {
            external_id: 0,
            output: ch,
        });

        let scheduler = EnergyAwareScheduler::new(SchedulingConfig {
            energy_budget_j: Some(1e-9), // Very small budget
            ..Default::default()
        });

        let scheduled = scheduler.schedule(&dfg, ThermalConstraint::Nominal);

        // Should detect budget exceeded
        let _exceeded = scheduler.check_energy_budget(&scheduled);
        // May or may not exceed depending on estimates
    }
}