onnx_runtime_optimizer/
lib.rs

1//! # `onnx-runtime-optimizer`
2//!
3//! Device-independent graph→graph optimization passes for the ORT 2.0 runtime
4//! (see `docs/architecture/ORT2.md` §18 "Optimization Passes"). This is the first Phase-2
5//! crate: pure, safe Rust graph rewriting over [`onnx_runtime_ir`] — **no**
6//! CUDA, no ORT C library, no FFI.
7//!
8//! ## What lives here
9//!
10//! | Concept | Type |
11//! |---------|------|
12//! | Pass contract | [`OptimizationPass`], [`PassContext`], [`run_passes`] |
13//! | Dead-code removal | [`DeadNodeElimination`] |
14//! | Bounded constant folding | [`ConstantFolding`] |
15//! | Reusable provider-scoped fusion machinery | [`OpFusion`], [`FusionPattern`], [`PatternMatch`] |
16//! | Errors | [`OptimizerError`], [`Result`] |
17//!
18//! ## Pipeline
19//!
20//! [`default_passes`] returns only the device-independent passes implemented
21//! here, in pipeline order: `ConstantFolding → DeadNodeElimination`.
22//! [`OpFusion`] remains in this crate as reusable machinery, but providers must
23//! schedule it themselves so private fused ops are introduced only when the
24//! selected provider can run them.
25//!
26//! ### Deferred (Phase 2b / Phase 3)
27//!
28//! The full pipeline in `docs/architecture/ORT2.md` §18.1 also lists passes that depend on
29//! crates or analyses not yet built. They are intentionally **not** implemented
30//! here and are listed in [`default_passes`]'s source in their eventual
31//! pipeline position: `ShapeInference` (the loader owns inference for now),
32//! `AttentionFusionPass`, `LayoutPropagation`, `PlacementOptimizer`,
33//! `TransferInsertion`, `InPlaceDetection`, `MemoryPlanning`,
34//! `CudaGraphRegionDetection`, and `OverlapScheduling`.
35
36#![forbid(unsafe_code)]
37
38mod constant_folding;
39mod dead_node;
40mod error;
41mod fusion;
42mod pass;
43
44pub use constant_folding::ConstantFolding;
45pub use dead_node::DeadNodeElimination;
46pub use error::{OptimizerError, Result};
47pub use fusion::{CONTRIB_DOMAIN, FusionPattern, OpFusion, PatternMatch, default_fusion_patterns};
48pub use pass::{InitializerResolver, OptimizationPass, PassContext, run_passes};
49
50/// The device-independent Phase-1 pass pipeline, in run order.
51///
52/// ```text
53/// ConstantFolding  →  DeadNodeElimination
54/// ```
55///
56/// Constant folding runs first so it can materialize shape-computation
57/// constants, then dead-node elimination prunes any node left unreachable, and
58/// then dead-node elimination prunes any node left unreachable. Operator fusion
59/// is provider-scoped because its replacements change the operator set.
60///
61/// **Deferred passes** (each in its eventual pipeline slot; see the crate-level
62/// docs for why): after `ConstantFolding` would come `ShapeInference`; after
63/// provider-scoped fusion would come `AttentionFusionPass`, then `LayoutPropagation`,
64/// `PlacementOptimizer`, `TransferInsertion`, `InPlaceDetection`,
65/// `MemoryPlanning`, `CudaGraphRegionDetection`, and `OverlapScheduling`.
66pub fn default_passes() -> Vec<Box<dyn OptimizationPass>> {
67    vec![
68        Box::new(ConstantFolding),
69        // ShapeInference — deferred (Phase 2b): the loader owns inference.
70        Box::new(DeadNodeElimination),
71        // AttentionFusionPass, LayoutPropagation, PlacementOptimizer,
72        // TransferInsertion, InPlaceDetection, MemoryPlanning,
73        // CudaGraphRegionDetection, OverlapScheduling — deferred (Phase 2b/3).
74    ]
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use onnx_runtime_ir::{DataType, Graph, Node, NodeId, static_shape};
81
82    #[test]
83    fn default_passes_lists_two() {
84        let passes = default_passes();
85        assert_eq!(passes.len(), 2);
86        assert_eq!(passes[0].name(), "ConstantFolding");
87        assert_eq!(passes[1].name(), "DeadNodeElimination");
88    }
89
90    #[test]
91    fn run_passes_pipeline_on_matmul_add_with_dead_branch() {
92        // MatMul+Add feeding an output, plus a dead Neg branch off `a`.
93        let mut g = Graph::new();
94        g.opset_imports.insert(String::new(), 17);
95        let mk =
96            |g: &mut Graph, n: &str| g.create_named_value(n, DataType::Float32, static_shape([4]));
97        let a = mk(&mut g, "a");
98        let w = mk(&mut g, "w");
99        let bias = mk(&mut g, "bias");
100        g.add_input(a);
101        g.add_input(w);
102        g.add_input(bias);
103        let m = mk(&mut g, "m");
104        g.insert_node(Node::new(
105            NodeId(0),
106            "MatMul",
107            vec![Some(a), Some(w)],
108            vec![m],
109        ));
110        let out = mk(&mut g, "out");
111        g.insert_node(Node::new(
112            NodeId(0),
113            "Add",
114            vec![Some(m), Some(bias)],
115            vec![out],
116        ));
117        g.add_output(out);
118        // Dead branch.
119        let dead = mk(&mut g, "dead");
120        g.insert_node(Node::new(NodeId(0), "Neg", vec![Some(a)], vec![dead]));
121
122        run_passes(&mut g, &default_passes(), &PassContext::new()).unwrap();
123
124        // Dead Neg removed by DCE; provider-scoped fusion is not part of the
125        // runtime default pipeline, so the standard MatMul/Add pair remains.
126        assert_eq!(g.num_nodes(), 2);
127        assert_eq!(
128            g.nodes
129                .values()
130                .filter(|node| matches!(node.op_type.as_str(), "MatMul" | "Add"))
131                .count(),
132            2
133        );
134        assert!(g.validate().is_ok());
135    }
136
137    #[test]
138    fn run_passes_is_ok_on_empty_graph() {
139        let mut g = Graph::new();
140        run_passes(&mut g, &default_passes(), &PassContext::new()).unwrap();
141        assert_eq!(g.num_nodes(), 0);
142    }
143}
onnx_runtime_optimizer/lib.rs

onnx_runtime_optimizer/
lib.rs
