[Trace Python Dispatcher] Support FuncTorchInterpreter (#144444)

yanboliang · pytorchmergebot · commit 43a00d73b364 · 2025-01-17T02:26:37.000Z
Pull Request resolved: #144444 Approved by: https://github.com/williamwen42, https://github.com/zou3519 ghstack dependencies: #144439
diff --git a/test/dynamo/test_python_dispatcher.py b/test/dynamo/test_python_dispatcher.py
@@ -102,6 +102,34 @@ def fn(x, dks):
         # Re-compile since the dispatch key set is different.
         self.assertEqual(counter.frame_count, 2)
 
+    def test_functorch_interpreter(self):
+        counter = CompileCounter()
+
+        def square_and_add(x, y):
+            interpreter = (
+                torch._functorch.pyfunctorch.retrieve_current_functorch_interpreter()
+            )
+            level = interpreter.level()
+            if interpreter.key() == torch._C._functorch.TransformType.Vmap:
+                return (x**2 + y) * level
+            else:
+                return x**2 * level
+
+        @torch.compile(backend=counter, fullgraph=True)
+        def fn(x, y):
+            return torch.vmap(square_and_add)(x, y)
+
+        x = torch.tensor([1, 2, 3, 4])
+        y = torch.tensor([10, 20, 30, 40])
+        self.assertEqual(fn(x, y), torch.tensor([11, 24, 39, 56]))
+        self.assertEqual(counter.frame_count, 1)
+
+        x = torch.tensor([1, 2, 3, 1])
+        y = torch.tensor([10, 20, 30, 10])
+        self.assertEqual(fn(x, y), torch.tensor([11, 24, 39, 11]))
+        # No recompile
+        self.assertEqual(counter.frame_count, 1)
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests
diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py
@@ -305,6 +305,7 @@
     "torch._C._functorch.peek_interpreter_stack": TorchInGraphFunctionVariable,
     "torch._C._functorch.unwrap_if_dead": TorchInGraphFunctionVariable,
     # everything else
+    "torch._functorch.pyfunctorch.coerce_cinterpreter": TorchInGraphFunctionVariable,
     "torch._higher_order_ops.triton_kernel_wrap.do_prune_configs": UserFunctionVariable,
     "torch._higher_order_ops.foreach_map.foreach_map": UserFunctionVariable,
     "torch._constrain_as_size": UserFunctionVariable,
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -217,6 +217,7 @@
 )
 from .torch import (
     DispatchKeySetVariable,
+    FuncTorchInterpreterVariable,
     TorchCtxManagerClassVariable,
     TorchInGraphFunctionVariable,
 )
@@ -668,7 +669,9 @@ def build_key_value(i, k, v):
             items = [SourcelessBuilder.create(self.tx, v) for v in value]
             self.install_guards(GuardBuilder.ID_MATCH)
             return FrozensetVariable(items, source=self.source)
-        elif isinstance(value, (enum.Enum, torch.DispatchKey)):
+        elif isinstance(
+            value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
+        ):
             self.install_guards(GuardBuilder.ID_MATCH)
             return EnumVariable(value=value, source=self.source)
         elif DebuggingVariable.is_reorderable_logging_function(value):
@@ -857,6 +860,9 @@ def build_key_value(i, k, v):
         elif isinstance(value, (torch._C._SDPAParams)):
             self.install_guards(GuardBuilder.TYPE_MATCH)
             return SDPAParamsVariable.create(self.tx, value, self.source)
+        elif isinstance(value, torch._functorch.pyfunctorch.FuncTorchInterpreter):
+            self.install_guards(GuardBuilder.ID_MATCH)
+            return FuncTorchInterpreterVariable(value)
         elif isinstance(value, torch.Event):
             self.install_guards(GuardBuilder.ID_MATCH)
             torch._dynamo.utils.store_user_object_weakref(value)
@@ -2960,7 +2966,9 @@ def create(tx: "InstructionTranslator", value) -> VariableTracker:
             return trace_rules.lookup_callable(value)(value)
         elif is_function_or_wrapper(value):
             return trace_rules.lookup(value)(value)
-        elif isinstance(value, (enum.Enum, torch.DispatchKey)):
+        elif isinstance(
+            value, (enum.Enum, torch.DispatchKey, torch._C._functorch.TransformType)
+        ):
             return EnumVariable(value)
         elif isinstance(value, (type, abc.ABCMeta)):
             return UserDefinedClassVariable(value)
@@ -3025,6 +3033,11 @@ def make_type_handlers():
         handlers[torch.DispatchKeySet] = lambda tx, value: DispatchKeySetVariable(
             value, mutation_type=ValueMutationNew()
         )
+        handlers[
+            torch._functorch.pyfunctorch.FuncTorchInterpreter
+        ] = lambda tx, value: FuncTorchInterpreterVariable(
+            value, mutation_type=ValueMutationNew()
+        )
 
         handlers[
             torch.distributions.constraints._Real
diff --git a/torch/_dynamo/variables/torch.py b/torch/_dynamo/variables/torch.py
@@ -826,6 +826,25 @@ def handle_unsafe_set_version_counter(
                 _unsafe_set_version_counter
             ).call_function(tx, [*args], kwargs)
 
+        @register(torch._C._functorch.peek_interpreter_stack)
+        def handle_functorch_peek_interpreter_stack(
+            self, tx: "InstructionTranslator", *args, **kwargs
+        ):
+            # Wrap C++ interpreter (torch._C._functorch.CInterpreter) as UserDefinedObjectVariable,
+            # but Python interpreter (torch._functorch.pyfunctorch.FuncTorchInterpreter) as FuncTorchInterpreterVariable.
+            return UserDefinedObjectVariable(
+                torch._C._functorch.peek_interpreter_stack()
+            )
+
+        @register(torch._functorch.pyfunctorch.coerce_cinterpreter)
+        def handle_functorch_pyfunctorch_coerce_cinterpreter(
+            self, tx: "InstructionTranslator", *args, **kwargs
+        ):
+            cinterpreter = args[0].value
+            return FuncTorchInterpreterVariable(
+                torch._functorch.pyfunctorch.coerce_cinterpreter(cinterpreter)
+            )
+
         @register(torch.tensor)
         def handle_torch_tensor(self, tx: "InstructionTranslator", *args, **kwargs):
             def check_any_unspec(x):
@@ -1260,3 +1279,31 @@ def call_method(
         elif name == "highestPriorityTypeId":
             return variables.EnumVariable(self.value.highestPriorityTypeId())
         return super().call_method(tx, name, args, kwargs)
+
+
+class FuncTorchInterpreterVariable(BaseTorchVariable):
+    """represents torch._functorch.pyfunctorch.FuncTorchInterpreter"""
+
+    @classmethod
+    def create_with_source(cls, value, source):
+        install_guard(source.make_guard(GuardBuilder.ID_MATCH))
+        return cls(value, source=source)
+
+    def call_method(
+        self,
+        tx,
+        name,
+        args: "List[VariableTracker]",
+        kwargs: "Dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        if name == "key":
+            return variables.EnumVariable(self.value.key())
+        elif name == "process":
+            return tx.inline_user_function_return(
+                variables.UserFunctionVariable(self.value.process.__func__),
+                [self] + args,
+                kwargs,
+            )
+        elif name in ["level", "batch_size", "randomness"]:
+            return variables.ConstantVariable.create(getattr(self.value, name)())
+        return super().call_method(tx, name, args, kwargs)
