[Dynamo] Guard serialization for RANGE_ITERATOR_MATCH (pytorch#152872)

Tests serialization for RANGE_ITERATOR_MATCH; includes no non-test changes.

This PR handles iterator exhaustion issues by utilizing the janky solution from pytorch#152865; it passes a function to generate kwargs and `frame_state.f_locals` is updated with fresh iterators through a second kwarg generation pass after initial tracing.

Pull Request resolved: pytorch#152872
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#152725, pytorch#152727, pytorch#152728, pytorch#152730, pytorch#152865

Author: jbschlosser

test/dynamo/test_guard_serialization.py

@@ -737,6 +737,48 @@ def _gen_kwargs(tup=tup, x=x):
             ref, loaded, {"t": iter((1, 2, 3, 4)), "x": torch.randn(4)}, False
         )
 
+    def test_range_iterator_match(self):
+        def fn(x, r):
+            y = x
+            for val in r:
+                y = x + val
+            return y
+
+        x = torch.randn(3)
+
+        def _gen_kwargs(x=x):
+            return {"x": x, "r": iter(range(2, 15, 3))}
+
+        ref, loaded = self._test_serialization(
+            "RANGE_ITERATOR_MATCH", fn, _gen_fn=_gen_kwargs
+        )
+
+        # same range
+        self._test_check_fn(ref, loaded, {"x": x, "r": iter(range(2, 15, 3))}, True)
+        self._test_check_fn(
+            ref, loaded, {"x": torch.randn(4), "r": iter(range(2, 15, 3))}, True
+        )
+        # equivalent even with different end
+        self._test_check_fn(ref, loaded, {"x": x, "r": iter(range(2, 16, 3))}, True)
+        self._test_check_fn(
+            ref, loaded, {"x": torch.randn(4), "r": iter(range(2, 16, 3))}, True
+        )
+        # different start
+        self._test_check_fn(ref, loaded, {"x": x, "r": iter(range(1, 15, 3))}, False)
+        self._test_check_fn(
+            ref, loaded, {"x": torch.randn(4), "r": iter(range(1, 15, 3))}, False
+        )
+        # different end resulting in different values
+        self._test_check_fn(ref, loaded, {"x": x, "r": iter(range(2, 18, 3))}, False)
+        self._test_check_fn(
+            ref, loaded, {"x": torch.randn(4), "r": iter(range(2, 18, 3))}, False
+        )
+        # different step
+        self._test_check_fn(ref, loaded, {"x": x, "r": iter(range(2, 15, 4))}, False)
+        self._test_check_fn(
+            ref, loaded, {"x": torch.randn(4), "r": iter(range(2, 15, 4))}, False
+        )
+
     def test_dict_version(self):
         def fn(x):
             return pytree.tree_leaves(x)[0] + 1