309 lines
8.7 KiB
Rust
309 lines
8.7 KiB
Rust
//! Optimization passes for GPU kernels
|
|
|
|
use crate::ir::{IRGraph, NodeId, OperationType};
|
|
use anyhow::Result;
|
|
use tracing::debug;
|
|
|
|
/// Optimization level
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub enum OptLevel {
|
|
/// No optimizations
|
|
None,
|
|
/// Basic optimizations
|
|
Basic,
|
|
/// Aggressive optimizations
|
|
Aggressive,
|
|
}
|
|
|
|
/// Statistics from optimization passes
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct OptimizationStats {
|
|
pub passes_applied: u32,
|
|
pub nodes_eliminated: u32,
|
|
pub operations_fused: u32,
|
|
pub memory_saved: u64,
|
|
}
|
|
|
|
/// Trait for optimization passes
|
|
pub trait OptimizationPass {
|
|
/// Apply the optimization pass to the graph
|
|
fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats>;
|
|
|
|
/// Get the name of this optimization pass
|
|
fn name(&self) -> &str;
|
|
}
|
|
|
|
/// Dead code elimination pass
|
|
pub struct DeadCodeEliminationPass;
|
|
|
|
impl OptimizationPass for DeadCodeEliminationPass {
|
|
fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
|
|
let mut stats = OptimizationStats::default();
|
|
let mut dead_nodes = Vec::new();
|
|
|
|
// Find nodes with no outputs and not in graph outputs
|
|
for (node_id, node) in &graph.nodes {
|
|
if node.outputs.is_empty() && !graph.outputs.contains(node_id) {
|
|
// Check if this node has side effects
|
|
match node.operation {
|
|
OperationType::Store | OperationType::Call => {
|
|
// These have side effects, keep them
|
|
}
|
|
_ => {
|
|
dead_nodes.push(*node_id);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Remove dead nodes
|
|
for node_id in dead_nodes {
|
|
graph.remove_node(node_id);
|
|
stats.nodes_eliminated += 1;
|
|
}
|
|
|
|
stats.passes_applied = 1;
|
|
Ok(stats)
|
|
}
|
|
|
|
fn name(&self) -> &'static str {
|
|
"DeadCodeElimination"
|
|
}
|
|
}
|
|
|
|
/// Constant folding pass
|
|
pub struct ConstantFoldingPass;
|
|
|
|
impl OptimizationPass for ConstantFoldingPass {
|
|
fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
|
|
let mut stats = OptimizationStats::default();
|
|
|
|
// Find constant operations that can be folded
|
|
let mut to_replace = Vec::new();
|
|
|
|
for (node_id, node) in &graph.nodes {
|
|
match node.operation {
|
|
OperationType::Add
|
|
| OperationType::Sub
|
|
| OperationType::Mul
|
|
| OperationType::Div => {
|
|
// Check if both inputs are constants
|
|
let input_nodes: Vec<_> = node
|
|
.inputs
|
|
.iter()
|
|
.filter_map(|id| graph.get_node(*id))
|
|
.collect();
|
|
|
|
if input_nodes.len() == 2
|
|
&& input_nodes
|
|
.iter()
|
|
.all(|n| matches!(n.operation, OperationType::Constant))
|
|
{
|
|
// This could be constant folded
|
|
to_replace.push(*node_id);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
// Replace with constant nodes
|
|
for node_id in to_replace {
|
|
if let Some(node) = graph.get_node_mut(node_id) {
|
|
node.operation = OperationType::Constant;
|
|
node.inputs.clear();
|
|
stats.nodes_eliminated += 1;
|
|
}
|
|
}
|
|
|
|
stats.passes_applied = 1;
|
|
Ok(stats)
|
|
}
|
|
|
|
fn name(&self) -> &'static str {
|
|
"ConstantFolding"
|
|
}
|
|
}
|
|
|
|
/// Operation fusion pass
|
|
pub struct OperationFusionPass;
|
|
|
|
impl OptimizationPass for OperationFusionPass {
|
|
fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
|
|
let mut stats = OptimizationStats::default();
|
|
|
|
// Look for fusable patterns
|
|
let mut fusion_groups = Vec::new();
|
|
|
|
for (node_id, node) in &graph.nodes {
|
|
// Look for simple fusion patterns
|
|
if matches!(node.operation, OperationType::Add | OperationType::Mul)
|
|
&& let Some(next_nodes) = self.find_fusable_sequence(*node_id, graph)
|
|
{
|
|
fusion_groups.push(next_nodes);
|
|
}
|
|
}
|
|
|
|
// Apply fusions
|
|
for group in fusion_groups {
|
|
if group.len() > 1 {
|
|
self.fuse_operations(graph, &group);
|
|
stats.operations_fused += (group.len() - 1) as u32;
|
|
}
|
|
}
|
|
|
|
stats.passes_applied = 1;
|
|
Ok(stats)
|
|
}
|
|
|
|
fn name(&self) -> &'static str {
|
|
"OperationFusion"
|
|
}
|
|
}
|
|
|
|
impl OperationFusionPass {
|
|
fn find_fusable_sequence(&self, start_node: NodeId, graph: &IRGraph) -> Option<Vec<NodeId>> {
|
|
let mut sequence = vec![start_node];
|
|
let mut current = start_node;
|
|
|
|
// Simple linear fusion for now
|
|
while let Some(node) = graph.get_node(current) {
|
|
if node.outputs.len() == 1 {
|
|
let next_id = node.outputs[0];
|
|
if let Some(next_node) = graph.get_node(next_id) {
|
|
if matches!(next_node.operation, OperationType::Add | OperationType::Mul) {
|
|
sequence.push(next_id);
|
|
current = next_id;
|
|
} else {
|
|
break;
|
|
}
|
|
} else {
|
|
break;
|
|
}
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
|
|
if sequence.len() > 1 {
|
|
Some(sequence)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
fn fuse_operations(&self, graph: &mut IRGraph, nodes: &[NodeId]) {
|
|
if nodes.is_empty() {
|
|
return;
|
|
}
|
|
|
|
// Create a fused operation node
|
|
let first_id = nodes[0];
|
|
if let Some(first_node) = graph.get_node_mut(first_id) {
|
|
// Mark as fused operation
|
|
first_node
|
|
.metadata
|
|
.insert("fused".to_string(), "true".to_string());
|
|
first_node
|
|
.metadata
|
|
.insert("fused_count".to_string(), nodes.len().to_string());
|
|
}
|
|
|
|
// Remove other nodes in the fusion group
|
|
for &node_id in &nodes[1..] {
|
|
graph.remove_node(node_id);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Memory optimization pass
|
|
pub struct MemoryOptimizationPass;
|
|
|
|
impl OptimizationPass for MemoryOptimizationPass {
|
|
fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
|
|
let mut stats = OptimizationStats::default();
|
|
|
|
// Analyze memory access patterns
|
|
let mut memory_ops = Vec::new();
|
|
|
|
for (node_id, node) in &graph.nodes {
|
|
match node.operation {
|
|
OperationType::Load | OperationType::Store => {
|
|
memory_ops.push(*node_id);
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
// Simple memory coalescing optimization
|
|
for &node_id in &memory_ops {
|
|
if let Some(node) = graph.get_node_mut(node_id) {
|
|
node.metadata
|
|
.insert("memory_optimized".to_string(), "true".to_string());
|
|
}
|
|
}
|
|
|
|
stats.memory_saved = (memory_ops.len() * 64) as u64; // Estimate
|
|
stats.passes_applied = 1;
|
|
Ok(stats)
|
|
}
|
|
|
|
fn name(&self) -> &'static str {
|
|
"MemoryOptimization"
|
|
}
|
|
}
|
|
|
|
/// Main optimizer that coordinates optimization passes
|
|
pub struct Optimizer {
|
|
passes: Vec<Box<dyn OptimizationPass>>,
|
|
}
|
|
|
|
impl Optimizer {
|
|
/// Create new optimizer
|
|
pub fn new() -> Self {
|
|
Self { passes: Vec::new() }
|
|
}
|
|
|
|
/// Add an optimization pass
|
|
pub fn add_pass(&mut self, pass: Box<dyn OptimizationPass>) {
|
|
self.passes.push(pass);
|
|
}
|
|
|
|
/// Optimize a graph
|
|
pub fn optimize(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
|
|
let mut total_stats = OptimizationStats::default();
|
|
|
|
for pass in &self.passes {
|
|
debug!("Applying optimization pass: {}", pass.name());
|
|
let pass_stats = pass.apply(graph)?;
|
|
|
|
total_stats.passes_applied += pass_stats.passes_applied;
|
|
total_stats.nodes_eliminated += pass_stats.nodes_eliminated;
|
|
total_stats.operations_fused += pass_stats.operations_fused;
|
|
total_stats.memory_saved += pass_stats.memory_saved;
|
|
}
|
|
|
|
Ok(total_stats)
|
|
}
|
|
}
|
|
|
|
impl Default for Optimizer {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Extended operation type for advanced optimizations
|
|
#[derive(Debug, Clone)]
|
|
pub enum ExtendedOperationType {
|
|
Basic(OperationType),
|
|
Fused {
|
|
operations: Vec<NodeId>,
|
|
pattern: String,
|
|
},
|
|
Optimized {
|
|
original: Box<Self>,
|
|
optimization: String,
|
|
},
|
|
}
|