Skip to content

Commit

Permalink
[docs] Update README.md for brief descriptions of optimizer and rewri…
Browse files Browse the repository at this point in the history
…ter tools. (#1702)

#TODO Add more details about function-based rewriting once tutorial for
function-based rewriting is merged
  • Loading branch information
shubhambhokare1 authored Jul 17, 2024
1 parent d27aede commit 9ced95d
Showing 1 changed file with 89 additions and 0 deletions.
89 changes: 89 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,16 @@ models using a subset of Python. ONNX Script is:
* **Debuggable:** allows for eager-mode evaluation that provides for a
more delightful ONNX model debugging experience.

This repo also covers:

* **ONNX IR:** an in-memory IR that supports the full ONNX spec, designed
for graph construction, analysis and transformation.
* **ONNX Script Optimizer:** provides functionality to optimize an ONNX
model by performing optimizations and clean-ups such as constant folding,
dead code elimination, etc.
* **ONNX Rewriter:** provides functionality to replace certain patterns in
an ONNX graph with replacement patterns based on user-defined rewrite rules.

Note however that ONNX Script does **not** intend to support the entirety
of the Python language.

Expand Down Expand Up @@ -142,6 +152,85 @@ result = Hardmax(v)

More examples can be found in the [docs/examples](docs/examples) directory.

## ONNX IR

An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.

### Features

* **Full ONNX spec support:** all valid models representable by ONNX protobuf,
and a subset of invalid models (so you can load and fix them).
* **Low memory footprint:** mmap'ed external tensors; unified interface for
ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size
limitation. Zero copies.
* **Straightforward access patterns:** Access value information and traverse the
graph topology at ease.
* **Robust mutation:** Create as many iterators as you like on the graph while mutating it.
* **Speed:** Performant graph manipulation, serialization/deserialization to Protobuf.
* **Pythonic and familiar APIs:** Classes define Pythonic apis and still map to
ONNX protobuf concepts in an intuitive way.

## ONNX Script Tools

### ONNX Optimizer

The ONNX Script Optimizer tool provides the user with the functionality to optimize an ONNX model by performing optimizations and clean-ups such as constant folding, dead code elimination, etc. In order to utilize the optimizer tool:

```python
import onnxscript

onnxscript.optimizer.optimize(onnx_model)
```

For a detailed summary of all the optimizations applied by the optimizer call, refer to the tutorial [Optimizing a Model using the Optimizer](https://onnxscript.ai/tutorial/optimizer/optimize.html)

### ONNX Rewriter

The ONNX Rewriter tool provides the user with the functionality to replace certain patterns in an ONNX graph with another pattern based on user-defined rewrite rules. The rewriter tools allows two different methods in which patterns in the graph can be rewritten.

### Pattern-based rewriting

For this style of rewriting, the user provides a `target_pattern` that is to be replaced, a `replacement_pattern` and a `match_condition` (pattern rewrite will occur only if the match condition is satisfied). A simple example on how to use the pattern-based rewriting tool is as follows:

```python
from onnxscript.rewriter import pattern

# The target pattern
def erf_gelu_pattern(op, x):
return 0.5 * (x * (op.Erf(x / math.sqrt(2)) + 1.0))

def erf_gelu_pattern_2(op, x):
return (x * (op.Erf(x / math.sqrt(2)) + 1.0)) * 0.5

# The replacement pattern
def gelu(op, x: ir.Value):
return op.Gelu(x, domain="com.microsoft")

# Create multiple rules
rule1 = pattern.RewriteRule(
erf_gelu_pattern, # Target Pattern
gelu, # Replacement
)
rule2 = pattern.RewriteRule(
erf_gelu_pattern_2, # Target Pattern
gelu, # Replacement
)
# Create a Rewrite Rule Set with multiple rules.
rewrite_rule_set = pattern.RewriteRuleSet([rule1, rule2])
# Apply rewrites
model_with_rewrite_applied = onnxscript.rewriter.rewrite(
model, # Original ONNX Model
pattern_rewrite_rules=rewrite_rule_set,
)
return model_with_rewrite_applied
```

For a detailed tutorial on how to create target_pattern, replacement_pattern and match_condition blocks in order to utilize the pattern-based rewriter, refer to the tutorial [Pattern-based Rewrite Using Rules](https://onnxscript.ai/tutorial/rewriter/rewrite_patterns.html)

### Function-based rewriting

This style of rewriting matches a `FUNCTION_KEYWORD` and `PACKAGE_NAME` provided by the user to an existing function within the graph and replaces it with a new function provided by the user.

## Development Guidelines

Every change impacting the converter or the eager evaluation must be
Expand Down

0 comments on commit 9ced95d

Please sign in to comment.