[Hierarchical Compilation] Use universal flatten APIs (pytorch#152505)

Pull Request resolved: pytorch#152505
Approved by: https://github.com/anijain2305
ghstack dependencies: pytorch#152389

Author: mlazos

test/dynamo/test_graph_deduplication.py

@@ -2,7 +2,6 @@
 # flake8: noqa: B950
 import torch
 import torch.fx
-from torch._dynamo.graph_deduplication import _flatten_args_kwargs
 from torch._dynamo.graph_utils import _detect_cycles
 from torch._dynamo.test_case import TestCase
 from torch._dynamo.testing import AotEagerAndRecordGraphs, normalize_gm
@@ -583,13 +582,6 @@ def forward(self, arg0_1: "f32[10, 10]", arg1_1: "f32[10, 20]"):
 """,
         )
 
-    def test_flatten_with_slices(self):
-        tree = [{"x": 3}, ["x", slice(1, 2, 3), 1], [4, 5, 6, [slice(3, 4, 5)]]]
-        out = _flatten_args_kwargs(tree)
-        self.assertExpectedInline(
-            str(out), """[3, 'x', 1, 2, 3, 1, 4, 5, 6, 3, 4, 5]"""
-        )
-
     def test_cycle_detection_no_cycle(self):
         def fn(x, y):
             x0 = x + 1

torch/_dynamo/graph_deduplication.py

@@ -11,7 +11,7 @@
 import operator
 from collections import defaultdict
 from collections.abc import Generator, Iterable
-from typing import Any, Optional
+from typing import Any
 
 import torch
 import torch.fx
@@ -20,7 +20,7 @@
 from torch.utils._ordered_set import OrderedSet
 
 from .graph_region_tracker import Node, Region
-from .graph_utils import _detect_cycles, _flatten_args_kwargs
+from .graph_utils import _detect_cycles, _get_flat_args, _get_flat_args_unique
 
 
 log = logging.getLogger(__name__)
@@ -92,7 +92,10 @@ def apply_graph_deduplication(output_graph) -> dict[str, torch.fx.GraphModule]:
                 node_to_additional_deps,
             )
 
-    _stable_topological_sort(output_graph.graph, node_to_additional_deps)
+    _stable_topological_sort(
+        output_graph.graph,
+        node_to_additional_deps,  # type: ignore[arg-type]
+    )
     return sub_gms
 
 
@@ -109,7 +112,7 @@ def _replace_region_with_subgraph(
     sub_args = []
     for node_ind, arg_ind in node_ind_arg_ind:
         node = region[node_ind]
-        flattened_args_kwargs = _flatten_args_kwargs((node.args, node.kwargs))
+        flattened_args_kwargs = _get_flat_args(node, {})
         sub_args.append(flattened_args_kwargs[arg_ind])
 
     invoke_args = (get_subgraph_node, subgraph_name, *sub_args)
@@ -162,7 +165,7 @@ def _get_external_inputs(
     external_node_to_indices = dict()
     region_unique = set(region)
     for node_ind, node in enumerate(region):
-        flattened_args_kwargs = _flatten_args_kwargs((node.args, node.kwargs))
+        flattened_args_kwargs = _get_flat_args(node, {})
         for arg_ind, in_node in enumerate(flattened_args_kwargs):
             if (
                 isinstance(in_node, Node)
@@ -237,23 +240,9 @@ def _create_subgraph(
     return subgraph, node_ind_input_inds
 
 
-def _args(
-    n: torch.fx.Node,
-    node_to_additional_deps: Optional[dict[torch.fx.Node, list[torch.fx.Node]]] = None,
-) -> list[torch.fx.node.Argument]:
-    if node_to_additional_deps is None:
-        node_to_additional_deps = {}
-
-    args: list[torch.fx.node.Argument] = []
-    torch.fx.map_arg((n.args, n.kwargs), args.append)
-    if n in node_to_additional_deps:
-        args.extend(node_to_additional_deps[n])
-    return args
-
-
 def _stable_topological_sort(
     graph: torch.fx.Graph,
-    node_to_additional_deps: dict[torch.fx.Node, list[torch.fx.Node]],
+    node_to_additional_deps: dict[torch.fx.Node, OrderedSet[torch.fx.Node]],
 ) -> None:
     # Nodes are in exactly one of these four collections:
 
@@ -283,7 +272,9 @@ def _stable_topological_sort(
             continue
 
         waiting_for = [
-            x for x in _args(node, node_to_additional_deps) if x not in ready
+            x
+            for x in _get_flat_args_unique(node, node_to_additional_deps)
+            if x not in ready
         ]
         if waiting_for:
             # We have unprocessed input nodes. Might as well wait for the last
@@ -328,7 +319,7 @@ def prev_cur_nodes(
             prev_nodes.append(cur_node)
 
     for prev_nodes, cur_node in prev_cur_nodes(all_nodes):
-        args_unique = _args(cur_node)
+        args_unique = _get_flat_args_unique(cur_node, {})
         additional_deps = node_to_additional_deps[cur_node]
         additional_deps.extend(n for n in all_nodes_dep_on if n not in args_unique)
         if cur_node.target in global_state_targets:

torch/_dynamo/graph_region_tracker.py

@@ -28,7 +28,7 @@
 from torch.utils._ordered_set import OrderedSet
 from torch.utils._pytree import tree_flatten
 
-from .graph_utils import _flatten_args_kwargs
+from .graph_utils import _get_flat_args_unique
 
 
 T = TypeVar("T")
@@ -416,7 +416,7 @@ def _populate_recursive_ancestor_map(graph: torch.fx.Graph) -> dict[Node, set[No
     for node in graph.nodes:
         node_to_recursive_ancestors[node] = set()
     for node in graph.nodes:
-        all_args = _flatten_args_kwargs((node.args, node.kwargs))
+        all_args = _get_flat_args_unique(node, {})
         for arg in all_args:
             if isinstance(arg, Node):
                 node_to_recursive_ancestors[node].update(

torch/_dynamo/graph_utils.py

@@ -1,32 +1,33 @@
 from collections import deque
 from typing import Any
 
-from torch.fx import Graph, Node
-from torch.utils._pytree import tree_flatten
+from torch.fx import Graph, map_arg, Node
+from torch.utils._ordered_set import OrderedSet
 
 
 # flattens with support for slices
 # Note: a better way to do this would
 # be register/unregister slices as pytree nodes
 # but there is no unregister API in the pytorch
 # pytree impl
-def _flatten_args_kwargs(args: Any) -> list[Node]:
-    fully_flattened = []
-
-    def flatten(args: Any) -> None:
-        flattened, _ = tree_flatten(args)
-        for arg in flattened:
-            if isinstance(arg, slice):
-                start = arg.start
-                stop = arg.stop
-                step = arg.step
-                flatten((start, stop, step))
-            else:
-                fully_flattened.append(arg)
-
-    flatten(args)
-
-    return fully_flattened
+def _get_flat_args(
+    node: Node, node_to_additional_deps: dict[Node, OrderedSet[Node]]
+) -> list[Node]:
+    args = list[Any]()
+    map_arg((node.args, node.kwargs), args.append)
+    if node in node_to_additional_deps:
+        args.extend(node_to_additional_deps[node])
+    return args
+
+
+def _get_flat_args_unique(
+    node: Node, node_to_additional_deps: dict[Node, OrderedSet[Node]]
+) -> OrderedSet[Node]:
+    args = OrderedSet[Node]()
+    map_arg((node.args, node.kwargs), args.add)
+    if node in node_to_additional_deps:
+        args.update(node_to_additional_deps[node])
+    return args
 
 
 def _detect_cycles(graph: Graph) -> str:

torch/_dynamo/utils.py

@@ -92,6 +92,8 @@
 from torch.utils._triton import has_triton, has_triton_package
 from torch.utils.hooks import RemovableHandle
 
+from .graph_utils import _get_flat_args
+
 
 if typing.TYPE_CHECKING:
     from collections.abc import (
@@ -3150,7 +3152,9 @@ def get_fake_value(node, tx, allow_non_graph_fake=False):
     args, kwargs = get_fake_values_from_nodes(
         tx, (node.args, node.kwargs), allow_non_graph_fake
     )
-    flat_args_kwargs, _ = pytree.tree_flatten((args, kwargs))
+    flat_args_kwargs = get_fake_values_from_nodes(
+        tx, _get_flat_args(node, {}), allow_non_graph_fake
+    )
     id_to_initial_version = {
         id(arg): arg._version for arg in flat_args_kwargs if is_fake(arg)
     }