[Hierarchical Compile] Sort all regions identically (#158814)

mlazos · pytorchmergebot · commit ecde76c76475 · 2025-08-13T11:55:23.000Z
Before we would topologically sort each region individually, this works well except if some nodes have no arguments, then their order may change. To rectify this, we sort the first region as the reference region and use that sort order to sort the remaining regions. Pull Request resolved: #158814 Approved by: https://github.com/williamwen42
diff --git a/test/dynamo/test_graph_region_tracker.py b/test/dynamo/test_graph_region_tracker.py
@@ -370,6 +370,14 @@ def fn(x, y):
             """[[['y', 'o1'], ['y_1', 'o2'], ['y_2', 'o3']]]""",
         )
 
+    def test_region_sorting(self):
+        from torch._dynamo.graph_region_tracker import _sort_with_ref_region
+
+        index_to_rank = {0: 0, 2: 1, 1: 2}
+        regions = [[0, 1, 2], [1, 2, 0]]
+        _sort_with_ref_region(index_to_rank, regions)
+        self.assertExpectedInline(regions, """[[0, 2, 1], [1, 0, 2]]""")
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/torch/_dynamo/graph_region_tracker.py b/torch/_dynamo/graph_region_tracker.py
@@ -123,6 +123,18 @@ def _normalize_args(
     return (sorted_keys, tuple(_extract_args(arg) for arg in all_args))
 
 
+def _sort_with_ref_region(
+    index_to_rank: dict[int, int], regions: list[list[Any]]
+) -> None:
+    # sort topologically
+    # we need to handle edge cases where some nodes have no dependencies
+    # so first we map each node to its ranking
+    ref_region = regions[0]
+    sorted_indices = sorted(range(len(ref_region)), key=lambda i: index_to_rank[i])
+    for region in regions:
+        region[:] = [region[i] for i in sorted_indices]
+
+
 def get_global_state_key() -> GlobalStateKey:
     return (
         torch.is_grad_enabled(),
@@ -327,8 +339,13 @@ def get_identical_regions(self, graph: torch.fx.Graph) -> list[list[Region]]:
                 self._is_identical,
             )
             # sort topologically
-            for region in region_group:
-                region.sort(key=lambda n: topological_ranking[n])
+            # we need to handle edge cases where some nodes have no dependencies
+            # so first we map each node to its ranking,
+            ref_region = region_group[0]
+            index_to_rank = {
+                index: topological_ranking[n] for index, n in enumerate(ref_region)
+            }
+            _sort_with_ref_region(index_to_rank, region_group)
 
         return [
             region_group for region_group in region_groups if len(region_group[0]) > 1
