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 onnx_runtime_ir::Graph;
6
7use crate::error::{OptimizerError, Result};
8
9/// Shared, read-only context threaded through every pass.
10///
11/// **Phase-1 minimalism.** The design in `docs/ORT2.md` §18.1 gives this struct
12/// `cost_model`, `ep_registry`, and `target_devices` fields. Those depend on
13/// the `onnx-runtime-cost-model` and EP-registry crates, which do not exist
14/// yet, and none of the device-independent Phase-1 passes
15/// ([`DeadNodeElimination`](crate::DeadNodeElimination),
16/// [`ConstantFolding`](crate::ConstantFolding), [`OpFusion`](crate::OpFusion))
17/// need them. The context is intentionally empty for now.
18///
19/// It is `#[non_exhaustive]` so the Phase-2b cost-model / EP-registry /
20/// placement fields can be added without breaking downstream construction.
21#[derive(Clone, Debug, Default)]
22#[non_exhaustive]
23pub struct PassContext {
24    // Phase 2b (deferred): pub cost_model: Arc<CostModel>,
25    //                      pub ep_registry: Arc<EpRegistry>,
26    //                      pub target_devices: Vec<DeviceId>,
27}
28
29impl PassContext {
30    /// A context with no device/cost information (the Phase-1 default).
31    pub fn new() -> Self {
32        Self::default()
33    }
34}
35
36/// A single graph→graph rewrite (see `docs/ORT2.md` §18.1).
37///
38/// Passes mutate the [`Graph`] in place and must preserve its structural
39/// invariants; [`postconditions`](OptimizationPass::postconditions) is checked
40/// after each pass in debug builds by [`run_passes`].
41pub trait OptimizationPass: Send + Sync {
42    /// A short, stable name for logging and error messages.
43    fn name(&self) -> &str;
44
45    /// Apply the rewrite in place.
46    fn run(&self, graph: &mut Graph, ctx: &PassContext) -> Result<()>;
47
48    /// Invariants that must hold after this pass. The default requires the
49    /// graph to pass full structural validation (`Graph::validate`).
50    fn postconditions(&self, graph: &Graph) -> Result<()> {
51        graph
52            .validate()
53            .map_err(|errors| OptimizerError::PostconditionFailed {
54                pass: self.name().to_string(),
55                errors,
56            })
57    }
58}
59
60/// Run `passes` over `graph` in order.
61///
62/// Each pass runs to completion, then — in debug builds only — its
63/// [`postconditions`](OptimizationPass::postconditions) are checked. In release
64/// builds the postcondition check is compiled out for speed, matching the
65/// "checked in debug builds" contract of `docs/ORT2.md` §18.1.
66pub fn run_passes(
67    graph: &mut Graph,
68    passes: &[Box<dyn OptimizationPass>],
69    ctx: &PassContext,
70) -> Result<()> {
71    for pass in passes {
72        pass.run(graph, ctx)?;
73        #[cfg(debug_assertions)]
74        pass.postconditions(graph)?;
75    }
76    Ok(())
77}
onnx_runtime_optimizer/pass.rs

onnx_runtime_optimizer/
pass.rs
