onnx/pass-writing

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

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

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

1. Preconditions (requires method): Check input model validity (optional)
2. Transformation (call method): Apply the transformation
3. Postconditions (ensures method): Validate output model (optional)

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)

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

# 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

# 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

# Always use safe=True to ensure proper cleanup
graph.remove(node, safe=True)

Create New Nodes

# 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

# 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

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

# 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

# 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

# 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

# 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

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

# 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

# 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).

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.

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.

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)

    def _eliminate_cse(self, graph: ir.Graph) -> bool:
        modified = False
        # Map from (op_identifier, inputs, attributes) to node
        existing_nodes: dict[tuple, ir.Node] = {}

        for node in graph:
            # Skip non-deterministic ops
            if self._is_non_deterministic(node):
                continue

            # Create a hashable key for the node
            node_key = (
                node.op_identifier(),
                tuple(id(inp) for inp in node.inputs),
                tuple(sorted(node.attributes.items())),
            )

            if node_key in existing_nodes:
                # Found duplicate - replace with existing
                existing_node = existing_nodes[node_key]
                ir.convenience.replace_all_uses_with(
                    node.outputs,
                    existing_node.outputs
                )
                graph.remove(node, safe=True)
                modified = True
            else:
                existing_nodes[node_key] = node

        return modified

    def _is_non_deterministic(self, node: ir.Node) -> bool:
        """Check if node is non-deterministic."""
        non_deterministic_ops = frozenset({
            "RandomUniform", "RandomNormal",
            "RandomUniformLike", "RandomNormalLike",
            "Multinomial"
        })
        return node.op_type in non_deterministic_ops and node.domain == ""

Pattern 4: Graph Normalization

Ensure graph is in a canonical form (e.g., topological sort, name fixing).

class TopologicalSortPass(ir.passes.InPlacePass):
    """Sort nodes in topological order."""

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        original_nodes = list(model.graph)
        model.graph.sort()  # Built-in method
        sorted_nodes = list(model.graph)

        # Check if order changed
        modified = False
        for node, new_node in zip(original_nodes, sorted_nodes):
            if node is not new_node:
                modified = True
                break

        # Also sort functions
        for function in model.functions.values():
            function.sort()

        return ir.passes.PassResult(model, modified=modified)

Pattern 5: Attribute/Metadata Manipulation

Modify node attributes or clear metadata.

class ClearMetadataPass(ir.passes.InPlacePass):
    """Clear metadata and doc strings from the model."""

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        modified = False

        # Clear model metadata
        if model.doc_string or model.metadata_props:
            model.doc_string = ""
            model.metadata_props.clear()
            modified = True

        # Clear graph metadata
        modified |= self._clear_graph_metadata(model.graph)

        # Clear function metadata
        for function in model.functions.values():
            modified |= self._clear_graph_metadata(function)

        return ir.passes.PassResult(model, modified=modified)

    def _clear_graph_metadata(self, graph_like: ir.Graph | ir.Function) -> bool:
        modified = False

        if graph_like.doc_string:
            graph_like.doc_string = ""
            modified = True

        for node in ir.traversal.RecursiveGraphIterator(graph_like):
            if node.doc_string or node.metadata_props:
                node.doc_string = ""
                node.metadata_props.clear()
                modified = True

        return modified

Testing Your Pass

Basic Test Structure

import onnx_ir as ir

def test_my_pass():
    # Create a test model
    model = create_test_model()

    # Apply the pass
    pass_instance = MyPass()
    result = pass_instance(model)

    # Verify the result
    assert result.modified == True
    assert len(result.model.graph) == expected_node_count

    # Verify specific transformations
    # ...

Common Pitfalls to Avoid

1. Modifying while iterating: ONNX IR's iterators are robust and support modification during iteration

   # Forward iteration with removal is safe in onnx_ir
   for node in graph:
       if should_remove(node):
           graph.remove(node, safe=True)
   
   # Reversed iteration is useful for dependency order
   for node in reversed(graph):
       if should_remove(node):
           graph.remove(node, safe=True)
   

   Note: Unlike standard Python iterators, onnx_ir's graph iterators are specifically designed to handle modifications during iteration. Choose forward or reverse iteration based on your algorithm's needs, not safety concerns.

2. Forgetting subgraphs: Always use RecursiveGraphIterator or manually process subgraphs

3. Not checking for None inputs: Nodes can have optional None inputs

   for input_value in node.inputs:
       if input_value is not None:  # Always check
           # Process input
           pass
   

4. Modifying graph outputs incorrectly: Be careful when replacing graph output values

   # Update graph outputs properly
   if output_value.is_graph_output():
       # Find and update in graph.outputs
       for idx, graph_output in enumerate(graph.outputs):
           if graph_output is output_value:
               graph.outputs[idx] = new_value
   

5. Not handling edge cases: Check for empty inputs/outputs, graph boundaries

6. Forgetting to process functions: Many passes should also process model.functions

Performance Considerations

1. Use in-place passes when possible (most efficient)

2. Minimize graph traversals: Combine multiple checks in one traversal

3. Use frozenset for lookups: When checking membership in graph inputs/outputs

   graph_outputs = frozenset(graph.outputs)
   if value in graph_outputs:  # O(1) lookup
       pass
   

4. Batch operations: Remove multiple nodes in one traversal rather than multiple passes

5. Early exit: Return early if no modifications are needed

Pass Composition

Sequential Pass Execution

# Combine multiple passes
passes = ir.passes.Sequential(
    RemoveUnusedNodesPass(),
    IdentityEliminationPass(),
    TopologicalSortPass(),
)

result = passes(model)

Pass Manager with Iteration

# Run passes multiple times until convergence
passes = ir.passes.PassManager(
    [
        CommonSubexpressionEliminationPass(),
        RemoveUnusedNodesPass(),
    ],
    steps=5,  # Maximum iterations
    early_stop=True,  # Stop if no changes
)

result = passes(model)

Security and Error Handling

1. Validate inputs: Check node structure, input/output counts
2. Handle exceptions gracefully: Catch and log errors in schema lookups
3. Use safe removal: Always call graph.remove(node, safe=True)
4. Avoid hardcoding assumptions: Check opset versions, schema availability
5. Log warnings: When skipping nodes due to unexpected conditions

try:
    schema = onnx.defs.get_schema(node.op_type, opset_version, domain=node.domain)
except Exception:
    logger.warning("Could not get schema for %s, skipping", node, exc_info=True)
    continue

Import Conventions

# Standard imports for pass files
from __future__ import annotations

import logging

import onnx_ir as ir

logger = logging.getLogger(__name__)

Summary

When creating a new pass:

1. Choose the right base class: InPlacePass (most common) or FunctionalPass
2. Implement call method: Main transformation logic
3. Return PassResult: With model and modified flag
4. Use RecursiveGraphIterator: To process all nodes including subgraphs
5. Process functions: Don't forget model.functions
6. Use convenience functions: replace_all_uses_with, etc.
7. Log appropriately: Use debug/info logging levels
8. Handle edge cases: None inputs, graph boundaries, subgraphs
9. Test thoroughly: Write tests for various scenarios
10. Document clearly: Explain what the pass does and any parameters

References

• Pass Infrastructure: src/onnx_ir/passes/_pass_infra.py
• Example Passes: src/onnx_ir/passes/common/
• Convenience Functions: src/onnx_ir/convenience.py
• Traversal Utilities: src/onnx_ir/traversal.py