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

onnx_runtime_optimizer/
pass.rs
