[dynamo] implement IteratorVariable and polyfill fallbacks for enumerate (pytorch#131725)

Fixes pytorch#112794.

Pull Request resolved: pytorch#131725
Approved by: https://github.com/anijain2305
ghstack dependencies: pytorch#131413, pytorch#131716

Author: williamwen42

test/dynamo/test_functions.py

@@ -2987,6 +2987,45 @@ def test_map_unpack_twice(a, b):
         l2 = list(m)
         return l1, l2
 
+    @make_test
+    def test_enumerate(a, b):
+        return list(enumerate([a, b], start=1)), a + 1
+
+    @make_test
+    def test_map_enumerate(a, b):
+        return list(enumerate(map(lambda x: x + 1, [a, b]), start=1)), a + 1
+
+    def test_enumerate_custom(self):
+        class MyClass:
+            def __iter__(self):
+                self.a = 1
+                return self
+
+            def __next__(self):
+                if self.a > 3:
+                    raise StopIteration
+                self.a += 1
+                return self.a
+
+        def fn(x):
+            for i, it in enumerate(MyClass()):
+                x += i + it
+            return x
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        self.assertEqual(fn(torch.ones(3, 3)), opt_fn(torch.ones(3, 3)))
+
+    def test_enumerate_reconstruct(self):
+        def fn(a, b):
+            return enumerate([a, b], start=1)
+
+        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
+        inps = (torch.randn(3, 3), torch.randn(3, 3))
+        it1 = fn(*inps)
+        it2 = opt_fn(*inps)
+        self.assertIsInstance(it2, enumerate)
+        self.assertEqual(list(it1), list(it2))
+
 
 def udf_mul(x, y):
     return x * y

torch/_dynamo/polyfill.py

@@ -24,7 +24,7 @@ def any(iterator):
 
 
 def index(iterator, item, start=0, end=None):
-    for i, elem in enumerate(list(iterator))[start:end]:
+    for i, elem in enumerate(list(iterator)[start:end], start):
         if item == elem:
             return i
     # This will not run in dynamo
@@ -124,3 +124,10 @@ def getattr_and_trace(*args, **kwargs):
     attr_name = args[1]
     fn = getattr(wrapper_obj, attr_name)
     return fn(*args[2:], **kwargs)
+
+
+def enumerate(iterable, start=0):
+    n = start
+    for elem in iterable:
+        yield n, elem
+        n += 1

torch/_dynamo/variables/__init__.py

@@ -42,6 +42,7 @@
 from .iter import (
     CountIteratorVariable,
     CycleIteratorVariable,
+    EnumerateVariable,
     IteratorVariable,
     ItertoolsVariable,
     MapVariable,

torch/_dynamo/variables/builtin.py

@@ -1419,22 +1419,29 @@ def call_zip(self, tx: "InstructionTranslator", *args, **kwargs):
         ]
         return variables.ZipVariable(args, strict=strict, mutable_local=MutableLocal())
 
-    def call_enumerate(self, tx: "InstructionTranslator", *args):
-        if len(args) == 1:
+    _call_enumerate_polyfill = _polyfill_call_impl("enumerate")
+
+    def call_enumerate(self, tx: "InstructionTranslator", iterable, start=_SENTINEL):
+        if start is self._SENTINEL:
             start = 0
         else:
-            assert len(args) == 2
-            assert isinstance(args[1], variables.ConstantVariable)
-            start = args[1].as_python_constant()
-        if args[0].has_unpack_var_sequence(tx):
-            items = [
-                variables.TupleVariable(
-                    [variables.ConstantVariable.create(idx), var],
-                )
-                for idx, var in enumerate(args[0].unpack_var_sequence(tx), start)
-            ]
-            return variables.TupleVariable(items)
-        # could have an iterable version
+            assert isinstance(start, variables.ConstantVariable)
+            start = start.as_python_constant()
+
+        if iterable.has_unpack_var_sequence(tx):
+            return variables.EnumerateVariable(
+                iterable.unpack_var_sequence(tx),
+                start,
+                mutable_local=MutableLocal(),
+            )
+        elif isinstance(iterable, variables.IteratorVariable):
+            return variables.EnumerateVariable(
+                iterable, start, mutable_local=MutableLocal()
+            )
+
+        return self._call_enumerate_polyfill(
+            tx, iterable, variables.ConstantVariable.create(start)
+        )
 
     def call_len(self, tx: "InstructionTranslator", *args, **kwargs):
         return args[0].call_method(tx, "__len__", args[1:], kwargs)

torch/_dynamo/variables/iter.py

@@ -480,3 +480,23 @@ def reconstruct(self, codegen):
                 create_instruction("CALL_FUNCTION_EX", arg=0),
             ]
         )
+
+
+class EnumerateVariable(ZipVariable):
+    def __init__(
+        self,
+        iterable: Union[List[VariableTracker], VariableTracker],
+        start: int = 0,
+        **kwargs,
+    ):
+        super().__init__(
+            [CountIteratorVariable(start, mutable_local=MutableLocal()), iterable],
+            **kwargs,
+        )
+
+    def reconstruct(self, codegen):
+        codegen.add_push_null(lambda: codegen.load_import_from("builtins", "enumerate"))
+        codegen(self.iterables[1])
+        assert isinstance(self.iterables[0], CountIteratorVariable)
+        codegen(self.iterables[0].item)
+        codegen.extend_output(codegen.create_call_function_kw(2, ("start",), False))