onnx_runtime_optimizer/
constant_folding.rs

1//! Constant folding: replace a node whose inputs are *all* constant
2//! (initializers) with a precomputed initializer (see `docs/ORT2.md` §18.1).
3//!
4//! ## Boundary (deliberately conservative)
5//!
6//! Fully general constant folding needs a kernel executor — the optimizer has
7//! none — so this pass folds only what the IR can compute *directly and
8//! exactly*, reusing the lessons of the loader's "const-fold-lite"
9//! (`crates/onnx-runtime-loader/src/shape_inference.rs`):
10//!
11//! * **`Constant`** nodes are materialized into initializers (always safe).
12//! * **`Shape`** on a fully-static input becomes an `int64` initializer.
13//! * **Elementwise integer `Add`/`Sub`/`Mul`** on two *same-shape* constant
14//!   `int32`/`int64` tensors are evaluated with **checked** arithmetic; any
15//!   overflow aborts the fold rather than emit a wrong constant.
16//!
17//! Everything else is left untouched. Folding is bounded to
18//! [`MAX_FOLD_ELEMS`] elements so large weight tensors are never materialized,
19//! and dispatch is purely on op type — no model-specific names. The invariant
20//! is: **never produce a wrong constant.** When in doubt, do not fold.
21
22use onnx_runtime_ir::{
23    Attribute, DataType, Graph, NodeId, TensorData, ValueId, WeightRef, as_static_shape,
24    static_shape,
25};
26
27use crate::error::Result;
28use crate::pass::{OptimizationPass, PassContext};
29
30/// Upper bound on the number of elements this pass will materialize. Keeps
31/// folding limited to shape-computation-sized tensors, never model weights.
32const MAX_FOLD_ELEMS: usize = 1024;
33
34/// Folds constant-input nodes into initializers (bounded, integer/shape only).
35#[derive(Clone, Copy, Debug, Default)]
36pub struct ConstantFolding;
37
38impl OptimizationPass for ConstantFolding {
39    fn name(&self) -> &str {
40        "ConstantFolding"
41    }
42
43    fn run(&self, graph: &mut Graph, _ctx: &PassContext) -> Result<()> {
44        // Iterate to a fixpoint: folding one node may make its consumers
45        // foldable in the next round (e.g. Constant -> Add).
46        loop {
47            let mut changed = false;
48            let node_ids: Vec<NodeId> = graph.nodes.keys().collect();
49            for nid in node_ids {
50                if !graph.nodes.contains(nid) {
51                    continue;
52                }
53                let node = graph.node(nid).clone();
54                if !matches!(node.domain.as_str(), "" | "ai.onnx") {
55                    continue;
56                }
57                if node.outputs.len() != 1 {
58                    continue;
59                }
60                let out = node.outputs[0];
61
62                let folded: Option<TensorData> = match node.op_type.as_str() {
63                    "Constant" => eval_constant(&node),
64                    "Shape" => fold_shape(graph, &node),
65                    "Add" | "Sub" | "Mul" => fold_binary_int(graph, &node),
66                    _ => None,
67                };
68
69                let Some(tensor) = folded else { continue };
70
71                // Only fold outputs that are still needed (have a consumer or
72                // are graph outputs); dead outputs are DCE's job and folding
73                // them would leave a stale initializer referencing a GC'd id.
74                let needed = graph.outputs.contains(&out)
75                    || graph
76                        .try_value(out)
77                        .is_some_and(|v| !v.consumers.is_empty());
78                if !needed {
79                    continue;
80                }
81
82                graph.remove_node(nid);
83                // The output survives because it is needed; retype it to the
84                // folded tensor and back it with an inline initializer.
85                if graph.try_value(out).is_some() {
86                    let dims = tensor.dims.clone();
87                    let dtype = tensor.dtype;
88                    let v = graph.value_mut(out);
89                    v.dtype = dtype;
90                    v.shape = static_shape(dims);
91                    graph.set_initializer(out, WeightRef::Inline(tensor));
92                    changed = true;
93                }
94            }
95            if !changed {
96                break;
97            }
98        }
99        Ok(())
100    }
101}
102
103/// The inline constant tensor backing `value`, if any (external weights, which
104/// are large, are never folded).
105fn inline_const(graph: &Graph, value: ValueId) -> Option<&TensorData> {
106    match graph.initializers.get(&value)? {
107        WeightRef::Inline(t) => Some(t),
108        WeightRef::External { .. } => None,
109    }
110}
111
112/// Materialize a `Constant` node's value into a concrete [`TensorData`].
113fn eval_constant(node: &onnx_runtime_ir::Node) -> Option<TensorData> {
114    if let Some(Attribute::Tensor(t)) = node.attr("value") {
115        return Some(t.clone());
116    }
117    if let Some(ints) = node.attr("value_ints").and_then(Attribute::as_ints) {
118        let mut data = Vec::with_capacity(ints.len() * 8);
119        for &i in ints {
120            data.extend_from_slice(&i.to_le_bytes());
121        }
122        return Some(TensorData::from_raw(DataType::Int64, vec![ints.len()], data));
123    }
124    if let Some(i) = node.attr("value_int").and_then(Attribute::as_int) {
125        return Some(TensorData::from_raw(
126            DataType::Int64,
127            Vec::new(),
128            i.to_le_bytes().to_vec(),
129        ));
130    }
131    None
132}
133
134/// Fold `Shape(x)` when `x` has a fully-static shape into an `int64` vector.
135///
136/// Conservative: bails if `start`/`end` attributes are present (a slice of the
137/// shape) so we never emit a partial result.
138fn fold_shape(graph: &Graph, node: &onnx_runtime_ir::Node) -> Option<TensorData> {
139    if node.attr("start").is_some() || node.attr("end").is_some() {
140        return None;
141    }
142    let input = node.inputs.first().copied().flatten()?;
143    let shape = &graph.try_value(input)?.shape;
144    let dims = as_static_shape(shape)?;
145    if dims.len() > MAX_FOLD_ELEMS {
146        return None;
147    }
148    let mut data = Vec::with_capacity(dims.len() * 8);
149    for &d in &dims {
150        data.extend_from_slice(&(d as i64).to_le_bytes());
151    }
152    Some(TensorData::from_raw(DataType::Int64, vec![dims.len()], data))
153}
154
155/// Fold elementwise integer `Add`/`Sub`/`Mul` on two same-shape constant
156/// tensors. Uses checked arithmetic; overflow aborts (returns `None`).
157fn fold_binary_int(graph: &Graph, node: &onnx_runtime_ir::Node) -> Option<TensorData> {
158    if node.inputs.len() != 2 {
159        return None;
160    }
161    let a = inline_const(graph, node.inputs[0]?)?;
162    let b = inline_const(graph, node.inputs[1]?)?;
163    if a.dtype != b.dtype || a.dims != b.dims {
164        return None; // no broadcasting / mixed dtype in v1
165    }
166    if !matches!(a.dtype, DataType::Int32 | DataType::Int64) {
167        return None;
168    }
169    let numel = a.numel();
170    if numel > MAX_FOLD_ELEMS {
171        return None;
172    }
173    let op = node.op_type.as_str();
174    let apply = |x: i64, y: i64| -> Option<i64> {
175        match op {
176            "Add" => x.checked_add(y),
177            "Sub" => x.checked_sub(y),
178            "Mul" => x.checked_mul(y),
179            _ => None,
180        }
181    };
182
183    match a.dtype {
184        DataType::Int64 => {
185            let (xs, ys) = (read_i64(a)?, read_i64(b)?);
186            let mut data = Vec::with_capacity(numel * 8);
187            for (x, y) in xs.into_iter().zip(ys) {
188                data.extend_from_slice(&apply(x, y)?.to_le_bytes());
189            }
190            Some(TensorData::from_raw(DataType::Int64, a.dims.clone(), data))
191        }
192        DataType::Int32 => {
193            let (xs, ys) = (read_i32(a)?, read_i32(b)?);
194            let mut data = Vec::with_capacity(numel * 4);
195            for (x, y) in xs.into_iter().zip(ys) {
196                let r = apply(x as i64, y as i64)?;
197                let r32: i32 = r.try_into().ok()?; // must fit back into i32
198                data.extend_from_slice(&r32.to_le_bytes());
199            }
200            Some(TensorData::from_raw(DataType::Int32, a.dims.clone(), data))
201        }
202        _ => None,
203    }
204}
205
206fn read_i64(t: &TensorData) -> Option<Vec<i64>> {
207    if t.data.len() != t.numel() * 8 {
208        return None;
209    }
210    Some(
211        t.data
212            .chunks_exact(8)
213            .map(|c| i64::from_le_bytes(c.try_into().unwrap()))
214            .collect(),
215    )
216}
217
218fn read_i32(t: &TensorData) -> Option<Vec<i32>> {
219    if t.data.len() != t.numel() * 4 {
220        return None;
221    }
222    Some(
223        t.data
224            .chunks_exact(4)
225            .map(|c| i32::from_le_bytes(c.try_into().unwrap()))
226            .collect(),
227    )
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use onnx_runtime_ir::{Node, NodeId};
234
235    fn int64_tensor(dims: Vec<usize>, vals: &[i64]) -> TensorData {
236        let mut data = Vec::new();
237        for &v in vals {
238            data.extend_from_slice(&v.to_le_bytes());
239        }
240        TensorData::from_raw(DataType::Int64, dims, data)
241    }
242
243    fn const_init(graph: &mut Graph, name: &str, dims: Vec<usize>, vals: &[i64]) -> ValueId {
244        let shape = static_shape(dims.clone());
245        let v = graph.create_named_value(name, DataType::Int64, shape);
246        graph.set_initializer(v, WeightRef::Inline(int64_tensor(dims, vals)));
247        v
248    }
249
250    #[test]
251    fn folds_add_of_two_const_inputs() {
252        let mut g = Graph::new();
253        g.opset_imports.insert(String::new(), 17);
254        let a = const_init(&mut g, "a", vec![3], &[1, 2, 3]);
255        let b = const_init(&mut g, "b", vec![3], &[10, 20, 30]);
256        let out = g.create_named_value("out", DataType::Int64, static_shape([3]));
257        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(a), Some(b)], vec![out]));
258        g.add_output(out);
259
260        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
261
262        assert_eq!(g.num_nodes(), 0, "Add should be folded away");
263        let t = inline_const(&g, out).expect("out is now an initializer");
264        assert_eq!(read_i64(t).unwrap(), vec![11, 22, 33]);
265        assert!(g.validate().is_ok());
266    }
267
268    #[test]
269    fn folds_sub_and_mul() {
270        for (op, expect) in [("Sub", vec![9, 18, 27]), ("Mul", vec![10, 40, 90])] {
271            let mut g = Graph::new();
272            g.opset_imports.insert(String::new(), 17);
273            let a = const_init(&mut g, "a", vec![3], &[10, 20, 30]);
274            let b = const_init(&mut g, "b", vec![3], &[1, 2, 3]);
275            let out = g.create_named_value("out", DataType::Int64, static_shape([3]));
276            g.insert_node(Node::new(NodeId(0), op, vec![Some(a), Some(b)], vec![out]));
277            g.add_output(out);
278
279            ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
280            let t = inline_const(&g, out).unwrap();
281            assert_eq!(read_i64(t).unwrap(), expect, "op {op}");
282        }
283    }
284
285    #[test]
286    fn does_not_fold_when_one_input_is_non_const() {
287        let mut g = Graph::new();
288        g.opset_imports.insert(String::new(), 17);
289        let a = const_init(&mut g, "a", vec![3], &[1, 2, 3]);
290        // `b` is a graph input, not a constant.
291        let b = g.create_named_value("b", DataType::Int64, static_shape([3]));
292        g.add_input(b);
293        let out = g.create_named_value("out", DataType::Int64, static_shape([3]));
294        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(a), Some(b)], vec![out]));
295        g.add_output(out);
296
297        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
298        assert_eq!(g.num_nodes(), 1, "must not fold with a non-const input");
299        assert!(inline_const(&g, out).is_none());
300        assert!(g.validate().is_ok());
301    }
302
303    #[test]
304    fn does_not_fold_mismatched_shapes() {
305        let mut g = Graph::new();
306        g.opset_imports.insert(String::new(), 17);
307        let a = const_init(&mut g, "a", vec![3], &[1, 2, 3]);
308        let b = const_init(&mut g, "b", vec![2], &[10, 20]);
309        let out = g.create_named_value("out", DataType::Int64, static_shape([3]));
310        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(a), Some(b)], vec![out]));
311        g.add_output(out);
312
313        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
314        assert_eq!(g.num_nodes(), 1, "no broadcasting in v1");
315    }
316
317    #[test]
318    fn does_not_fold_overflow() {
319        let mut g = Graph::new();
320        g.opset_imports.insert(String::new(), 17);
321        let a = const_init(&mut g, "a", vec![1], &[i64::MAX]);
322        let b = const_init(&mut g, "b", vec![1], &[1]);
323        let out = g.create_named_value("out", DataType::Int64, static_shape([1]));
324        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(a), Some(b)], vec![out]));
325        g.add_output(out);
326
327        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
328        assert_eq!(g.num_nodes(), 1, "overflow must abort the fold");
329    }
330
331    #[test]
332    fn folds_constant_node_to_initializer() {
333        let mut g = Graph::new();
334        g.opset_imports.insert(String::new(), 17);
335        let out = g.create_named_value("c", DataType::Int64, static_shape([2]));
336        let mut node = Node::new(NodeId(0), "Constant", vec![], vec![out]);
337        node.attributes
338            .insert("value".into(), Attribute::Tensor(int64_tensor(vec![2], &[7, 8])));
339        g.insert_node(node);
340        // Keep `out` alive with a consumer.
341        let sink = g.create_named_value("sink", DataType::Int64, static_shape([2]));
342        g.insert_node(Node::new(NodeId(0), "Identity", vec![Some(out)], vec![sink]));
343        g.add_output(sink);
344
345        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
346        assert!(g.try_node(NodeId(0)).is_none(), "Constant folded away");
347        let t = inline_const(&g, out).unwrap();
348        assert_eq!(read_i64(t).unwrap(), vec![7, 8]);
349        assert!(g.validate().is_ok());
350    }
351
352    #[test]
353    fn folds_shape_of_static_input() {
354        let mut g = Graph::new();
355        g.opset_imports.insert(String::new(), 17);
356        let x = g.create_named_value("x", DataType::Float32, static_shape([2, 3, 4]));
357        g.add_input(x);
358        let out = g.create_named_value("s", DataType::Int64, static_shape([3]));
359        g.insert_node(Node::new(NodeId(0), "Shape", vec![Some(x)], vec![out]));
360        g.add_output(out);
361
362        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
363        let t = inline_const(&g, out).expect("Shape folded to initializer");
364        assert_eq!(read_i64(t).unwrap(), vec![2, 3, 4]);
365        assert!(g.validate().is_ok());
366    }
367
368    #[test]
369    fn folds_transitively_to_fixpoint() {
370        // Constant c1, Constant c2, then Add(c1, c2) -> out. All should fold.
371        let mut g = Graph::new();
372        g.opset_imports.insert(String::new(), 17);
373
374        let c1 = g.create_named_value("c1", DataType::Int64, static_shape([2]));
375        let mut n1 = Node::new(NodeId(0), "Constant", vec![], vec![c1]);
376        n1.attributes
377            .insert("value".into(), Attribute::Tensor(int64_tensor(vec![2], &[1, 2])));
378        g.insert_node(n1);
379
380        let c2 = g.create_named_value("c2", DataType::Int64, static_shape([2]));
381        let mut n2 = Node::new(NodeId(0), "Constant", vec![], vec![c2]);
382        n2.attributes
383            .insert("value".into(), Attribute::Tensor(int64_tensor(vec![2], &[3, 4])));
384        g.insert_node(n2);
385
386        let out = g.create_named_value("out", DataType::Int64, static_shape([2]));
387        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(c1), Some(c2)], vec![out]));
388        g.add_output(out);
389
390        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
391        assert_eq!(g.num_nodes(), 0, "both constants and the Add fold away");
392        let t = inline_const(&g, out).unwrap();
393        assert_eq!(read_i64(t).unwrap(), vec![4, 6]);
394        assert!(g.validate().is_ok());
395    }
396
397    #[test]
398    fn does_not_fold_float_binary() {
399        let mut g = Graph::new();
400        g.opset_imports.insert(String::new(), 17);
401        let mk = |g: &mut Graph, name: &str| {
402            let v = g.create_named_value(name, DataType::Float32, static_shape([2]));
403            g.set_initializer(
404                v,
405                WeightRef::Inline(TensorData::from_raw(
406                    DataType::Float32,
407                    vec![2],
408                    vec![0u8; 8],
409                )),
410            );
411            v
412        };
413        let a = mk(&mut g, "a");
414        let b = mk(&mut g, "b");
415        let out = g.create_named_value("out", DataType::Float32, static_shape([2]));
416        g.insert_node(Node::new(NodeId(0), "Mul", vec![Some(a), Some(b)], vec![out]));
417        g.add_output(out);
418
419        ConstantFolding.run(&mut g, &PassContext::new()).unwrap();
420        assert_eq!(g.num_nodes(), 1, "float folding is out of scope in v1");
421    }
422}
onnx_runtime_optimizer/constant_folding.rs

onnx_runtime_optimizer/
constant_folding.rs
