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(mut self, resolver: Arc<dyn InitializerResolver>) -> Self {
59        self.initializer_resolver = Some(resolver);
60        self
61    }
62
63    /// Resolve an initializer's raw bytes. Inline initializers are always
64    /// available; external references require an attached resolver.
65    pub fn initializer_bytes<'b>(&'b self, weight: &'b WeightRef) -> Option<&'b [u8]> {
66        match weight {
67            WeightRef::Inline(tensor) => Some(&tensor.data),
68            WeightRef::External { .. } => self.initializer_resolver.as_deref()?.bytes(weight),
69        }
70    }
71}
72
73/// A single graph→graph rewrite (see `docs/ORT2.md` §18.1).
74///
75/// Passes mutate the [`Graph`] in place and must preserve its structural
76/// invariants; [`postconditions`](OptimizationPass::postconditions) is checked
77/// after each pass in debug builds by [`run_passes`].
78pub trait OptimizationPass: Send + Sync {
79    /// A short, stable name for logging and error messages.
80    fn name(&self) -> &str;
81
82    /// Apply the rewrite in place.
83    fn run(&self, graph: &mut Graph, ctx: &PassContext) -> Result<()>;
84
85    /// Invariants that must hold after this pass. The default requires the
86    /// graph to pass full structural validation (`Graph::validate`).
87    fn postconditions(&self, graph: &Graph) -> Result<()> {
88        graph
89            .validate()
90            .map_err(|errors| OptimizerError::PostconditionFailed {
91                pass: self.name().to_string(),
92                errors,
93            })
94    }
95}
96
97/// Run `passes` over `graph` in order.
98///
99/// Each pass runs to completion, then — in debug builds only — its
100/// [`postconditions`](OptimizationPass::postconditions) are checked. In release
101/// builds the postcondition check is compiled out for speed, matching the
102/// "checked in debug builds" contract of `docs/ORT2.md` §18.1.
103pub fn run_passes(
104    graph: &mut Graph,
105    passes: &[Box<dyn OptimizationPass>],
106    ctx: &PassContext,
107) -> Result<()> {
108    for pass in passes {
109        pass.run(graph, ctx)?;
110        #[cfg(debug_assertions)]
111        pass.postconditions(graph)?;
112    }
113    Ok(())
114}
onnx_runtime_optimizer/pass.rs

onnx_runtime_optimizer/
pass.rs
