IterativePassGroup: warn when the step cap is hit (#24759)

Summary:

`IterativePassGroup` re-runs its passes until they stop changing the
graph, up to `steps` rounds. Until now it handed `steps` to
`ExportedProgramPassManager`, which stops quietly at the cap, so a group that
never converged looked the same as one that did.

The group now loops itself, one manager step per round, and returns as soon as
a round leaves the graph unchanged. If a multi-step group is still changing the
graph after its last allowed round, it logs a warning naming the group and the
cap. Single-step groups run once by design and never warn.

`AdvanceQuantizeOpAboveDefInBranchPass` now reports whether it changed the
graph. It returned `ExportPass.call`'s result, which always says modified, so
a group containing it could never converge: it ran all of its rounds and hit
this warning. Like `AdvanceQuantizeOpAboveDefChainPass`, it now re-traces and
reports a change only when it advanced a quantize. The rewrite itself is
unchanged.

Differential Revision: D129194624

Pull Request resolved: #24759

Author: mcremon-meta

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -288,7 +288,7 @@ def test_advance_branched_quantize(self) -> None:
             0,
         )
 
-    def _advance_two_branch_quantize(
+    def _two_branch_quantize_graph(
         self,
         first_qparams: Tuple[float, int],
         second_qparams: Tuple[float, int],
@@ -311,13 +311,34 @@ def _advance_two_branch_quantize(
                 )
             )
         builder.output(quantized)
+        return builder.get_graph_module()
+
+    def _advance_two_branch_quantize(
+        self,
+        first_qparams: Tuple[float, int],
+        second_qparams: Tuple[float, int],
+    ) -> torch.fx.GraphModule:
         graph_module = cast(
             PassResult,
-            AdvanceQuantizeOpAboveDefInBranchPass()(builder.get_graph_module()),
+            AdvanceQuantizeOpAboveDefInBranchPass()(
+                self._two_branch_quantize_graph(first_qparams, second_qparams)
+            ),
         ).graph_module
         graph_module.graph.eliminate_dead_code()
         return graph_module
 
+    def test_reports_whether_it_advanced_a_quantize(self) -> None:
+        for second_qparams, advanced in (((0.1, 10), True), ((2.0, 4), False)):
+            with self.subTest(advanced=advanced):
+                result = cast(
+                    PassResult,
+                    AdvanceQuantizeOpAboveDefInBranchPass()(
+                        self._two_branch_quantize_graph((0.1, 10), second_qparams)
+                    ),
+                )
+
+                self.assertEqual(result.modified, advanced)
+
     def test_does_not_advance_quantize_when_qparams_differ(self) -> None:
         # Requantizing from one hoisted quantize can clip or coarsen a branch, so
         # the branches are left alone unless they already share qparams

backends/transforms/advance_quantize_op_above_def_in_branch.py

@@ -107,7 +107,8 @@ def get_descendent_quant_ops(self, node: torch.fx.Node) -> List[torch.fx.Node]:
 
         return descendent_quant_ops
 
-    def advance_quantize_op(self, graph_module: torch.fx.GraphModule):
+    def advance_quantize_op(self, graph_module: torch.fx.GraphModule) -> bool:
+        modified = False
         graph = graph_module.graph
         for node in graph.nodes:
             # We are only interested in call functions and placeholders
@@ -150,6 +151,7 @@ def advance_quantize_op(self, graph_module: torch.fx.GraphModule):
                     node.replace_all_uses_with(dom_quant_node)
                     dom_quant_node.args = dom_quant_args
                 graph.erase_node(quant_node)
+                modified = True
                 continue
 
             # Otherwise we have multiple quant descendents. Requantizing them from
@@ -185,12 +187,16 @@ def advance_quantize_op(self, graph_module: torch.fx.GraphModule):
                     )
                     dequant_node.args = (quant_node.args[0],) + rep_args
                     quant_node.args = (dequant_node,) + quant_node.args[1:]
+            modified = True
 
-        graph_module.recompile()
-        graph_module.graph.eliminate_dead_code()
+        return modified
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         self.graph_module = graph_module
-        self.advance_quantize_op(graph_module)
-        result = super().call(graph_module)
-        return result
+        modified = self.advance_quantize_op(graph_module)
+        if modified:
+            graph_module.recompile()
+            graph_module.graph.eliminate_dead_code()
+            return super().call(graph_module)
+
+        return PassResult(graph_module, False)

backends/transforms/fused_quant/pass_base.py

@@ -5,10 +5,14 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import logging
+
 from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
 from executorch.exir.pass_manager import ExportedProgramPassManager, PassType
 from torch.export import ExportedProgram
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 class IterativePassGroup(ExportedProgramPassBase):
     """Runs a group of passes as a single ExportedProgramPassBase.
@@ -22,8 +26,23 @@ def __init__(self, passes: list[PassType], steps: int) -> None:
         # rebuild this group with wrapped leaf passes.
         self._passes = passes
         self._steps = steps
-        self._pm = ExportedProgramPassManager(passes, steps=steps)
+        self._pm = ExportedProgramPassManager(passes, steps=1)
 
     def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
-        result = self._pm(exported_program)
-        return ExportedProgramPassResult(result.exported_program, result.modified)
+        modified = False
+        for _ in range(self._steps):
+            result = self._pm(exported_program)
+            exported_program = result.exported_program
+            if not result.modified:
+                return ExportedProgramPassResult(exported_program, modified)
+            modified = True
+        # A single-step group runs its passes once by design; only a group meant
+        # to iterate to a fixed point can fail to converge.
+        if self._steps > 1:
+            logger.warning(
+                "%s did not converge within %d steps; the graph was still changing "
+                "on the last step",
+                type(self).__name__,
+                self._steps,
+            )
+        return ExportedProgramPassResult(exported_program, modified)