gh-106529: Split FOR_ITER_RANGE into uops (#106638)

For an example of what this does for Tier 1 and Tier 2, see #106529 (comment)
python · Jul 12, 2023 · dd1884d · dd1884d
1 parent 7f55f58
commit dd1884d
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 28 deletions.
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
@@ -2443,7 +2443,6 @@ def testfunc(x):
                 i += 1
 
         opt = _testinternalcapi.get_uop_optimizer()
-
         with temporary_optimizer(opt):
             testfunc(1000)
 
@@ -2580,13 +2579,33 @@ def testfunc(n):
 
         ex = get_first_executor(testfunc)
         self.assertIsNotNone(ex)
-        # for i, (opname, oparg) in enumerate(ex):
-        #     print(f"{i:4d}: {opname:<20s} {oparg:4d}")
         uops = {opname for opname, _ in ex}
         # Since there is no JUMP_FORWARD instruction,
         # look for indirect evidence: the += operator
         self.assertIn("_BINARY_OP_ADD_INT", uops)
 
+    def test_for_iter_range(self):
+        def testfunc(n):
+            total = 0
+            for i in range(n):
+                total += i
+            return total
+        # import dis; dis.dis(testfunc)
+
+        opt = _testinternalcapi.get_uop_optimizer()
+        with temporary_optimizer(opt):
+            total = testfunc(10)
+            self.assertEqual(total, 45)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        # for i, (opname, oparg) in enumerate(ex):
+        #     print(f"{i:4d}: {opname:<20s} {oparg:3d}")
+        uops = {opname for opname, _ in ex}
+        self.assertIn("_ITER_EXHAUSTED_RANGE", uops)
+        # Verification that the jump goes past END_FOR
+        # is done by manual inspection of the output
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -2451,9 +2451,14 @@ dummy_func(
             // Common case: no jump, leave it to the code generator
         }
 
-        inst(FOR_ITER_RANGE, (unused/1, iter -- iter, next)) {
+        op(_ITER_CHECK_RANGE, (iter -- iter)) {
             _PyRangeIterObject *r = (_PyRangeIterObject *)iter;
             DEOPT_IF(Py_TYPE(r) != &PyRangeIter_Type, FOR_ITER);
+        }
+
+        op(_ITER_JUMP_RANGE, (iter -- iter)) {
+            _PyRangeIterObject *r = (_PyRangeIterObject *)iter;
+            assert(Py_TYPE(r) == &PyRangeIter_Type);
             STAT_INC(FOR_ITER, hit);
             if (r->len <= 0) {
                 STACK_SHRINK(1);
@@ -2463,15 +2468,29 @@ dummy_func(
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
+        }
+
+        // Only used by Tier 2
+        op(_ITER_EXHAUSTED_RANGE, (iter -- iter, exhausted)) {
+            _PyRangeIterObject *r = (_PyRangeIterObject *)iter;
+            assert(Py_TYPE(r) == &PyRangeIter_Type);
+            exhausted = r->len <= 0 ? Py_True : Py_False;
+        }
+
+        op(_ITER_NEXT_RANGE, (iter -- iter, next)) {
+            _PyRangeIterObject *r = (_PyRangeIterObject *)iter;
+            assert(Py_TYPE(r) == &PyRangeIter_Type);
+            assert(r->len > 0);
             long value = r->start;
             r->start = value + r->step;
             r->len--;
             next = PyLong_FromLong(value);
-            if (next == NULL) {
-                goto error;
-            }
+            ERROR_IF(next == NULL, error);
         }
 
+        macro(FOR_ITER_RANGE) =
+            unused/1 + _ITER_CHECK_RANGE + _ITER_JUMP_RANGE + _ITER_NEXT_RANGE;
+
         inst(FOR_ITER_GEN, (unused/1, iter -- iter, unused)) {
             DEOPT_IF(tstate->interp->eval_frame, FOR_ITER);
             PyGenObject *gen = (PyGenObject *)iter;

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/optimizer.c b/Python/optimizer.c
@@ -479,6 +479,28 @@ translate_bytecode_to_trace(
                 break;
             }
 
+            case FOR_ITER_RANGE:
+            {
+                // Assume jump unlikely (can a for-loop exit be likely?)
+                // Reserve 9 entries (4 here, 3 stub, plus SAVE_IP + EXIT_TRACE)
+                if (trace_length + 9 > max_length) {
+                    DPRINTF(1, "Ran out of space for FOR_ITER_RANGE\n");
+                    goto done;
+                }
+                _Py_CODEUNIT *target_instr =  // +1 at the end skips over END_FOR
+                    instr + 1 + _PyOpcode_Caches[_PyOpcode_Deopt[opcode]] + oparg + 1;
+                max_length -= 3;  // Really the start of the stubs
+                ADD_TO_TRACE(_ITER_CHECK_RANGE, 0);
+                ADD_TO_TRACE(_ITER_EXHAUSTED_RANGE, 0);
+                ADD_TO_TRACE(_POP_JUMP_IF_TRUE, max_length);
+                ADD_TO_TRACE(_ITER_NEXT_RANGE, 0);
+
+                ADD_TO_STUB(max_length + 0, POP_TOP, 0);
+                ADD_TO_STUB(max_length + 1, SAVE_IP, INSTR_IP(target_instr, code));
+                ADD_TO_STUB(max_length + 2, EXIT_TRACE, 0);
+                break;
+            }
+
             default:
             {
                 const struct opcode_macro_expansion *expansion = &_PyOpcode_macro_expansion[opcode];
@@ -574,8 +596,8 @@ translate_bytecode_to_trace(
                     }
                 }
             }
-            trace_length += buffer_size - max_length;
         }
+        trace_length += buffer_size - max_length;
         return trace_length;
     }
     else {