Overview
The ONNX IR pass infrastructure is designed for graph construction, analysis, and transformation. Passes are composable units that transform ONNX models in a well-defined way.
Key Concepts
- Pass: A transformation that takes an
ir.Model and returns a PassResult containing the transformed model and a boolean indicating if modifications were made - InPlacePass: A pass that modifies the input model directly and returns it (most efficient)
- FunctionalPass: A pass that returns a new model without modifying the input
- PassManager/Sequential: Composes multiple passes to run in sequence
Pass Infrastructure
Base Classes
All passes inherit from one of these base classes defined in onnx_ir.passes:
InPlacePass
python
class MyPass(ir.passes.InPlacePass):
"""Most common pass type - modifies model in place."""
def call(self, model: ir.Model) -> ir.passes.PassResult:
modified = False
# Transform the model
return ir.passes.PassResult(model, modified=modified)
Use when: You want efficient in-place mutation (recommended for most passes)
Properties:
in_place = True (automatically set)changes_input = True (automatically set)
FunctionalPass
python
class MyPass(ir.passes.FunctionalPass):
"""Pure functional pass - does not modify input."""
def call(self, model: ir.Model) -> ir.passes.PassResult:
# Must return a different model object
new_model = model.clone()
# Transform new_model
return ir.passes.PassResult(new_model, modified=True)
Use when: You need to preserve the original model unchanged
Properties:
in_place = False (automatically set)changes_input = False (automatically set)
Pass Lifecycle
- Preconditions (
requires method): Check input model validity (optional) - Transformation (
call method): Apply the transformation - Postconditions (
ensures method): Validate output model (optional)
python
class MyPass(ir.passes.InPlacePass):
def requires(self, model: ir.Model) -> None:
"""Validate preconditions. Raise PreconditionError if violated."""
# Example: Ensure specific opset version
if model.graph.opset_imports.get("", 0) < 13:
raise ir.passes.PreconditionError("Requires opset >= 13")
def call(self, model: ir.Model) -> ir.passes.PassResult:
"""Main transformation logic."""
modified = False
# ... transformation code ...
return ir.passes.PassResult(model, modified=modified)
def ensures(self, model: ir.Model) -> None:
"""Validate postconditions. Raise PostconditionError if violated."""
# Example: Check model validity
pass
Best Practices for Pass Implementation
1. Graph Traversal
Traverse All Nodes (Including Subgraphs)
python
import onnx_ir as ir
# Use RecursiveGraphIterator to process all nodes including subgraphs
for node in ir.traversal.RecursiveGraphIterator(model.graph):
# Process node
if node.op_type == "Identity":
# ... handle identity node ...
pass
Process Functions
python
# Don't forget to process functions in the model
for function in model.functions.values():
for node in ir.traversal.RecursiveGraphIterator(function):
# Process node in function
pass
Simple Graph Iteration
python
# For non-recursive iteration of the main graph
for node in model.graph:
# Process only direct nodes (no subgraphs)
pass
# Reverse iteration (useful for removal)
for node in reversed(model.graph):
# Process in reverse topological order
pass
2. Node Manipulation
Safely Remove Nodes
python
# Always use safe=True to ensure proper cleanup
graph.remove(node, safe=True)
Create New Nodes
python
# Create a node with the ir.node() helper
new_node = ir.node(
"Identity",
inputs=[input_value],
outputs=[
ir.Value(
name="output_name",
type=ir.TensorType(ir.DataType.FLOAT),
shape=ir.Shape([1, 3, 224, 224]),
)
],
)
# Insert the node at a specific position
graph.insert_before(reference_node, new_node)
graph.insert_after(reference_node, new_node)
# Or append to the end
graph.append(new_node)
Modify Node Attributes
python
# Access attributes as a dictionary
if "training_mode" in node.attributes:
node.attributes.pop("training_mode")
# Add new attributes
node.attributes["new_attr"] = ir.Attr("new_attr", ir.AttributeType.STRING, "value")
3. Value Manipulation
Replace All Uses of a Value
python
import onnx_ir.convenience as convenience
# Replace all uses of old_value with new_value
convenience.replace_all_uses_with(
old_value,
new_value,
replace_graph_outputs=True # Also replace in graph outputs if present
)
# Replace multiple values at once
convenience.replace_all_uses_with(
[old_value1, old_value2],
[new_value1, new_value2],
)
Check Value Usage
python
# Check if a value is used
if output_value.uses():
# Value has consumers
pass
# Check if value is a graph output
if value.is_graph_output():
# Special handling for outputs
pass
# Check if value is a graph input
if value.is_graph_input():
# Special handling for inputs
pass
Merge Value Information
python
# When eliminating nodes, preserve shape/type information
def merge_shapes(shape1: ir.Shape | None, shape2: ir.Shape | None) -> ir.Shape | None:
if shape1 is None:
return shape2
if shape2 is None:
return shape1
# More sophisticated merging logic...
return shape1
# Copy shape and type information
input_value.shape = merge_shapes(input_value.shape, output_value.shape)
if input_value.type is None:
input_value.type = output_value.type
4. Initializers and Constants
Work with Initializers
python
# Access initializers by name
initializers = graph.initializers
if "weight" in initializers:
weight_initializer = initializers["weight"]
# Register a new initializer
new_initializer = ir.Value(
name="new_weight",
shape=ir.Shape([3, 3, 64, 64]),
type=ir.TensorType(ir.DataType.FLOAT),
const_value=tensor_data,
)
graph.register_initializer(new_initializer)
# Remove unused initializers
graph_outputs = frozenset(graph.outputs)
graph_inputs = frozenset(graph.inputs)
for init in list(initializers.values()):
if not (init.uses() or init in graph_outputs or init in graph_inputs):
assert init.name is not None
del initializers[init.name]
Lift Constants to Initializers
python
# Check if node is a Constant
if node.op_type == "Constant" and node.domain in ("", "onnx.ai"):
# Get the tensor from the value attribute
attr_value = node.attributes.get("value")
if attr_value:
tensor = attr_value.as_tensor()
# Create initializer
initializer = ir.Value(
name=node.outputs[0].name,
shape=tensor.shape,
type=ir.TensorType(tensor.dtype),
const_value=tensor,
)
graph.register_initializer(initializer)
# Replace uses and remove node
node.outputs[0].replace_all_uses_with(initializer)
graph.remove(node, safe=True)
5. Logging and Debugging
python
import logging
logger = logging.getLogger(__name__)
class MyPass(ir.passes.InPlacePass):
def call(self, model: ir.Model) -> ir.passes.PassResult:
count = 0
for node in model.graph:
# Use debug for detailed information
logger.debug("Processing node: %s", node)
# Use info for important changes
logger.info("Removed node: %s", node.name)
count += 1
# Summarize at the end
if count:
logger.info("MyPass removed %s nodes", count)
return ir.passes.PassResult(model, modified=bool(count))
6. Handling Subgraphs
python
# Process attributes that may contain subgraphs
for attr in node.attributes.values():
if not isinstance(attr, ir.Attr):
continue
if attr.type == ir.AttributeType.GRAPH:
subgraph = attr.as_graph()
# Process the subgraph recursively
modified |= self._process_graph(subgraph)
elif attr.type == ir.AttributeType.GRAPHS:
for subgraph in attr.as_graphs():
# Process each subgraph
modified |= self._process_graph(subgraph)
7. Opset Version Handling
python
# Get opset version for the graph
onnx_opset_version = model.graph.opset_imports.get("", None)
# Check if a specific opset is available
if onnx_opset_version is not None and onnx_opset_version >= 13:
# Use features from opset 13+
pass
# Get schema information for a node
import onnx
try:
op_schema = onnx.defs.get_schema(
node.op_type,
onnx_opset_version,
domain=node.domain
)
# Use schema information
except Exception:
logger.info("Failed to get schema for %s", node)
Common Pass Patterns
Pattern 1: Node Elimination
Eliminate nodes that match certain criteria (e.g., Identity, unused nodes).
python
class NodeEliminationPass(ir.passes.InPlacePass):
def call(self, model: ir.Model) -> ir.passes.PassResult:
modified = False
for node in ir.traversal.RecursiveGraphIterator(model.graph):
if self._should_eliminate(node):
if self._try_eliminate_node(node):
modified = True
return ir.passes.PassResult(model, modified=modified)
def _should_eliminate(self, node: ir.Node) -> bool:
"""Check if node should be eliminated."""
return node.op_type == "Identity" and node.domain == ""
def _try_eliminate_node(self, node: ir.Node) -> bool:
"""Try to eliminate node. Returns True if successful."""
# Validate node structure
if len(node.inputs) != 1 or len(node.outputs) != 1:
return False
input_value = node.inputs[0]
output_value = node.outputs[0]
if input_value is None:
return False
# Replace uses
ir.convenience.replace_all_uses_with(
output_value, input_value, replace_graph_outputs=True
)
# Remove node
assert node.graph is not None
node.graph.remove(node, safe=True)
return True
Pattern 2: Dead Code Elimination
Remove unused nodes and values.
python
class DeadCodeEliminationPass(ir.passes.InPlacePass):
def call(self, model: ir.Model) -> ir.passes.PassResult:
count = self._remove_unused_nodes(model.graph)
for function in model.functions.values():
count += self._remove_unused_nodes(function)
return ir.passes.PassResult(model, modified=bool(count))
def _remove_unused_nodes(self, graph_like: ir.Graph | ir.Function) -> int:
"""Remove nodes that produce no used outputs."""
graph_outputs = frozenset(graph_like.outputs)
count = 0
# Iterate in reverse to handle dependencies
for node in reversed(graph_like):
removable = True
for output in node.outputs:
if output in graph_outputs or output.uses():
removable = False
break
if removable:
graph_like.remove(node, safe=True)
count += 1
return count
Pattern 3: Common Subexpression Elimination
Eliminate duplicate computations.
python
class CSEPass(ir.passes.InPlacePass):
def call(self, model: ir.Model) -> ir.passes.PassResult:
modified = self._eliminate_cse(model.graph)
return ir.passes.PassResult(model, modified=modified)