Merge branch 'main' into titaiwang/add_bfloat16_offline_pass

microsoft · May 3, 2024 · 5c63bfa · 5c63bfa
2 parents e0cf5e3 + a1e46e4
commit 5c63bfa
Show file tree

Hide file tree

Showing 24 changed files with 91 additions and 150 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -26,6 +26,7 @@
     "sphinx.ext.ifconfig",
     "sphinx.ext.viewcode",
     "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
     "sphinx.ext.githubpages",
     "sphinx_gallery.gen_gallery",
     "sphinx.ext.autodoc",

diff --git a/docs/index.md b/docs/index.md
@@ -100,7 +100,6 @@ tutorial/index
 api/index
 intermediate_representation/index
 auto_examples/index
-rewriter/index
 articles/index
 ```
 

diff --git a/docs/test/test_documentation_examples.py b/docs/test/test_documentation_examples.py
@@ -50,7 +50,7 @@ def test(*relpath):
 
         test("..", "..", "docs", "examples")
         test("..", "..", "docs", "tutorial", "examples")
-        test("..", "..", "docs", "rewriter", "examples")
+        test("..", "..", "docs", "tutorial", "rewriter", "examples")
 
 
 if __name__ == "__main__":

diff --git a/docs/tutorial/index.md b/docs/tutorial/index.md
@@ -6,6 +6,13 @@ myst:
 
 # Tutorial
 
+## Welcome to ONNX Script Tutorials
+
+For extended tutorials on how to use the **Optimizer** and **Rewriter** tools, refer to the relevant sub-sections within the Tutorial section.
+- <project:optimizer/optimize.md>
+- <project:rewriter/rewrite_patterns.md>
+
+
 In this tutorial, we illustrate the features supported by {{ onnxscript }} using examples.
 
 ## Basic Features
@@ -257,3 +264,11 @@ ONNX perspective, the two assignments to *g* represent two distinct tensors
 
 ```{literalinclude} examples/outerscope_redef_error.py
 ```
+
+```{toctree}
+:maxdepth: 1
+
+optimizer/index
+rewriter/index
+```
+
diff --git a/docs/tutorial/optimizer/index.md b/docs/tutorial/optimizer/index.md
@@ -0,0 +1,5 @@
+# Optimizer Tutorials
+
+```{toctree}
+optimize
+```
diff --git a/docs/tutorial/optimizer/optimize.md b/docs/tutorial/optimizer/optimize.md
@@ -0,0 +1,53 @@
+# Optimizing a Model using the Optimizer
+
+## Introduction
+
+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.
+
+## Usage
+
+In order to utilize the optimizer tool,
+
+```python
+import onnxscript
+
+onnxscript.optimizer.optimize(model)
+```
+
+### optimize API
+The `onnxscript.optimizer.optimize` call takes in several optional parameters that allows the caller to further fine-tune the process of optimization.
+
+```{eval-rst}
+.. autofunction:: onnxscript.optimizer.optimize
+    :noindex:
+```
+
+## Description of optimizations applied by `onnxscript.optimizer.optimize`
+
+:::{table}
+:widths: auto
+:align: center
+
+| Optimization 'onnxscript.optimizer.` + .. | Description |
+| - | - |
+| **Constant folding** <br>`constant_folding.fold_constants` | Applies constant folding optimization to the model. |
+| **Constant propagation** <br>`constant_folding.fold_constants` | Applies constant propagation optimization to the model. Applied as part of the constant folding optimization. |
+| **Sequence simplification** <br>`constant_folding.fold_constants` | Simplifies Sequence based ops (SequenceConstruct, ConcatFromSequence) present in the model. Applied as part of the constant folding optimization. |
+| **Remove unused nodes** <br>`remove_unused.remove_unused_nodes` | Removes unused nodes from the model. |
+| **Remove unused functions** <br>`remove_unused_function.remove_unused_functions` | Removes unused function protos from the model. |
+| **Inline functions with unused outputs** <br>`simple_function_folding.inline_functions_with_unused_outputs` | Inlines function nodes that have unused outputs. |
+| **Inline simple functions** <br>`simple_function_folding.inline_simple_functions` | Inlines simple functions based on a node count threshold. |
+:::
+
+## List of pattern rewrite rules applied by `onnxscript.optimizer.optimize`
+
+```{eval-rst}
+.. autosummary::
+    :nosignatures:
+
+    onnxscript.rewriter.broadcast_to_matmul
+    onnxscript.rewriter.cast_constant_of_shape
+    onnxscript.rewriter.gemm_to_matmul_add
+    onnxscript.rewriter.no_op
+
+```
diff --git a/docs/rewriter/examples/broadcast_matmul.py → ...ial/rewriter/examples/broadcast_matmul.py b/docs/rewriter/examples/broadcast_matmul.py → ...ial/rewriter/examples/broadcast_matmul.py
diff --git a/docs/rewriter/examples/erfgelu.py → docs/tutorial/rewriter/examples/erfgelu.py b/docs/rewriter/examples/erfgelu.py → docs/tutorial/rewriter/examples/erfgelu.py
diff --git a/docs/rewriter/examples/img/broadcast_01.png → ...al/rewriter/examples/img/broadcast_01.png b/docs/rewriter/examples/img/broadcast_01.png → ...al/rewriter/examples/img/broadcast_01.png
diff --git a/docs/rewriter/examples/img/broadcast_02.png → ...al/rewriter/examples/img/broadcast_02.png b/docs/rewriter/examples/img/broadcast_02.png → ...al/rewriter/examples/img/broadcast_02.png
diff --git a/docs/rewriter/examples/img/erfgelu_01.png → ...rial/rewriter/examples/img/erfgelu_01.png b/docs/rewriter/examples/img/erfgelu_01.png → ...rial/rewriter/examples/img/erfgelu_01.png
diff --git a/docs/rewriter/examples/img/erfgelu_02.png → ...rial/rewriter/examples/img/erfgelu_02.png b/docs/rewriter/examples/img/erfgelu_02.png → ...rial/rewriter/examples/img/erfgelu_02.png
diff --git a/...riter/examples/img/erfgelu_03_commute.png → ...riter/examples/img/erfgelu_03_commute.png b/...riter/examples/img/erfgelu_03_commute.png → ...riter/examples/img/erfgelu_03_commute.png
diff --git a/...riter/examples/img/erfgelu_04_commute.png → ...riter/examples/img/erfgelu_04_commute.png b/...riter/examples/img/erfgelu_04_commute.png → ...riter/examples/img/erfgelu_04_commute.png
diff --git a/...riter/examples/img/erfgelu_05_commute.png → ...riter/examples/img/erfgelu_05_commute.png b/...riter/examples/img/erfgelu_05_commute.png → ...riter/examples/img/erfgelu_05_commute.png
diff --git a/...riter/examples/img/erfgelu_06_commute.png → ...riter/examples/img/erfgelu_06_commute.png b/...riter/examples/img/erfgelu_06_commute.png → ...riter/examples/img/erfgelu_06_commute.png
diff --git a/...riter/examples/img/erfgelu_07_commute.png → ...riter/examples/img/erfgelu_07_commute.png b/...riter/examples/img/erfgelu_07_commute.png → ...riter/examples/img/erfgelu_07_commute.png
diff --git a/docs/rewriter/index.md → docs/tutorial/rewriter/index.md b/docs/rewriter/index.md → docs/tutorial/rewriter/index.md
diff --git a/docs/rewriter/rewrite_patterns.md → docs/tutorial/rewriter/rewrite_patterns.md b/docs/rewriter/rewrite_patterns.md → docs/tutorial/rewriter/rewrite_patterns.md
@@ -28,10 +28,10 @@ We will show how we can find a subgraph matching this computation and replace it
 Firstly, include all the rewriter relevant imports.
 
 ```python
-    from onnxscript.rewriter import pattern
-    from onnxscript import ir
+from onnxscript.rewriter import pattern
+from onnxscript import ir
 
-    _op = pattern.onnxop
+_op = pattern.onnxop
 ```
 
 Then create a target pattern that needs to be replaced using onnxscript operators.
@@ -55,10 +55,10 @@ The inputs to the replacement pattern are of type `ir.Value`. For detailed usage
 For this example, we do not require a `match_condition` so that option is skipped for now. Then the rewrite rule is created using the `RewriteRule` function.
 
 ```python
-    rule = pattern.RewriteRule(
-        erf_gelu_pattern,  # Target Pattern
-        gelu,  # Replacement Pattern
-    )
+rule = pattern.RewriteRule(
+    erf_gelu_pattern,  # Target Pattern
+    gelu,  # Replacement Pattern
+)
 ```
 
 Now that the rewrite rule has been created, the next step is to apply these pattern-based rewrite rules. The `rewriter.rewrite` call consists of three main components:
@@ -117,8 +117,8 @@ Only one of the patterns has been successfully matched and replaced by a `GELU`
 This method requires creating two separate rules and packing them into either a sequence of `PatternRewriteRule`s or a `RewriteRuleSet`. Creating a `RewriteRuleSet` is the preferable option but either can be used. In order to create a `RewriteRuleSet` with multiple rules `rule1` and `rule2` for example:
 
 ```python
-    from onnxscript.rewriter import pattern
-    rewrite_rule_set = pattern.RewriteRuleSet(rules=[rule1, rule2])
+from onnxscript.rewriter import pattern
+rewrite_rule_set = pattern.RewriteRuleSet(rules=[rule1, rule2])
 ```
 
 In order to apply this method to the example above, first create the two separate target patterns as follows:
@@ -171,7 +171,7 @@ First, write a target pattern and replacement pattern in a similar way to the fi
 ```
 
 ```{literalinclude} examples/broadcast_matmul.py
-:pyobject: matmul
+:pyobject: matmul_pattern
 ```
 
 :::{note}

diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py
@@ -6,10 +6,6 @@
 
 from onnxscript import rewriter
 from onnxscript.optimizer.constant_folding import fold_constants
-from onnxscript.optimizer.copy_propagation import (
-    do_copy_propagation,
-    do_sequence_simplification,
-)
 from onnxscript.optimizer.remove_unused import remove_unused_nodes
 from onnxscript.optimizer.remove_unused_function import remove_unused_functions
 from onnxscript.optimizer.simple_function_folding import (
@@ -108,14 +104,11 @@ def optimize(
                 node.name,
             )
 
-    # do_sequence_simplification(model)
     return model
 
 
 __all__ = [
     "fold_constants",
     "remove_unused_nodes",
     "optimize",
-    "do_copy_propagation",
-    "do_sequence_simplification",
 ]
diff --git a/onnxscript/optimizer/constant_folding.py b/onnxscript/optimizer/constant_folding.py
@@ -263,7 +263,10 @@ def fold_constants(
     *,
     onnx_shape_inference: bool = False,
 ) -> bool:
-    """Returns true iff the model was modified."""
+    """
+    Applies constant folding optimization to the model.
+    Returns true iff the model was modified.
+    """
     folder = ConstantFolder(
         evaluator.registry,
         external_data_folder,

diff --git a/onnxscript/optimizer/copy_propagation.py b/onnxscript/optimizer/copy_propagation.py
diff --git a/onnxscript/optimizer/copy_propagation_test.py b/onnxscript/optimizer/copy_propagation_test.py
diff --git a/onnxscript/optimizer/simple_function_folding.py b/onnxscript/optimizer/simple_function_folding.py
@@ -200,6 +200,7 @@ def function_with_unused_outputs(self) -> dict[ir.FunctionId, onnx.FunctionProto
 
 
 def inline_simple_functions(model: onnx.ModelProto, node_count: int = 2) -> bool:
+    """Inlines simple functions based on a node count threshold"""
     inliner = FunctionInliner(node_count)
     inliner.visit_model(model)
     logger.info(
@@ -218,6 +219,7 @@ def inline_simple_functions(model: onnx.ModelProto, node_count: int = 2) -> bool
 
 
 def inline_functions_with_unused_outputs(model: onnx.ModelProto) -> bool:
+    """Inlines function nodes that have unused outputs."""
     # TODO: Use onnx.inliner after 1.16.
     # This visitor based inliner is used to ensure the function inner value info remains consistent.
     visitor = FindFunctionWithUnusedOutputsVisitor()