Specialize STORE_FAST by reference and type in JIT

python · Fidget-Spinner · Jun 23, 2025 · Jun 20, 2025 · Jun 20, 2025 · Jun 20, 2025
commit 4b415f58dabada7ab08af9e404c6d7406a3f3d0c
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Include/internal/pycore_uop_ids.h b/Include/internal/pycore_uop_ids.h
diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h
diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
@@ -2349,6 +2349,66 @@ def testfunc(n):
         assert ex is not None
         """))
 
+    def test_store_fast_pop_top_specialize_immortal(self):
+        def testfunc(n):
+            for _ in range(n):
+                x = None # _POP_TOP, as x's type is not yet known by optimizer.
+                x = None # _POP_TOP_NOP, as x = None
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_POP_TOP_NOP", uops)
+
+    def test_store_fast_pop_top_specialize_int(self):
+        def testfunc(n):
+            y = int(1e6) # Big number so no int caching
+            for _ in range(n):
+                x = y + y  # _POP_TOP, as x's type is not yet known by optimizer.
+                x = None # _POP_TOP_INT, as x = int
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_POP_TOP_INT", uops)
+
+    def test_store_fast_pop_top_specialize_float(self):
+        def testfunc(n):
+            y = 1.0
+            for _ in range(n):
+                x = y + y  # _POP_TOP, as x's type is not yet known by optimizer.
+                x = None # _POP_TOP_FLOAT, as x = int
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_POP_TOP_FLOAT", uops)
+
+    def test_store_fast_pop_top_specialize_unicode(self):
+        def testfunc(n):
+            y = "hi"
+            for _ in range(n):
+                x = y + y  # _POP_TOP, as x's type is not yet known by optimizer.
+                x = None # _POP_TOP_STR, as x = int
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_POP_TOP_UNICODE", uops)
+
+
 
 def global_identity(x):
     return x

@@ -306,13 +306,14 @@ dummy_func(
             value = PyStackRef_FromPyObjectBorrow(obj);
         }
 
-        replicate(8) inst(STORE_FAST, (value --)) {
-            _PyStackRef tmp = GETLOCAL(oparg);
+        replicate(8) op(_SWAP_FAST, (value -- trash)) {
+            trash = GETLOCAL(oparg);
             GETLOCAL(oparg) = value;
             DEAD(value);
-            PyStackRef_XCLOSE(tmp);
         }
 
+        macro(STORE_FAST) = _SWAP_FAST + POP_TOP;
+
         pseudo(STORE_FAST_MAYBE_NULL, (unused --)) = {
             STORE_FAST,
         };
@@ -344,6 +345,31 @@ dummy_func(
             PyStackRef_XCLOSE(value);
         }
 
+        op(_POP_TOP_NOP, (value --)) {
+            // TODO (gh-134584): Consider moving this to a function pointer table and replicate.
+            assert(!PyStackRef_RefcountOnObject(value) ||
+                _Py_IsImmortal((PyStackRef_AsPyObjectBorrow(value))));
+            DEAD(value);
+        }
+
+        op(_POP_TOP_INT, (value --)) {
+            // TODO (gh-134584): Consider moving this to a function pointer table and replicate.
+            assert(PyLong_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
+            PyStackRef_CLOSE_SPECIALIZED(value, _PyLong_ExactDealloc);
+        }
+
+        op(_POP_TOP_FLOAT, (value --)) {
+            // TODO (gh-134584): Consider moving this to a function pointer table and replicate.
+            assert(PyFloat_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
+            PyStackRef_CLOSE_SPECIALIZED(value, _PyFloat_ExactDealloc);
+        }
+
+        op(_POP_TOP_UNICODE, (value --)) {
+            // TODO (gh-134584): Consider moving this to a function pointer table and replicate.
+            assert(PyUnicode_CheckExact(PyStackRef_AsPyObjectBorrow(value)));
+            PyStackRef_CLOSE_SPECIALIZED(value, _PyUnicode_ExactDealloc);
+        }
+
         tier2 op(_POP_TWO, (nos, tos --)) {
             PyStackRef_CLOSE(tos);
             PyStackRef_CLOSE(nos);