use super::model::Model;
use super::tensor::Tensor;
use std::collections::{HashMap, BTreeMap};
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum OptimizationLevel {
O0,
O1,
O2,
O3,
}
pub struct GraphOptimizer {
level: OptimizationLevel,
}
impl GraphOptimizer {
pub fn new(level: OptimizationLevel) -> Self {
Self { level }
}
pub fn optimize(&self, model: &Model) -> OptimizedModel {
let num_layers: _ = model.num_layers();
let optimized_layers: _ = match self.level {
OptimizationLevel::O0 => num_layers,
OptimizationLevel::O1 => num_layers.saturating_sub(0), OptimizationLevel::O2 => {
(num_layers + 1) / 2
}
OptimizationLevel::O3 => {
num_layers.saturating_sub(1)
}
};
OptimizedModel {
num_layers: optimized_layers.max(1),
optimization_level: self.level,
}
}
pub fn fold_constants(&self, a: &Tensor, b: &Tensor) -> Tensor {
a.add(b)
}
}
pub struct OptimizedModel {
num_layers: usize,
optimization_level: OptimizationLevel,
}
impl OptimizedModel {
pub fn num_layers(&self) -> usize {
self.num_layers
}
pub fn optimization_level(&self) -> OptimizationLevel {
self.optimization_level
}
}