Merge branch 'main' into justinchu/node-graph

Author: justinchuby

docs/tutorial/rewriter/rewrite_patterns.md

@@ -152,6 +152,15 @@ In order to apply this method to the example above, first create the two separat
 :pyobject: erf_gelu_pattern_2
 ```
 
+:::{note}
+:name: rule-application-order-matters
+
+When you pass multiple rules in `pattern_rewrite_rules`, the **order in which they appear is important**.
+This is because some rules may depend on patterns created or modified by earlier rules. For example, if `rule2` can only match after `rule1` has made a specific change in the model, then `rule1` must come **before** `rule2` in the list.
+If you're not seeing expected results, try adjusting the order or applying the rule set in a loop until no more changes occur.
+:::
+
+
 Then, create two separate `PatternRewriteRule`s, one for each target pattern. Pack these rules into a `RewriteRuleSet` object and apply rewrites by passing the created `RewriteRuleSet` for the `pattern_rewrite_rules` parameter.
 
 ```{literalinclude} examples/erfgelu.py

onnxscript/ir/passes/_pass_infra.py

@@ -136,6 +136,20 @@ def __call__(self, model: ir.Model) -> PassResult:
                 f"The result of the pass '{self.__class__.__name__}' should be type PassResult. "
                 "Please create one with ir.passes.PassResult()."
             )
+
+        # Checks that the declared in-place property is respected
+        if self.in_place and result.model is not model:
+            raise PassError(
+                f"The pass '{self.__class__.__name__}' is declared in-place, "
+                "but the model returned is *not* the same object as the input model. "
+                "Pass developer: Pass should return the same model object or the in_place property should return False."
+            )
+        if not self.in_place and result.model is model:
+            raise PassError(
+                f"The pass '{self.__class__.__name__}' is declared not in-place, "
+                "but the model returned *is* the same object as the input model. "
+                "Pass developer: Pass should return a new model object or the in_place property should return True."
+            )
         return result
 
     @abc.abstractmethod

onnxscript/ir/passes/common/constant_manipulation.py

@@ -23,11 +23,14 @@ class LiftConstantsToInitializersPass(ir.passes.InPlacePass):
     Attributes:
         lift_all_constants: Whether to lift all Constant nodes, including those that does not contain a tensor attribute (e.g. with value_ints etc.)
             Default to False, where only Constants with the ``value`` attribute are lifted.
+        size_limit: The minimum size of the tensor to be lifted. If the tensor contains
+            number of elements less than size_limit, it will not be lifted. Default is 16.
     """
 
-    def __init__(self, lift_all_constants: bool = False):
+    def __init__(self, lift_all_constants: bool = False, size_limit: int = 16):
         super().__init__()
-        self._lift_all_constants = lift_all_constants
+        self.lift_all_constants = lift_all_constants
+        self.size_limit = size_limit
 
     def call(self, model: ir.Model) -> ir.passes.PassResult:
         count = 0
@@ -79,16 +82,17 @@ def call(self, model: ir.Model) -> ir.passes.PassResult:
 
     def _constant_node_attribute_to_tensor(
         self, node, attr_name: str, attr_value: ir.Attr, initializer_name: str
-    ) -> ir.Tensor | None:
+    ) -> ir.TensorProtocol | None:
         """Convert constant node attribute to tensor."""
-        if not self._lift_all_constants and attr_name != "value":
+        if not self.lift_all_constants and attr_name != "value":
             logger.debug(
                 "Constant node '%s' has non-tensor attribute '%s'", node.name, attr_name
             )
             return None
 
+        tensor: ir.TensorProtocol
         if attr_name == "value":
-            tensor = attr_value.as_tensor()  # type: ignore[union-attr]
+            tensor = attr_value.as_tensor()
         elif attr_name == "value_int":
             tensor = ir.tensor(
                 attr_value.as_int(), dtype=ir.DataType.INT64, name=initializer_name
@@ -110,5 +114,15 @@ def _constant_node_attribute_to_tensor(
                 np.array(attr_value.value, dtype=np.bytes_), name=initializer_name
             )
         else:
-            tensor = None
-        return tensor  # type: ignore[return-value]
+            raise ValueError(
+                f"Unsupported constant node '{node.name}' attribute '{attr_name}'"
+            )
+
+        if tensor.size < self.size_limit:
+            logger.debug(
+                "Tensor from node '%s' has less than %s elements",
+                node.name,
+                self.size_limit,
+            )
+            return None
+        return tensor

onnxscript/ir/passes/common/constant_manipulation_test.py

@@ -56,7 +56,7 @@ def test_pass_with_lifting_float_and_int_constants_to_initializers(
 
         # Perform lift constants to initializers
         result = constant_manipulation.LiftConstantsToInitializersPass(
-            lift_all_constants=lift_all_constants
+            lift_all_constants=lift_all_constants, size_limit=0
         )(model)
         self.assertTrue(result.modified)
         # Check that the constant node is lifted to an initializer
@@ -130,7 +130,7 @@ def test_pass_with_lifting_constants_to_initializers_within_subgraph(
             ir_version=10,
         )
         result = constant_manipulation.LiftConstantsToInitializersPass(
-            lift_all_constants=lift_all_constants
+            lift_all_constants=lift_all_constants, size_limit=0
         )(model)
         self.assertTrue(result.modified)
         # Check that the constant node is lifted to the subgraph initializers
@@ -206,7 +206,7 @@ def test_pass_with_lifting_constants_to_initializers_with_floats_ints_strings(
 
         # Perform lift constants to initializers
         result = constant_manipulation.LiftConstantsToInitializersPass(
-            lift_all_constants=lift_all_constants
+            lift_all_constants=lift_all_constants, size_limit=0
         )(model)
         if lift_all_constants:
             self.assertTrue(result.modified)
@@ -249,3 +249,7 @@ def test_not_lifting_constants_to_initializers_when_it_is_output(self):
         self.assertFalse(result.modified)
         # Check that the constant node is not lifted to an initializer
         self.assertEqual(len(result.model.graph.initializers), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

onnxscript/ir/passes/common/topological_sort.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Pass for topologically sorting the graphs."""
+
+from __future__ import annotations
+
+__all__ = [
+    "TopologicalSortPass",
+]
+
+
+from onnxscript import ir
+
+
+class TopologicalSortPass(ir.passes.InPlacePass):
+    """Topologically sort graphs and functions in a model."""
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        original_nodes = list(model.graph)
+        model.graph.sort()
+        sorted_nodes = list(model.graph)
+        for function in model.functions.values():
+            original_nodes.extend(function)
+            function.sort()
+            sorted_nodes.extend(function)
+
+        # Compare node orders to determine if any changes were made
+        modified = False
+        for node, new_node in zip(original_nodes, sorted_nodes):
+            if node is not new_node:
+                modified = True
+                break
+        return ir.passes.PassResult(model=model, modified=modified)

onnxscript/ir/passes/common/topological_sort_test.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Unit tests for the TopologicalSortPass."""
+
+import unittest
+
+from onnxscript import ir
+from onnxscript.ir.passes.common import topological_sort
+
+
+class TopologicalSortPassTest(unittest.TestCase):
+    def setUp(self):
+        self.node_a = ir.node("A", inputs=[], name="node_a")
+        self.node_b = ir.node("B", inputs=self.node_a.outputs, name="node_b")
+        self.node_c = ir.node("C", inputs=self.node_b.outputs, name="node_c")
+
+    def test_topological_sort_modified_true(self):
+        graph = ir.Graph(
+            inputs=self.node_a.inputs,
+            outputs=self.node_c.outputs,
+            nodes=[self.node_c, self.node_b, self.node_a],  # Unsorted nodes
+            name="test_graph",
+        )
+        model = ir.Model(graph, ir_version=10)
+        result = topological_sort.TopologicalSortPass()(model)
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            tuple(result.model.graph),
+            (self.node_a, self.node_b, self.node_c),
+        )
+
+    def test_topological_sort_modified_false(self):
+        """Test that modified is False when the input model is already sorted."""
+        sorted_graph = ir.Graph(
+            inputs=self.node_a.inputs,
+            outputs=self.node_c.outputs,
+            nodes=[self.node_a, self.node_b, self.node_c],  # Sorted nodes
+            name="test_graph",
+        )
+        sorted_model = ir.Model(sorted_graph, ir_version=10)
+        result = topological_sort.TopologicalSortPass()(sorted_model)
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            tuple(result.model.graph),
+            (self.node_a, self.node_b, self.node_c),
+        )
+
+    def test_topological_sort_on_functions(self):
+        """Test that TopologicalSortPass works on functions in a model."""
+        # Create a function with unsorted nodes
+        func_graph = ir.Graph(
+            inputs=self.node_a.inputs,
+            outputs=self.node_c.outputs,
+            nodes=[self.node_c, self.node_b, self.node_a],  # Unsorted nodes
+        )
+        function = ir.Function(
+            domain="test_domain",
+            name="test_function",
+            graph=func_graph,
+            attributes=[],
+        )
+
+        # Create a model with the function
+        graph = ir.Graph(
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            name="test_graph",
+        )
+        model = ir.Model(graph, ir_version=10, functions=[function])
+
+        # Apply the TopologicalSortPass
+        result = topological_sort.TopologicalSortPass()(model)
+
+        # Verify that the nodes in the function are sorted
+        sorted_func_nodes = (self.node_a, self.node_b, self.node_c)
+        self.assertTrue(result.modified)
+        self.assertEqual(
+            tuple(result.model.functions[function.identifier()]),
+            sorted_func_nodes,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

requirements/ci/requirements-ort-nightly.txt

@@ -1,3 +1,3 @@
 # https://aiinfra.visualstudio.com/PublicPackages/_artifacts/feed/ORT-Nightly/PyPI/onnxruntime/overview
 --index-url=https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-onnxruntime==1.22.0.dev20250303002
+onnxruntime==1.22.0.dev20250402004