Add CLEANUP_THROW bytecode

test/dynamo/test_generator.py

@@ -593,6 +593,38 @@ def fn(t):
         self.assertEqual(i, 3)
         self.assertEqual(y, [(0, t), (1, t + 1), (2, t + 2)])
 
+    @unittest.skipIf(sys.version_info < (3, 12), "Test CLEANUP_THROW")
+    @unittest.expectedFailure
+    def test_cleanup_throw(self):
+        def nested_generator():
+            try:
+                yield 1
+                yield 2
+            except StopIteration:
+                return 123  # noqa: B901
+
+        def outer_generator():
+            yield from nested_generator()
+            yield 3
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def fn(t):
+            gen = outer_generator()
+            next(gen)  # Start the outer generator and enter the nested generato
+
+            i = 0
+            try:
+                # Force an exception while the generator is running
+                i = gen.throw(StopIteration("stop"))
+            except RuntimeError:
+                pass
+            return (i, t.sin())
+
+        t = torch.randn(2)
+        i, y = self._compile_check(fn, args=(t,))
+        self.assertEqual(i, 3)
+        self.assertEqual(y, t.sin())
+
     def test_iter(self):
         def whoo():
             i = 0
@@ -670,8 +702,6 @@ def fn(t):
         y = fn(t)
         self.assertEqual(y, t.sin())
 
-    @unittest.expectedFailure
-    @unittest.skipIf(sys.version_info < (3, 11), "Python 3.11+")
     def test_close_subgen(self):
         z = 0
 
@@ -844,8 +874,6 @@ def fn(t):
         with self.assertRaises(exc):
             fn(t)
 
-    @unittest.expectedFailure
-    @unittest.skipIf(sys.version_info < (3, 11), "Python 3.11+")
     def test_close_with_subgen(self):
         L = []
         z = 0
@@ -1450,8 +1478,6 @@ def fn(t):
 
         self._compile_check(fn)
 
-    @unittest.skipIf(sys.version_info < (3, 12), "Test CLEANUP_THROW")
-    @unittest.expectedFailure
     def test_exception_context_with_yield_from_with_context_cycle(self):
         # Check trying to create an exception context cycle:
         # https://bugs.python.org/issue40696

torch/_dynamo/symbolic_convert.py

@@ -1511,6 +1511,15 @@ def RAISE_VARARGS(self, inst):
                 self._raise_exception_variable(inst)
             unimplemented("raise ... from ...")
 
+    def CLEANUP_THROW(self, inst):
+        # https://github.com/python/cpython/pull/96010
+        tos = self.stack[-1]
+        assert isinstance(tos, ExceptionVariable)
+        if tos.exc_type is StopIteration:
+            unimplemented("CLEANUP_THROW with StopIteration")
+        else:
+            self.RERAISE(inst)
+
     def RERAISE(self, inst):
         if sys.version_info >= (3, 11):
             # RERAISE is currently supported in a narrow case of `raise ... from None`