[dynamo][pytree][1/N] make CXX pytree traceable: tree_iter / tree_leaves (pytorch#137397)

Pull Request resolved: pytorch#137397
Approved by: https://github.com/jansel

Author: XuehaiPan

test/dynamo/test_misc.py

@@ -32,7 +32,7 @@
 import torch._dynamo.testing
 import torch._inductor.test_case
 import torch.onnx.operators
-import torch.utils._pytree as pytree
+import torch.utils._pytree as python_pytree
 import torch.utils.cpp_extension
 from torch import Tensor
 from torch._C import FileCheck
@@ -89,9 +89,11 @@
 from torch.testing._internal.logging_utils import logs_to_string
 
 
-HAS_OPTREE = importlib.util.find_spec("optree")
+HAS_OPTREE = python_pytree._cxx_pytree_exists
 if HAS_OPTREE:
-    import optree
+    import torch.utils._cxx_pytree as cxx_pytree
+else:
+    cxx_pytree = None
 
 MyTuple = collections.namedtuple("MyTuple", ["a", "b", "ab"])
 T = typing.TypeVar("T")
@@ -293,9 +295,9 @@ def fn(x):
 
     @unittest.skipIf(not HAS_OPTREE, "missing optree package")
     def test_optree_graph_break_message(self):
-        @torch.compile(
-            backend="eager",
-        )
+        import optree
+
+        @torch.compile(backend="eager")
         def fn(x):
             d = {"a": 1}
             optree.tree_flatten(d)
@@ -8722,9 +8724,9 @@ def fn():
 
     def test_tracing_py_tree(self):
         def fn(xs):
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             res = [x.clone() for x in flat_xs]
-            return pytree.tree_unflatten(res, spec)
+            return python_pytree.tree_unflatten(res, spec)
 
         xs = [torch.tensor(i) for i in range(3)]
 
@@ -8734,12 +8736,10 @@ def fn(xs):
         self.assertEqual(counter.op_count, 3)
 
     def test_tracing_nested_py_tree(self):
-        import torch.utils._pytree as pytree
-
         def fn(xs):
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             res = [x.clone() for x in flat_xs]
-            return pytree.tree_unflatten(res, spec)
+            return python_pytree.tree_unflatten(res, spec)
 
         xs = [torch.tensor(i) for i in range(3)]
         xsl = [xs, xs, xs, xs]
@@ -8752,12 +8752,10 @@ def fn(xs):
         self.assertEqual(counter.op_count, 12)
 
     def test_tracing_nested_py_tree_tuples(self):
-        import torch.utils._pytree as pytree
-
         def fn(xs):
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             res = [x.clone() for x in flat_xs]
-            return pytree.tree_unflatten(res, spec)
+            return python_pytree.tree_unflatten(res, spec)
 
         xs = [torch.tensor(i) for i in range(3)]
         xsl = (xs, xs, xs, xs)
@@ -8770,12 +8768,10 @@ def fn(xs):
         self.assertEqual(counter.op_count, 12)
 
     def test_tracing_nested_py_tree_dicts(self):
-        import torch.utils._pytree as pytree
-
         def fn(xs):
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             res = [x.clone() for x in flat_xs]
-            return pytree.tree_unflatten(res, spec)
+            return python_pytree.tree_unflatten(res, spec)
 
         xs = [torch.tensor(i) for i in range(3)]
         xsl = {
@@ -8808,12 +8804,10 @@ def fn(x):
         self.assertEqual(counter.op_count, 2)
 
     def test_tracing_nested_py_tree_mixed_all(self):
-        import torch.utils._pytree as pytree
-
         def fn(xs):
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             res = [x.clone() for x in flat_xs]
-            return pytree.tree_unflatten(res, spec)
+            return python_pytree.tree_unflatten(res, spec)
 
         xs = [torch.tensor(i) for i in range(3)]
         xsa = (xs, xs)
@@ -8858,13 +8852,12 @@ def fn(x):
         self.assertEqual(cnt.frame_count, 2)
 
     def test_tracing_py_tree_tensor_subclass(self):
-        import torch.utils._pytree as pytree
         from torch.testing._internal.two_tensor import TwoTensor
         from torch.utils.checkpoint import checkpoint
 
         def fn(xs):
             nested_xs = [[xs]]
-            flat_xs, spec = pytree.tree_flatten(xs)
+            flat_xs, spec = python_pytree.tree_flatten(xs)
             return flat_xs[0].clone()
 
         # use checkpoint to trigger a "sourceless" tensor subclass
@@ -8879,13 +8872,11 @@ def checkpoint_fn(xs):
         self.assertEqual(counter.op_count, 2)
 
     def test_tracing_tree_map_only(self):
-        import torch.utils._pytree as pytree
-
         def fn(xs):
             def mapper(x):
                 return x.clone()
 
-            y = pytree.tree_map_only(torch.Tensor, mapper, xs)
+            y = python_pytree.tree_map_only(torch.Tensor, mapper, xs)
             return y
 
         xs = [torch.tensor(i) for i in range(3)] + ["hi"]
@@ -10235,7 +10226,9 @@ def fn(x, y):
         self.assertEqual(actual, expected)
 
     def test_pytree_tree_leaves(self):
-        implemtations = [("python", pytree)]
+        implemtations = [("python", python_pytree)]
+        if cxx_pytree is not None:
+            implemtations.append(("cxx", cxx_pytree))
 
         for name, module in implemtations:
             with self.subTest(f"pytree implement: {name}"):
@@ -10267,7 +10260,7 @@ def fn(x):
                 self.assertEqual(actual, expected)
 
     def test_pytree_tree_flatten_unflatten(self):
-        implemtations = [("python", pytree)]
+        implemtations = [("python", python_pytree)]
 
         for name, module in implemtations:
             with self.subTest(f"pytree implement: {name}"):
@@ -10316,7 +10309,7 @@ def fn(x, y):
             self.assertEqual(actual, expected)
 
     def test_pytree_tree_map(self):
-        implemtations = [("python", pytree)]
+        implemtations = [("python", python_pytree)]
 
         for name, module in implemtations:
             with self.subTest(f"pytree implement: {name}"):

test/dynamo/test_trace_rules.py

@@ -323,10 +323,16 @@ def _check_set_equality(self, generated, used, rule_map, ignored_set):
     # or loaded in case there is typo in the strings.
     def test_skipfiles_inlinelist(self):
         for m in LEGACY_MOD_INLINELIST.union(MOD_INLINELIST):
-            self.assertTrue(
-                isinstance(importlib.import_module(m), types.ModuleType),
-                f"{m} from trace_rules.MOD_INLINELIST/LEGACY_MOD_INLINELIST is not a python module, please check and correct it.",
-            )
+            try:
+                mod = importlib.import_module(m)
+            except ImportError:
+                continue
+            else:
+                self.assertTrue(
+                    isinstance(mod, types.ModuleType),
+                    f"{m} from trace_rules.MOD_INLINELIST/LEGACY_MOD_INLINELIST "
+                    "is not a python module, please check and correct it.",
+                )
 
     @unittest.skip(
         "This test keeps getting broken and our disable infra is not handling well. see #120627"

torch/_dynamo/guards.py

@@ -2090,10 +2090,11 @@ def _set_guard_export_info(self, guard, code_list, provided_guarded_object=None)
         obj_ref = None
         # Not necessary to have weakref for Enum type, but there is a bug that
         # makes hasattr(guarded_object.__class__, "__weakref__") return True.
+        supports_weakref = (
+            getattr(guarded_object.__class__, "__weakrefoffset__", 0) != 0
+        )
         # See D64140537 for why we are checking for tuple.
-        if hasattr(guarded_object.__class__, "__weakref__") and not isinstance(
-            guarded_object, (enum.Enum, tuple)
-        ):
+        if supports_weakref and not isinstance(guarded_object, (enum.Enum, tuple)):
             obj_ref = weakref.ref(guarded_object)
 
         guard.set_export_info(

torch/_dynamo/polyfills/__init__.py

@@ -23,6 +23,7 @@
         itertools as itertools,
         operator as operator,
         os as os,
+        pytree as pytree,
         sys as sys,
     )
 

torch/_dynamo/polyfills/loader.py

@@ -18,6 +18,7 @@
     "itertools",
     "operator",
     "os",
+    "pytree",
     "sys",
 )
 POLYFILLED_MODULES: Tuple["ModuleType", ...] = tuple(

torch/_dynamo/polyfills/pytree.py

@@ -0,0 +1,97 @@
+"""
+Python polyfills for torch.utils.pytree
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable, TYPE_CHECKING
+
+import torch.utils._pytree as python_pytree
+
+from ..decorators import substitute_in_graph
+
+
+if TYPE_CHECKING:
+    from torch.utils._cxx_pytree import PyTree
+
+
+__all__: list[str] = []
+
+
+if python_pytree._cxx_pytree_exists:
+    import optree
+    import optree._C
+
+    import torch.utils._cxx_pytree as cxx_pytree
+
+    @substitute_in_graph(
+        optree._C.is_dict_insertion_ordered,
+        can_constant_fold_through=True,
+    )
+    def _(*args: Any, **kwargs: Any) -> bool:
+        # In namespace 'torch', the dictionary is always traversed in insertion order.
+        # This function returns True.
+        raise ValueError(
+            "Should not be called directly "
+            "because the original function will be called in the constant fold path."
+        )
+
+    __name = ""
+    for __name in (
+        "is_namedtuple",
+        "is_namedtuple_class",
+        "is_namedtuple_instance",
+        "is_structseq",
+        "is_structseq_class",
+        "is_structseq_instance",
+        "namedtuple_fields",
+        "structseq_fields",
+    ):
+        __func = getattr(optree, __name)
+        substitute_in_graph(__func, can_constant_fold_through=True)(
+            __func.__python_implementation__
+        )
+        del __func
+    del __name
+
+    @substitute_in_graph(cxx_pytree.tree_is_leaf, can_constant_fold_through=True)
+    def tree_is_leaf(
+        tree: PyTree,
+        is_leaf: Callable[[PyTree], bool] | None = None,
+    ) -> bool:
+        if tree is None or (is_leaf is not None and is_leaf(tree)):
+            return True
+        if optree.register_pytree_node.get(type(tree), namespace="torch") is None:  # type: ignore[attr-defined]
+            return True
+        return False
+
+    @substitute_in_graph(cxx_pytree.tree_iter, can_constant_fold_through=False)
+    def tree_iter(
+        tree: PyTree,
+        is_leaf: Callable[[PyTree], bool] | None = None,
+    ) -> Iterable[Any]:
+        stack = [tree]
+        while stack:
+            node = stack.pop()
+            if tree_is_leaf(node, is_leaf=is_leaf):
+                yield node
+                continue
+
+            children, *_ = optree.tree_flatten_one_level(
+                node,
+                is_leaf=is_leaf,
+                none_is_leaf=True,
+                namespace="torch",
+            )
+            stack.extend(reversed(children))
+
+    __all__ += ["tree_iter"]
+
+    @substitute_in_graph(cxx_pytree.tree_leaves, can_constant_fold_through=True)
+    def tree_leaves(
+        tree: PyTree,
+        is_leaf: Callable[[PyTree], bool] | None = None,
+    ) -> list[Any]:
+        return list(tree_iter(tree, is_leaf=is_leaf))
+
+    __all__ += ["tree_leaves"]

torch/_dynamo/trace_rules.py

@@ -3311,6 +3311,7 @@ def _module_dir(m: types.ModuleType):
     "torch.testing",
     "torch.utils._content_store",
     "torch.utils._contextlib",
+    "torch.utils._cxx_pytree",
     "torch.utils._device",
     "torch.utils._foreach_utils",
     "torch.utils._python_dispatch",