Create a tutorial for IR passes

[justinchuby]

Here is an AI generated draft

Tutorial: IR Passes in ONNX IR Python

Introduction

The ONNX IR Python library provides a comprehensive pass framework for implementing graph transformations and optimizations. This tutorial covers the different base classes available and when to use each one.

Pass Base Classes

1. PassBase - The Abstract Foundation

PassBase is the abstract base class for all passes 1 . It defines the core contract that all passes must follow, including two key properties:

• in_place: Whether the pass modifies the model in place and returns it
• changes_input: Whether the pass modifies the input model 2

The relationship between these properties determines the pass behavior: 3

2. InPlacePass - For In-Place Modifications

InPlacePass is designed for passes that modify the input model directly and return the same model object 4 . This is the most commonly used base class in the codebase.

When to use InPlacePass:

• When you want to modify the existing model structure
• For performance-critical optimizations where copying the model is expensive
• When implementing transformations like node removal, attribute modification, or graph restructuring

Examples in the codebase:

• LiftConstantsToInitializersPass 5
• CommonSubexpressionEliminationPass 6
• ShapeInferencePass 7

3. FunctionalPass - For Pure Transformations

FunctionalPass creates a new model without modifying the input model 8 . This follows functional programming principles where transformations don't have side effects.

When to use FunctionalPass:

• When you need to preserve the original model
• For analysis passes that shouldn't modify the input
• When implementing transformations that benefit from immutable data structures
• When working in environments where the original model must remain unchanged

Core Implementation Pattern

The call() Method

Every pass must implement the call() method which takes a model and returns a PassResult 9 . The PassResult contains the transformed model and a boolean indicating whether modifications occurred 10 .

Preconditions and Postconditions

Passes can optionally implement requires() and ensures() methods for validation 11 . The framework automatically calls these methods and raises appropriate exceptions if conditions are violated 12 .

Advanced Pass Composition

Sequential Pass Execution

Sequential allows chaining multiple passes together 13 . The properties of the sequential pass are determined by the constituent passes 14 .

PassManager for Iterative Optimization

PassManager extends Sequential with additional features like multiple iterations and early stopping 15 . It runs passes repeatedly until no more changes occur or a maximum number of steps is reached 16 .

Implementation Examples

Example 1: Constant Lifting (InPlacePass)

The LiftConstantsToInitializersPass demonstrates a typical InPlacePass implementation. It converts Constant nodes to graph initializers, modifying the graph structure in place 17 .

Example 2: Common Subexpression Elimination (InPlacePass)

The CommonSubexpressionEliminationPass shows how to implement graph analysis and transformation. It identifies duplicate computations and removes redundant nodes 18 .

Example 3: Shape Inference (InPlacePass)

The ShapeInferencePass demonstrates integration with external tools (ONNX shape inference) while maintaining the pass interface 19 .

Error Handling

The pass framework provides a hierarchy of exception types for different error conditions:

• InvariantError: Base class for constraint violations 20
• PreconditionError: For requirement violations 21
• PostconditionError: For postcondition failures 22
• PassError: For general pass execution errors 23

Best Practices

1. Choose the right base class: Use InPlacePass for most transformations, FunctionalPass when preservation of the original model is required
2. Return accurate modification flags: Always set the modified flag correctly in PassResult to enable optimization frameworks to make informed decisions
3. Implement proper error handling: Use the provided exception types for clear error reporting
4. Consider pass composition: Use Sequential and PassManager to build complex optimization pipelines
5. Add logging: Use the logging framework for debugging and monitoring pass execution

Notes

The pass framework is inspired by PyTorch's pass infrastructure 24 and provides a robust foundation for implementing graph transformations. The framework handles many common concerns like input validation, property verification, and error handling, allowing pass developers to focus on the transformation logic.

While the codebase currently contains many examples of InPlacePass implementations, the FunctionalPass base class is available for use cases requiring immutable transformations. The choice between these depends on your specific requirements regarding model preservation and performance characteristics.