Merge branch 'main' into titaiwang/add_constant_to_initilizer_pass

Author: titaiwangms

onnxscript/_internal/version_utils.py

@@ -43,26 +43,6 @@ def transformers_older_than(version: str) -> bool | None:
     )
 
 
-def is_onnxruntime_training() -> bool:
-    """Returns True if the onnxruntime is onnxruntime-training."""
-    try:
-        from onnxruntime import training  # pylint: disable=import-outside-toplevel
-
-        assert training
-    except ImportError:
-        # onnxruntime not training
-        return False
-
-    try:
-        from onnxruntime.capi.onnxruntime_pybind11_state import (  # pylint: disable=import-outside-toplevel
-            OrtValueVector,
-        )
-    except ImportError:
-        return False
-
-    return hasattr(OrtValueVector, "push_back_batch")
-
-
 def onnxruntime_older_than(version: str) -> bool:
     """Returns True if the onnxruntime version is older than the given version."""
     import onnxruntime  # pylint: disable=import-outside-toplevel

onnxscript/optimizer/__init__.py

@@ -35,7 +35,9 @@ def optimize(model: ir.Model, *args, **kwargs) -> ir.Model:
         return legacy_optimizer.optimize(model, *args, **kwargs)
 
 
-def fold_constants(model: ir.Model | onnx.ModelProto, *args, **kwargs) -> bool:
+def fold_constants(
+    model: ir.Model | onnx.ModelProto, *args, **kwargs
+) -> constant_folding.FoldConstantsResult | bool:
     if isinstance(model, ir.Model):
         return constant_folding.fold_constants(model, *args, **kwargs)
     else:

onnxscript/optimizer/_constant_folding.py

@@ -148,18 +148,24 @@ class Replacement:
 # Currently, we assume that symbolic dimensions are also guaranteed to be non-negative.
 # TODO: Add support for negative symbolic dimensions.
 
+SymbolicValue = Union[ir.Value, list[ir.Value], ir.Shape]
+
 
 class OptimizerState:
     def __init__(self):
-        self._sym_value_map: dict[ir.Value, Any] = {}
+        self._sym_value_map: dict[ir.Value, SymbolicValue] = {}
         self._initializer_inputs: list[set[ir.Value]] = []
 
-    def get_sym_value(self, value: ir.Value | None) -> Any:
+    @property
+    def symbolic_value_map(self) -> dict[ir.Value, SymbolicValue]:
+        return self._sym_value_map
+
+    def get_sym_value(self, value: ir.Value | None) -> SymbolicValue | None:
         if value is None:
             return None
         return self._sym_value_map.get(value)
 
-    def set_sym_value(self, value: ir.Value, sym_value: Any) -> None:
+    def set_sym_value(self, value: ir.Value, sym_value: SymbolicValue) -> None:
         self._sym_value_map[value] = sym_value
 
     def push_initializer_inputs(self) -> None:
@@ -1094,7 +1100,17 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
         for function in model.functions.values():
             # TODO(rama): Should we specialize functions?
             self.visit_function(function)
-        return ir.passes.PassResult(model, self.modified)
+        return FoldConstantsResult(model, self.modified, self._state.symbolic_value_map)
+
+
+@dataclasses.dataclass
+class FoldConstantsResult(ir.passes.PassResult):
+    symbolic_value_map: dict[ir.Value, SymbolicValue]
+
+    # Add conversion to bool for backward compatibility. The previously returned value
+    # for the fold_constants method was a boolean indicating whether the model was modified.
+    def __bool__(self) -> bool:
+        return self.modified
 
 
 def fold_constants(
@@ -1104,23 +1120,31 @@ def fold_constants(
     onnx_shape_inference: bool = False,
     input_size_limit: int = DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT,
     output_size_limit: int = DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT,
-) -> bool:
+) -> FoldConstantsResult:
     """
     Applies constant folding optimization to the model.
-    Returns true iff the model was modified.
+
+    Args:
+        model: The ONNX model to optimize.
+        external_data_folder: Path to the folder containing external data
+            for the model. Defaults to an empty string.
+        onnx_shape_inference: Whether to enable ONNX shape inference during
+            constant folding. Defaults to False.
+        input_size_limit: The maximum size (in bytes) of input tensors
+            that can be considered for constant folding. Defaults to
+            `DEFAULT_CONSTANT_FOLD_INPUT_SIZE_LIMIT`.
+        output_size_limit: The maximum size (in bytes) of output tensors
+            that can be stored after constant folding. Defaults to
+            `DEFAULT_CONSTANT_FOLD_OUTPUT_SIZE_LIMIT`.
+
+    Returns:
+        An instance of `FoldConstantsResult`.
+
     """
     folder_pass = FoldConstantsPass(
         external_data_folder=external_data_folder,
         shape_inference=onnx_shape_inference,
         input_size_limit=input_size_limit,
         output_size_limit=output_size_limit,
     )
-    folder_pass(model)
-    for op in folder_pass.counts:
-        logger.info(
-            "Constant-folded '%s' %s times, with %s size.",
-            op,
-            folder_pass.counts[op],
-            folder_pass.sizes[op],
-        )
-    return folder_pass.modified
+    return folder_pass(model)  # type: ignore[return-value]

onnxscript/optimizer/_inliner.py

@@ -190,6 +190,7 @@ def id_abbreviation(id: ir.OperatorIdentifier) -> str:
 
 class InlinePass(ir.passes.InPlacePass):
     def __init__(self) -> None:
+        super().__init__()
         self._functions: dict[ir.OperatorIdentifier, ir.Function] = {}
         self._function_id_abbreviations: dict[ir.OperatorIdentifier, str] = {}
         self._opset_imports: dict[str, int] = {}

onnxscript/optimizer/_legacy/_optimizer.py

@@ -15,7 +15,6 @@
     inline_simple_functions,
 )
 from onnxscript.optimizer._legacy.constant_folding import fold_constants
-from onnxscript.optimizer._optimizer import _DEFAULT_REWRITE_RULES
 
 logger = logging.getLogger(__name__)
 
@@ -75,7 +74,7 @@ def optimize(
         onnxscript.optimizer.remove_unused_functions(model)
         inline_functions_with_unused_outputs(model)
         # NOTE: This is general rewrite rules
-        model = rewriter.rewrite(model, pattern_rewrite_rules=_DEFAULT_REWRITE_RULES)
+        model = rewriter.rewrite(model)
         if stop_if_no_change and not modified:
             logger.debug("Stopping after %d iterations.", _)
             break

onnxscript/optimizer/_optimizer.py

@@ -6,29 +6,11 @@
 
 import onnxscript.ir.passes.common.constant_manipulation
 import onnxscript.ir.passes.common.unused_removal
-import onnxscript.optimizer
 from onnxscript import ir, rewriter
 from onnxscript.optimizer import _constant_folding, _inliner
-from onnxscript.rewriter import (
-    broadcast_to_matmul,
-    cast_constant_of_shape,
-    collapse_slices,
-    gemm_to_matmul_add,
-    llama_rule_sets,
-    no_op,
-)
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_REWRITE_RULES: tuple[rewriter.pattern.RewriteRule, ...] = (
-    *no_op.rules.rules,  # TODO: merge this rule into constant folding?
-    *broadcast_to_matmul.rules.rules,
-    gemm_to_matmul_add.rule,  # type: ignore[has-type]
-    *cast_constant_of_shape.rules.rules,
-    *collapse_slices.rules.rules,
-    *llama_rule_sets.llama_p0_rule_set().rules,
-)
-
 
 def optimize_ir(
     model: ir.Model,
@@ -62,7 +44,7 @@ def optimize_ir(
                     input_size_limit=input_size_limit,
                     output_size_limit=output_size_limit,
                 ),
-                rewriter.RewritePass(_DEFAULT_REWRITE_RULES),
+                rewriter.RewritePass(rewriter._DEFAULT_REWRITE_RULES),
                 onnxscript.ir.passes.common.unused_removal.RemoveUnusedNodesPass(),
                 onnxscript.ir.passes.common.unused_removal.RemoveUnusedFunctionsPass(),
                 onnxscript.ir.passes.common.unused_removal.RemoveUnusedOpsetsPass(),

onnxscript/rewriter/__init__.py

@@ -5,46 +5,81 @@
 from typing import Sequence, TypeVar, Union
 
 __all__ = [
-    # Modules
     "pattern",
-    # Functions
     "rewrite",
+    "RewritePass",
 ]
 
 import onnx
 
 from onnxscript import ir
 from onnxscript.ir.passes.common import unused_removal
-from onnxscript.rewriter import pattern
+from onnxscript.rewriter import (
+    broadcast_to_matmul,
+    cast_constant_of_shape,
+    collapse_slices,
+    gemm_to_matmul_add,
+    llama_rule_sets,
+    no_op,
+    pattern,
+)
 
-PatternRewriteRule = pattern.RewriteRule
-
-ModelProtoOrIr = TypeVar("ModelProtoOrIr", onnx.ModelProto, ir.Model)
+_ModelProtoOrIr = TypeVar("_ModelProtoOrIr", onnx.ModelProto, ir.Model)
+_DEFAULT_REWRITE_RULES: tuple[pattern.RewriteRule, ...] = (
+    *no_op.rules.rules,  # TODO: merge this rule into constant folding?
+    *broadcast_to_matmul.rules.rules,
+    gemm_to_matmul_add.rule,  # type: ignore[has-type]
+    *cast_constant_of_shape.rules.rules,
+    *collapse_slices.rules.rules,
+    *llama_rule_sets.llama_p0_rule_set().rules,
+)
 
 
 class RewritePass(ir.passes.InPlacePass):
     def __init__(
         self,
-        pattern_rewrite_rules: Sequence[PatternRewriteRule] | pattern.RewriteRuleSet = (),
+        rules: Sequence[pattern.RewriteRule] | pattern.RewriteRuleSet,
+        /,
     ) -> None:
-        if pattern_rewrite_rules:
-            if not isinstance(pattern_rewrite_rules, pattern.RewriteRuleSet):
-                # Create a pattern rule-set using provided rules
-                pattern_rewrite_rules = pattern.RewriteRuleSet(pattern_rewrite_rules)
-        assert isinstance(pattern_rewrite_rules, pattern.RewriteRuleSet)
-        self.pattern_rewrite_rules: pattern.RewriteRuleSet = pattern_rewrite_rules
+        super().__init__()
+        if isinstance(rules, Sequence):
+            if not rules:
+                raise ValueError("rules must not be empty")
+            # Create a pattern rule-set using provided rules
+            rules = pattern.RewriteRuleSet(rules)
+        assert isinstance(rules, pattern.RewriteRuleSet)
+        self.rules: pattern.RewriteRuleSet = rules
 
     def call(self, model: ir.Model) -> ir.passes.PassResult:
-        count = self.pattern_rewrite_rules.apply_to_model(model)
+        count = self.rules.apply_to_model(model)
         if count:
             print(f"Applied {count} of general pattern rewrite rules.")
         return ir.passes.PassResult(model, bool(count))
 
 
 def rewrite(
-    model: ModelProtoOrIr,
-    pattern_rewrite_rules: Union[Sequence[PatternRewriteRule], pattern.RewriteRuleSet] = (),
-) -> ModelProtoOrIr:
+    model: _ModelProtoOrIr,
+    pattern_rewrite_rules: Union[Sequence[pattern.RewriteRule], pattern.RewriteRuleSet]
+    | None = None,
+) -> _ModelProtoOrIr:
+    """Rewrite the model using the provided pattern rewrite rules.
+
+    Unused nodes, functions, and opsets will be removed after the rewrite.
+
+    Args:
+        model: The model to be rewritten. Can be an ONNX ModelProto or an ir.Model.
+        pattern_rewrite_rules: A sequence of pattern rewrite rules or a RewriteRuleSet.
+            If not provided, default rules will be applied. If empty, no rules will be applied
+            and the original model will be returned.
+
+    Returns:
+        The rewritten model as the same type as the input model.
+    """
+    if pattern_rewrite_rules is None:
+        pattern_rewrite_rules = _DEFAULT_REWRITE_RULES
+    elif not pattern_rewrite_rules:
+        return model
+
     if isinstance(model, onnx.ModelProto):
         model_ir = ir.serde.deserialize_model(model)
         proto = True

onnxscript/rewriter/pattern.py

@@ -1664,13 +1664,18 @@ def _get_new_overload(model: ir.Model, domain: str, name: str) -> str:
 
 class RewriteRuleSet:
     def __init__(self, rules: Sequence[RewriteRule], *, commute: bool = False) -> None:
+        if not rules:
+            raise ValueError("rules must contain at least one rule")
         if commute:
             rules = list(itertools.chain.from_iterable([rule.commute() for rule in rules]))
         self.rules = rules
         # We call remove_unused_nodes at end of rewriting if there is any rule that does
         # NOT remove nodes (immediately when it is applied)
         self.remove_unused_nodes = any(not rule.remove_nodes for rule in rules)
 
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.rules})"
+
     def _apply_to_graph_or_function(
         self,
         model: ir.Model,