Refactors onnxscript.rewriter.onnxruntime.rewrite to call onnx.rewrit…

…er.rewrite (#1628) Signed-off-by: Xavier Dupre <[email protected]>
microsoft · Jun 17, 2024 · d620466 · d620466
1 parent dc31a6e
commit d620466
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 23 deletions.
diff --git a/docs/api/tools.md b/docs/api/tools.md
@@ -7,9 +7,9 @@
 ```
 
 ```{eval-rst}
-.. autofunction:: onnxscript.tools.transformers_models.phi.get_phi_model_config
+.. autofunction:: onnxscript.tools.transformers_models.phi.get_phi_model_from_config
 ```
 
 ```{eval-rst}
-.. autofunction:: onnxscript.tools.transformers_models.llama.get_llama_model_config
+.. autofunction:: onnxscript.tools.transformers_models.llama.get_llama_model_from_config
 ```
diff --git a/onnxscript/_legacy_ir/visitor.py b/onnxscript/_legacy_ir/visitor.py
@@ -590,6 +590,8 @@ def get_constant_value(i: int) -> onnx.TensorProto | None:
                     )
 
         for output in node.output:
+            if output == "":
+                continue
             info = self.lookup_or_create(output)
             if output in output_types:
                 if info.type is not None:

diff --git a/onnxscript/optimizer/remove_unused_ir.py b/onnxscript/optimizer/remove_unused_ir.py
@@ -32,8 +32,10 @@ def is_used_output(i: int) -> bool:
 
         if is_used_output(1) or is_used_output(2):
             return
-        node.outputs[1].name = ""
-        node.outputs[2].name = ""
+        if len(node.outputs) > 1:
+            node.outputs[1].name = ""
+        if len(node.outputs) > 2:
+            node.outputs[2].name = ""
         node.attributes.pop("training_mode", None)
         return
 

diff --git a/onnxscript/rewriter/__init__.py b/onnxscript/rewriter/__init__.py
@@ -40,7 +40,7 @@ def rewrite(
             pattern_rewrite_rules = pattern.RewriteRuleSet(pattern_rewrite_rules)
         count = pattern_rewrite_rules.apply_to_model(model_ir)
         print(f"Applied {count} of general pattern rewrite rules.")
+    remove_unused.remove_unused_nodes(model_ir)
+    model_ir = remove_unused_function.remove_unused_functions(model_ir)
     model = ir.serde.serialize_model(model_ir)
-    remove_unused.remove_unused_nodes(model)
-    model = remove_unused_function.remove_unused_functions(model)
     return model
diff --git a/onnxscript/rewriter/onnxruntime/__init__.py b/onnxscript/rewriter/onnxruntime/__init__.py
@@ -4,9 +4,8 @@
 
 import onnx
 
-from onnxscript import ir
-from onnxscript.optimizer import remove_unused, remove_unused_function
 from onnxscript.rewriter import function_rule, pattern
+from onnxscript.rewriter import rewrite as _rewrite
 from onnxscript.rewriter.onnxruntime import (
     group_normalization_merge_silu,
     instance_to_group_normalization,
@@ -44,18 +43,6 @@ def rewrite(
     """
     function_rules = function_rules or ORT_FUNCTION_REWRITE_RULES
     pattern_rules = pattern_rules or ORT_PATTERN_REWRITE_RULES
-    model = ir.serde.deserialize_model(model_proto)
-    # TODO(bowenbao): Function rules first, or pattern rules first?
-    if function_rules:
-        for rule_cls in function_rules:
-            count, model = rule_cls().apply_to_model(model)
-            if count > 0:
-                print(f"Applied {count} of onnxruntime specific function rewrite rules.")
-    if pattern_rules:
-        count = pattern.RewriteRuleSet(pattern_rules).apply_to_model(model)
-        print(f"Applied {count} of onnxruntime specific pattern rewrite rules.")
-
-    model_proto = ir.serde.serialize_model(model)
-    remove_unused.remove_unused_nodes(model_proto)
-    model_proto = remove_unused_function.remove_unused_functions(model_proto)
-    return model_proto
+    return _rewrite(
+        model_proto, function_rewrite_rules=function_rules, pattern_rewrite_rules=pattern_rules
+    )