From 5c930a26fb78c40929f1b894efee1b07c6d828fd Mon Sep 17 00:00:00 2001
From: "T. Wouters" <thomas@python.org>
Date: Wed, 29 Jan 2025 01:07:56 +0100
Subject: gh-115999: Enable free-threaded specialization of LOAD_CONST
 (#129365)

Enable free-threaded specialization of LOAD_CONST.
---
 Lib/test/test_opcache.py          | 14 ++++++++++++++
 Python/bytecodes.c                | 14 ++++++++++++--
 Python/generated_cases.c.h        | 14 ++++++++++++--
 Tools/cases_generator/analyzer.py |  1 +
 4 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
index 4d7304b..e8ea21f 100644
--- a/Lib/test/test_opcache.py
+++ b/Lib/test/test_opcache.py
@@ -1773,6 +1773,20 @@ class TestSpecializer(TestBase):
         self.assert_specialized(compare_op_str, "COMPARE_OP_STR")
         self.assert_no_opcode(compare_op_str, "COMPARE_OP")
 
+    @cpython_only
+    @requires_specialization_ft
+    def test_load_const(self):
+        def load_const():
+            def unused(): pass
+            # Currently, the empty tuple is immortal, and the otherwise
+            # unused nested function's code object is mortal. This test will
+            # have to use different values if either of that changes.
+            return ()
+
+        load_const()
+        self.assert_specialized(load_const, "LOAD_CONST_IMMORTAL")
+        self.assert_specialized(load_const, "LOAD_CONST_MORTAL")
+        self.assert_no_opcode(load_const, "LOAD_CONST")
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 5f0be8d..7d46351 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -294,11 +294,21 @@ dummy_func(
              * marshalling can intern strings and make them immortal. */
             PyObject *obj = GETITEM(FRAME_CO_CONSTS, oparg);
             value = PyStackRef_FromPyObjectNew(obj);
-#if ENABLE_SPECIALIZATION
+#if ENABLE_SPECIALIZATION_FT
+#ifdef Py_GIL_DISABLED
+            uint8_t expected = LOAD_CONST;
+            if (!_Py_atomic_compare_exchange_uint8(
+                    &this_instr->op.code, &expected,
+                    _Py_IsImmortal(obj) ? LOAD_CONST_IMMORTAL : LOAD_CONST_MORTAL)) {
+                // We might lose a race with instrumentation, which we don't care about.
+                assert(expected >= MIN_INSTRUMENTED_OPCODE);
+            }
+#else
             if (this_instr->op.code == LOAD_CONST) {
                 this_instr->op.code = _Py_IsImmortal(obj) ? LOAD_CONST_IMMORTAL : LOAD_CONST_MORTAL;
             }
 #endif
+#endif
         }
 
         inst(LOAD_CONST_MORTAL, (-- value)) {
@@ -2558,7 +2568,7 @@ dummy_func(
             }
             OPCODE_DEFERRED_INC(COMPARE_OP);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-            #endif  /* ENABLE_SPECIALIZATION */
+            #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
         op(_COMPARE_OP, (left, right -- res)) {
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index ad044e6..5dd2f37 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -3318,7 +3318,7 @@
                 }
                 OPCODE_DEFERRED_INC(COMPARE_OP);
                 ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-                #endif  /* ENABLE_SPECIALIZATION */
+                #endif  /* ENABLE_SPECIALIZATION_FT */
             }
             // _COMPARE_OP
             {
@@ -6035,11 +6035,21 @@
              * marshalling can intern strings and make them immortal. */
             PyObject *obj = GETITEM(FRAME_CO_CONSTS, oparg);
             value = PyStackRef_FromPyObjectNew(obj);
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
+            #ifdef Py_GIL_DISABLED
+            uint8_t expected = LOAD_CONST;
+            if (!_Py_atomic_compare_exchange_uint8(
+                    &this_instr->op.code, &expected,
+                    _Py_IsImmortal(obj) ? LOAD_CONST_IMMORTAL : LOAD_CONST_MORTAL)) {
+                // We might lose a race with instrumentation, which we don't care about.
+                assert(expected >= MIN_INSTRUMENTED_OPCODE);
+            }
+            #else
             if (this_instr->op.code == LOAD_CONST) {
                 this_instr->op.code = _Py_IsImmortal(obj) ? LOAD_CONST_IMMORTAL : LOAD_CONST_MORTAL;
             }
             #endif
+            #endif
             stack_pointer[0] = value;
             stack_pointer += 1;
             assert(WITHIN_STACK_BOUNDS());
diff --git a/Tools/cases_generator/analyzer.py b/Tools/cases_generator/analyzer.py
index bc9c42e..b9293ff 100644
--- a/Tools/cases_generator/analyzer.py
+++ b/Tools/cases_generator/analyzer.py
@@ -634,6 +634,7 @@ NON_ESCAPING_FUNCTIONS = (
     "_Py_STR",
     "_Py_TryIncrefCompare",
     "_Py_TryIncrefCompareStackRef",
+    "_Py_atomic_compare_exchange_uint8",
     "_Py_atomic_load_ptr_acquire",
     "_Py_atomic_load_uintptr_relaxed",
     "_Py_set_eval_breaker_bit",
-- 
cgit v0.12