onnx_runtime_optimizer/
pass.rs

1//! Pass infrastructure: the [`OptimizationPass`] trait, a minimal
2//! [`PassContext`], and the [`run_passes`] pipeline runner (see
3//! `docs/ORT2.md` §18.1).
4
5use std::sync::Arc;
6
7use onnx_runtime_ir::{Graph, WeightRef};
8use onnx_runtime_tracer::{Args, SpanGuard};
9
10use crate::error::{OptimizerError, Result};
11
12fn optimizer_span(pass_name: &str) -> Option<SpanGuard> {
13    onnx_runtime_tracer::global_context()
14        .filter(|trace| trace.is_enabled())
15        .map(|trace| trace.span(format!("optimize.{pass_name}"), "optimize"))
16}
17
18/// Resolves graph initializer descriptors to their backing bytes.
19pub trait InitializerResolver: Send + Sync {
20    /// Resolve an initializer descriptor to its raw little-endian bytes.
21    fn bytes<'a>(&'a self, weight: &'a WeightRef) -> Option<&'a [u8]>;
22}
23
24/// Shared, read-only context threaded through every pass.
25///
26/// **Phase-1 minimalism.** The design in `docs/ORT2.md` §18.1 gives this struct
27/// `cost_model`, `ep_registry`, and `target_devices` fields. Those depend on
28/// crates or analyses that do not exist yet, and none of the device-independent
29/// Phase-1 passes
30/// ([`DeadNodeElimination`](crate::DeadNodeElimination),
31/// [`ConstantFolding`](crate::ConstantFolding), [`OpFusion`](crate::OpFusion))
32/// need them. The only current service is an optional initializer resolver for
33/// EP-scoped passes that physically rewrite immutable weights.
34///
35/// It is `#[non_exhaustive]` so the Phase-2b cost-model / EP-registry /
36/// placement fields can be added without breaking downstream construction.
37#[derive(Clone, Default)]
38#[non_exhaustive]
39pub struct PassContext {
40    initializer_resolver: Option<Arc<dyn InitializerResolver>>,
41    // Phase 2b (deferred): pub cost_model: Arc<CostModel>,
42    //                      pub ep_registry: Arc<EpRegistry>,
43    //                      pub target_devices: Vec<DeviceId>,
44}
45
46impl std::fmt::Debug for PassContext {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        f.debug_struct("PassContext")
49            .field(
50                "initializer_resolver",
51                &self.initializer_resolver.as_ref().map(|_| "<resolver>"),
52            )
53            .finish()
54    }
55}
56
57impl PassContext {
58    /// A context with no device/cost information (the Phase-1 default).
59    pub fn new() -> Self {
60        Self::default()
61    }
62
63    /// Attach the resolver that exposes inline or externally mapped initializer
64    /// bytes to passes that rewrite immutable weights.
65    pub fn with_initializer_resolver(mut self, resolver: Arc<dyn InitializerResolver>) -> Self {
66        self.initializer_resolver = Some(resolver);
67        self
68    }
69
70    /// Resolve an initializer's raw bytes. Inline initializers are always
71    /// available; external references require an attached resolver.
72    pub fn initializer_bytes<'b>(&'b self, weight: &'b WeightRef) -> Option<&'b [u8]> {
73        match weight {
74            WeightRef::Inline(tensor) => Some(&tensor.data),
75            WeightRef::External { .. } => self.initializer_resolver.as_deref()?.bytes(weight),
76        }
77    }
78}
79
80/// A single graph→graph rewrite (see `docs/ORT2.md` §18.1).
81///
82/// Passes mutate the [`Graph`] in place and must preserve its structural
83/// invariants; [`postconditions`](OptimizationPass::postconditions) is checked
84/// after each pass in debug builds by [`run_passes`].
85pub trait OptimizationPass: Send + Sync {
86    /// A short, stable name for logging and error messages.
87    fn name(&self) -> &str;
88
89    /// Apply the rewrite in place.
90    fn run(&self, graph: &mut Graph, ctx: &PassContext) -> Result<()>;
91
92    /// Invariants that must hold after this pass. The default requires the
93    /// graph to pass full structural validation (`Graph::validate`).
94    fn postconditions(&self, graph: &Graph) -> Result<()> {
95        graph
96            .validate()
97            .map_err(|errors| OptimizerError::PostconditionFailed {
98                pass: self.name().to_string(),
99                errors,
100            })
101    }
102}
103
104/// Run `passes` over `graph` in order.
105///
106/// Each pass runs to completion, then — in debug builds only — its
107/// [`postconditions`](OptimizationPass::postconditions) are checked. In release
108/// builds the postcondition check is compiled out for speed, matching the
109/// "checked in debug builds" contract of `docs/ORT2.md` §18.1.
110pub fn run_passes(
111    graph: &mut Graph,
112    passes: &[Box<dyn OptimizationPass>],
113    ctx: &PassContext,
114) -> Result<()> {
115    for pass in passes {
116        let nodes_before = graph.num_nodes();
117        let values_before = graph.values.len();
118        let mut span = optimizer_span(pass.name());
119        if let Some(span) = span.as_mut() {
120            span.set_args(
121                Args::new()
122                    .with("pass", pass.name().to_string())
123                    .with("nodes_before", nodes_before as u64)
124                    .with("values_before", values_before as u64),
125            );
126        }
127        pass.run(graph, ctx)?;
128        #[cfg(debug_assertions)]
129        pass.postconditions(graph)?;
130        if let Some(span) = span.as_mut() {
131            span.set_args(
132                Args::new()
133                    .with("pass", pass.name().to_string())
134                    .with("nodes_before", nodes_before as u64)
135                    .with("nodes_after", graph.num_nodes() as u64)
136                    .with(
137                        "nodes_delta",
138                        graph.num_nodes() as i64 - nodes_before as i64,
139                    )
140                    .with("values_before", values_before as u64)
141                    .with("values_after", graph.values.len() as u64),
142            );
143        }
144    }
145    Ok(())
146}
onnx_runtime_optimizer/pass.rs

onnx_runtime_optimizer/
pass.rs
