Change loop order during rewrite (#2427)

Change the loop order when applying a collection of rewrite-rules to all
nodes in graph. The order in this PR is preferable, for a couple of
reasons.

First: rules often have mutual dependences, with one rule enabling
another. This loop-ordering handles such dependences better, and does so
more efficiently (requiring fewer iterations of yet another outer loop
to invoke the rewriter multiple times).

It also sets it up for another optimization planned for originally (but
not yet implemented): For patterns that end with a definite op (like
ScatterND), which are most of the rules/patterns, we can create a
dictionary mapping the op-identifier to rules applicable to that op.
Then, it is sufficient to iterate only over rules applicable to current
node's op.

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/rewriter/_rewrite_rule.py

@@ -465,12 +465,12 @@ def _apply_to_graph_or_function(
         """
         count = 0
 
-        # NOTE: Rules should be prioritized in the order they are added to the RewriteRuleSet.
-        # And the graph is applied in order.
         for rule in self.rules:
             if rule.graph_pre_visitor:
                 rule.graph_pre_visitor()
-            for node in graph_or_function:
+
+        for node in graph_or_function:
+            for rule in self.rules:
                 delta = rule.try_rewrite(
                     model, graph_or_function, node, verbose=verbose, tracer=tracer
                 )
@@ -549,6 +549,9 @@ def _apply_to_graph_or_function(
                 )
 
                 count += 1
+                break
+
+        for rule in self.rules:
             if rule.graph_post_visitor:
                 rule.graph_post_visitor()