Raise exception if vmap (eager) calls compiled function (pytorch#140439)

Fixes pytorch#138422

This is not a proper fix for pytorch#140439, but more of a way to prevent a user from seeing a nasty error inside the C++ code.

Pull Request resolved: pytorch#140439
Approved by: https://github.com/zou3519

test/dynamo/test_higher_order_ops.py

@@ -5967,6 +5967,60 @@ def wrapped_fn(x):
         ):
             torch.func.vmap(fn)(x)
 
+    def test_vmap_call_compiled_backward_fn(self):
+        # See PyTorch issue #138422
+        @torch.compile
+        def f(x):
+            return x**2
+
+        x = torch.randn(2, requires_grad=True)
+        y = f(x)
+
+        def get_vjp(v):
+            return torch.autograd.grad(y, x, v)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "It looks like you're trying to call a compiled backward function within vmap/grad/vjp, which isn't supported",
+        ):
+            torch.func.vjp(get_vjp, x)
+
+    def test_vjp_call_compiled_backward_fn(self):
+        # See PyTorch issue #138422
+        @torch.compile
+        def f(x):
+            return x**2
+
+        x = torch.randn(2, requires_grad=True)
+        y = f(x)
+
+        def get_vjp(v):
+            return torch.autograd.grad(y, x, v)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "It looks like you're trying to call a compiled backward function within vmap/grad/vjp, which isn't supported",
+        ):
+            torch.func.vjp(get_vjp, x)
+
+    def test_grad_call_compiled_backward_fn(self):
+        # See PyTorch issue #138422
+        @torch.compile
+        def f(x):
+            return x**2
+
+        x = torch.randn(2, requires_grad=True)
+        y = f(x)
+
+        def get_vjp(v):
+            return torch.autograd.grad(y, x, v)
+
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "It looks like you're trying to call a compiled backward function within vmap/grad/vjp, which isn't supported",
+        ):
+            torch.func.grad(get_vjp)(x)
+
     def test_grad_call_torch_compile_fn(self):
         def wrapped_fn(x):
             return x.sin().sum()

torch/_functorch/_aot_autograd/runtime_wrappers.py

@@ -1732,6 +1732,19 @@ def backward(double_ctx, *args):
 
                 return CompiledFunctionBackward.apply(*all_args)
 
+            @staticmethod
+            def _raise_if_functorch_active():
+                # not ideal but prevent the user from seeing a nasty traceback - See #138422
+                stack = torch._C._functorch.peek_interpreter_stack()
+                torch._check(
+                    stack is None,
+                    lambda: (
+                        "It looks like you're trying to call a compiled backward function within vmap/grad/vjp, "
+                        "which isn't supported. Try wrapping vmap inside torch.compile, or skip compiling the "
+                        "backward function."
+                    ),
+                )
+
             @staticmethod
             def _backward_prologue(ctx, *flat_args):
                 # Calling convention: we expect a grad_out passed to the backward:
@@ -1743,6 +1756,8 @@ def _backward_prologue(ctx, *flat_args):
                 # - updated inputs due to metadata-only mutations.
                 # We need to return them in the forward, but ensure that they all do not get gradients in the backward,
                 # and we filter them out here before passing the remaining grad_outputs into the compiled backward.
+                CompiledFunction._raise_if_functorch_active()
+
                 num_intermediate_bases = (
                     CompiledFunction.metadata.num_intermediate_bases
                 )