Return ConstantVariable(None) from WithExitFunctionVariable.exit to prevent NoneType crash inside autocast exception path (pytorch#153612)

Return ConstantVariable(None) from WithExitFunctionVariable.exit to prevent NoneType crash inside autocast exception path (pytorch#152503)

Copy of pytorch#152013 with PR time benchmarks updated (regressions seem unrelated)

Pull Request resolved: pytorch#152503
Approved by: https://github.com/anijain2305, https://github.com/Skylion007

Co-authored-by: Jason Ansel <[email protected]>

Author: wdziurdz

test/dynamo/test_exceptions.py

@@ -128,6 +128,28 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(ref, res)
 
+    def test_autocast_with_exception(self):
+        class Optimizer(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x):
+                raise NotImplementedError("Not implemented")
+
+            @staticmethod
+            def backward(ctx, grad_out):
+                return grad_out
+
+        @torch.compile
+        def f(x: torch.Tensor):
+            try:
+                with torch.autocast(device_type="cpu", dtype=None):
+                    Optimizer.apply(x)
+            except NotImplementedError:
+                return x + 1
+
+        inp = torch.ones(3)
+        out = f(inp)
+        self.assertTrue(torch.equal(out, inp + 1))
+
     @make_dynamo_test
     def test_propagate_exception_inside_ctx_manager(self):
         @contextlib.contextmanager

torch/_dynamo/variables/ctx_manager.py

@@ -839,6 +839,7 @@ def exit(self, tx: "InstructionTranslator", *args):
         tx.output.create_node(
             "call_function", torch.amp._exit_autocast, (self.state.proxy,), {}
         )
+        return variables.ConstantVariable.create(None)
 
     def enter(self, tx):
         ctx = torch.amp._enter_autocast(*self.target_values)