onnx-ir 0.21.0

ONNX-IR is a pure Rust library for parsing ONNX models into an intermediate representation that can be used to generate code for various ML/DL frameworks
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

//! Phase 5: Finalization
//!
//! Removes unused constants and builds the OnnxGraphBuilder.

use std::{
    cell::RefCell,
    collections::{HashMap, HashSet},
    rc::Rc,
};

use crate::{
    graph_state::GraphState,
    ir::{Argument, NodeType, OnnxGraphBuilder, RawNode},
};

/// Finalize the graph by removing unused constants and building OnnxGraphBuilder
///
/// Returns OnnxGraphBuilder which still contains RawNode instances.
/// The caller should convert this to OnnxGraph using convert_to_graph().
pub(crate) fn finalize(
    nodes: &mut Vec<RawNode>,
    inputs: Vec<Argument>,
    outputs: &mut Vec<Argument>,
    state_rc: Rc<RefCell<GraphState>>,
) -> OnnxGraphBuilder {
    remove_unreferenced_constants(nodes, outputs);

    OnnxGraphBuilder {
        nodes: std::mem::take(nodes),
        inputs,
        outputs: std::mem::take(outputs),
        graph_state: Some(state_rc),
    }
}

/// Remove constant nodes that have zero runtime references
fn remove_unreferenced_constants(nodes: &mut Vec<RawNode>, outputs: &[Argument]) {
    // Build map of constant output names to node indices
    let mut constant_output_to_idx: HashMap<String, usize> = HashMap::new();
    for (idx, node) in nodes.iter().enumerate() {
        if node.node_type == NodeType::Constant {
            for output in &node.outputs {
                constant_output_to_idx.insert(output.name.clone(), idx);
            }
        }
    }

    // Count references (only Constant/Dynamic, not Static)
    let mut constant_references: HashMap<String, usize> = HashMap::new();

    for node in nodes.iter() {
        for input in &node.inputs {
            if (input.is_constant() || input.is_dynamic())
                && constant_output_to_idx.contains_key(&input.name)
            {
                *constant_references.entry(input.name.clone()).or_insert(0) += 1;
            }
        }
    }

    for output in outputs {
        if (output.is_constant() || output.is_dynamic())
            && constant_output_to_idx.contains_key(&output.name)
        {
            *constant_references.entry(output.name.clone()).or_insert(0) += 1;
        }
    }

    // Mark constants with zero references for removal
    let mut constants_to_remove = HashSet::new();
    for (output_name, &node_idx) in &constant_output_to_idx {
        let ref_count = constant_references.get(output_name).unwrap_or(&0);
        if *ref_count == 0 {
            constants_to_remove.insert(node_idx);
        }
    }

    // Filter out unreferenced constants
    let mut i = 0;
    nodes.retain(|_node| {
        let keep = !constants_to_remove.contains(&i);
        i += 1;
        keep
    });
}