nxpu_opt/
lib.rs

1//! IR optimization passes for NxPU.
2//!
3//! Provides a [`Pass`] trait, a [`PassManager`] with fixed-point iteration,
4//! and built-in optimization passes (constant folding, FMA fusion, dead code
5//! elimination).
6
7pub mod calibrate;
8mod const_fold;
9mod cse;
10mod dce;
11mod fma_fusion;
12pub mod fusion;
13pub mod layout;
14pub mod memory;
15pub mod quantize;
16pub mod schedule;
17pub mod shape;
18pub mod tiling;
19mod validation;
20pub mod vectorize;
21pub mod workgroup;
22
23pub use calibrate::{
24    CalibrationDataset, CalibrationError, CalibrationMethod, CalibrationResult, HistogramCollector,
25    TensorHistogram, calibrate, calibrate_kl_divergence, calibrate_minmax, calibrate_percentile,
26    per_channel_quantize, run_calibration,
27};
28pub use const_fold::ConstantFolding;
29pub use cse::CommonSubexprElimination;
30pub use dce::DeadCodeElimination;
31pub use fma_fusion::FmaFusion;
32pub use fusion::OperatorFusion;
33pub use layout::{
34    LayoutTransform, TransposeInsertion, TransposeRecord, layout_permutation, reorder_dims,
35};
36pub use memory::{LiveInterval, MemoryPlanning, plan_memory};
37// Re-export canonical memory plan types from nxpu-backend-core (via memory module).
38pub use memory::{BufferAllocation, MemoryPlan, TensorId};
39pub use quantize::{
40    CalibrationData, F32ToBf16, F32ToF16, F32ToInt8, MixedPrecisionPass, MixedPrecisionPolicy,
41    PerChannelQuantParams, QuantizationParams,
42};
43pub use schedule::{Schedule, SchedulePass, ScheduleSlot, compute_schedules, format_schedule};
44pub use shape::ShapeInference;
45pub use tiling::{
46    Conv2DShape as TilingConv2DShape, MatMulShape as TilingMatMulShape, TileConfig, TilingDefaults,
47    TilingPass, TilingPlan, tile_conv2d, tile_matmul,
48};
49pub use validation::{IrValidation, ValidationWarning, collect_warnings};
50pub use vectorize::{VectorWidth, VectorizationHint, VectorizationPass, analyze_vectorization};
51pub use workgroup::{
52    OccupancyResult, WorkgroupHwParams, WorkgroupOptimization, calculate_occupancy,
53    optimize_workgroup_size,
54};
55
56use std::fmt::Debug;
57
58use nxpu_ir::Module;
59
60/// An optimization pass that transforms an IR module.
61pub trait Pass: Debug {
62    /// Human-readable name of the pass.
63    fn name(&self) -> &str;
64
65    /// Run the pass on a module. Returns `true` if anything was modified.
66    fn run(&self, module: &mut Module) -> bool;
67}
68
69/// Optimization level.
70#[derive(Clone, Copy, Debug, PartialEq, Eq)]
71pub enum OptLevel {
72    /// No optimizations.
73    O0,
74    /// Basic optimizations (constant folding, FMA fusion, operator fusion, DCE).
75    O1,
76    /// Aggressive optimizations (same as O1 for now).
77    O2,
78}
79
80/// Maximum number of fixed-point iterations before giving up.
81const MAX_ITERATIONS: usize = 10;
82
83/// Runs passes in sequence with fixed-point iteration.
84pub struct PassManager {
85    /// Passes run once before the fixed-point loop (e.g. validation).
86    pre_passes: Vec<Box<dyn Pass>>,
87    /// Passes run in the fixed-point loop.
88    passes: Vec<Box<dyn Pass>>,
89}
90
91impl Default for PassManager {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97impl PassManager {
98    /// Creates an empty pass manager with no passes.
99    pub fn new() -> Self {
100        Self {
101            pre_passes: Vec::new(),
102            passes: Vec::new(),
103        }
104    }
105
106    /// Creates a pass manager with passes appropriate for the given level.
107    pub fn for_level(level: OptLevel) -> Self {
108        let mut pm = Self::new();
109        match level {
110            OptLevel::O0 => {}
111            OptLevel::O1 | OptLevel::O2 => {
112                pm.add_pre_pass(Box::new(IrValidation));
113                pm.add_pass(Box::new(ConstantFolding));
114                pm.add_pass(Box::new(FmaFusion));
115                pm.add_pass(Box::new(CommonSubexprElimination));
116                pm.add_pass(Box::new(OperatorFusion));
117                pm.add_pass(Box::new(DeadCodeElimination));
118            }
119        }
120        pm
121    }
122
123    /// Adds a pass to run once before the fixed-point loop.
124    pub fn add_pre_pass(&mut self, pass: Box<dyn Pass>) {
125        self.pre_passes.push(pass);
126    }
127
128    /// Adds a pass to the fixed-point pipeline.
129    pub fn add_pass(&mut self, pass: Box<dyn Pass>) {
130        self.passes.push(pass);
131    }
132
133    /// Runs pre-passes once, then iterates the main passes until a fixed point
134    /// is reached or the iteration limit.
135    pub fn run(&self, module: &mut Module) {
136        for pass in &self.pre_passes {
137            pass.run(module);
138            log::debug!("pre-pass '{}' completed", pass.name());
139        }
140
141        for iteration in 0..MAX_ITERATIONS {
142            let mut changed = false;
143            for pass in &self.passes {
144                let pass_changed = pass.run(module);
145                log::debug!(
146                    "pass '{}' iteration {}: changed={}",
147                    pass.name(),
148                    iteration,
149                    pass_changed
150                );
151                changed |= pass_changed;
152            }
153            if !changed {
154                log::debug!("fixed point reached after {} iteration(s)", iteration + 1);
155                break;
156            }
157        }
158    }
159}
160
161/// Convenience function: runs O1 optimization passes on a module.
162pub fn optimize(module: &mut Module) {
163    PassManager::for_level(OptLevel::O1).run(module);
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    #[test]
171    fn optimize_empty_module() {
172        let mut module = Module::default();
173        optimize(&mut module);
174        // Should not panic; module remains empty.
175        assert_eq!(module.entry_points.len(), 0);
176    }
177
178    #[test]
179    fn pass_manager_o0_is_noop() {
180        let pm = PassManager::for_level(OptLevel::O0);
181        let mut module = Module::default();
182        pm.run(&mut module);
183        assert_eq!(module.entry_points.len(), 0);
184    }
185
186    #[test]
187    fn pass_manager_o1_runs() {
188        let pm = PassManager::for_level(OptLevel::O1);
189        let mut module = Module::default();
190        pm.run(&mut module);
191        // No crash on empty module.
192    }
193}
nxpu_opt/lib.rs

nxpu_opt/
lib.rs
