bpo-26280: Port BINARY_SUBSCR to PEP 659 adaptive interpreter (GH-27043)

7 files changed, 203 insertions, 23 deletions

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
index f5e814d..06a8b1b 100644
--- a/Include/internal/pycore_code.h
+++ b/Include/internal/pycore_code.h
@@ -310,17 +310,18 @@ too_many_cache_misses(_PyAdaptiveEntry *entry) {
     return entry->counter == saturating_zero();
 }
 
-#define BACKOFF 64
+#define ADAPTIVE_CACHE_BACKOFF 64
 
 static inline void
 cache_backoff(_PyAdaptiveEntry *entry) {
-    entry->counter = BACKOFF;
+    entry->counter = ADAPTIVE_CACHE_BACKOFF;
 }
 
 /* Specialization functions */
 
 int _Py_Specialize_LoadAttr(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name, SpecializedCacheEntry *cache);
 int _Py_Specialize_LoadGlobal(PyObject *globals, PyObject *builtins, _Py_CODEUNIT *instr, PyObject *name, SpecializedCacheEntry *cache);
+int _Py_Specialize_BinarySubscr(PyObject *sub, PyObject *container, _Py_CODEUNIT *instr);
 
 #define SPECIALIZATION_STATS 0
 #define SPECIALIZATION_STATS_DETAILED 0
diff --git a/Include/opcode.h b/Include/opcode.h
index 7f8376f..7bebb87 100644
--- a/Include/opcode.h
+++ b/Include/opcode.h
@@ -136,15 +136,19 @@ extern "C" {
 #define DICT_MERGE              164
 #define DICT_UPDATE             165
 #define CALL_METHOD_KW          166
-#define JUMP_ABSOLUTE_QUICK       7
-#define LOAD_ATTR_ADAPTIVE        8
-#define LOAD_ATTR_SPLIT_KEYS     13
-#define LOAD_ATTR_WITH_HINT      14
-#define LOAD_ATTR_SLOT           18
-#define LOAD_ATTR_MODULE         21
-#define LOAD_GLOBAL_ADAPTIVE     36
-#define LOAD_GLOBAL_MODULE       38
-#define LOAD_GLOBAL_BUILTIN      39
+#define BINARY_SUBSCR_ADAPTIVE    7
+#define BINARY_SUBSCR_LIST_INT    8
+#define BINARY_SUBSCR_TUPLE_INT  13
+#define BINARY_SUBSCR_DICT       14
+#define JUMP_ABSOLUTE_QUICK      18
+#define LOAD_ATTR_ADAPTIVE       21
+#define LOAD_ATTR_SPLIT_KEYS     36
+#define LOAD_ATTR_WITH_HINT      38
+#define LOAD_ATTR_SLOT           39
+#define LOAD_ATTR_MODULE         40
+#define LOAD_GLOBAL_ADAPTIVE     41
+#define LOAD_GLOBAL_MODULE       42
+#define LOAD_GLOBAL_BUILTIN      43
 #ifdef NEED_OPCODE_JUMP_TABLES
 static uint32_t _PyOpcode_RelativeJump[8] = {
     0U,
diff --git a/Lib/opcode.py b/Lib/opcode.py
index 7e5916a..7ba1519 100644
--- a/Lib/opcode.py
+++ b/Lib/opcode.py
@@ -220,6 +220,10 @@ def_op('CALL_METHOD_KW', 166)
 del def_op, name_op, jrel_op, jabs_op
 
 _specialized_instructions = [
+    "BINARY_SUBSCR_ADAPTIVE",
+    "BINARY_SUBSCR_LIST_INT",
+    "BINARY_SUBSCR_TUPLE_INT",
+    "BINARY_SUBSCR_DICT",
     "JUMP_ABSOLUTE_QUICK",
     "LOAD_ATTR_ADAPTIVE",
     "LOAD_ATTR_SPLIT_KEYS",
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-07-14-10-31-10.bpo-26280.cgpM4B.rst b/Misc/NEWS.d/next/Core and Builtins/2021-07-14-10-31-10.bpo-26280.cgpM4B.rst
new file mode 100644
index 0000000..cb561e7
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-07-14-10-31-10.bpo-26280.cgpM4B.rst
@@ -0,0 +1,9 @@
+Implement adaptive specialization for BINARY_SUBSCR

+

+ Three specialized forms of BINARY_SUBSCR are added:

+

+ * BINARY_SUBSCR_LIST_INT

+

+ * BINARY_SUBSCR_TUPLE_INT

+

+ * BINARY_SUBSCR_DICT
\ No newline at end of file
diff --git a/Python/ceval.c b/Python/ceval.c
index 2218405..1467c12 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -15,6 +15,7 @@
 #include "pycore_ceval.h"         // _PyEval_SignalAsyncExc()
 #include "pycore_code.h"
 #include "pycore_initconfig.h"    // _PyStatus_OK()
+#include "pycore_long.h"          // _PyLong_GetZero()
 #include "pycore_object.h"        // _PyObject_GC_TRACK()
 #include "pycore_moduleobject.h"
 #include "pycore_pyerrors.h"      // _PyErr_Fetch()
@@ -1398,6 +1399,8 @@ eval_frame_handle_pending(PyThreadState *tstate)
 
 #define DEOPT_IF(cond, instname) if (cond) { goto instname ## _miss; }
 
+#define UPDATE_PREV_INSTR_OPARG(instr, oparg) ((uint8_t*)(instr))[-1] = (oparg)
+
 #define GLOBALS() specials[FRAME_SPECIALS_GLOBALS_OFFSET]
 #define BUILTINS() specials[FRAME_SPECIALS_BUILTINS_OFFSET]
 #define LOCALS() specials[FRAME_SPECIALS_LOCALS_OFFSET]
@@ -1913,6 +1916,8 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
         }
 
         case TARGET(BINARY_SUBSCR): {
+            PREDICTED(BINARY_SUBSCR);
+            STAT_INC(BINARY_SUBSCR, unquickened);
             PyObject *sub = POP();
             PyObject *container = TOP();
             PyObject *res = PyObject_GetItem(container, sub);
@@ -1924,6 +1929,91 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             DISPATCH();
         }
 
+        case TARGET(BINARY_SUBSCR_ADAPTIVE): {
+            if (oparg == 0) {
+                PyObject *sub = TOP();
+                PyObject *container = SECOND();
+                next_instr--;
+                if (_Py_Specialize_BinarySubscr(container, sub, next_instr) < 0) {
+                    goto error;
+                }
+                DISPATCH();
+            }
+            else {
+                STAT_INC(BINARY_SUBSCR, deferred);
+                // oparg is the adaptive cache counter
+                UPDATE_PREV_INSTR_OPARG(next_instr, oparg - 1);
+                assert(_Py_OPCODE(next_instr[-1]) == BINARY_SUBSCR_ADAPTIVE);
+                assert(_Py_OPARG(next_instr[-1]) == oparg - 1);
+                JUMP_TO_INSTRUCTION(BINARY_SUBSCR);
+            }
+        }
+
+        case TARGET(BINARY_SUBSCR_LIST_INT): {
+            PyObject *sub = TOP();
+            PyObject *list = SECOND();
+            DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
+            DEOPT_IF(!PyList_CheckExact(list), BINARY_SUBSCR);
+
+            // Deopt unless 0 <= sub < PyList_Size(list)
+            Py_ssize_t signed_magnitude = Py_SIZE(sub);
+            DEOPT_IF(((size_t)signed_magnitude) > 1, BINARY_SUBSCR);
+            assert(((PyLongObject *)_PyLong_GetZero())->ob_digit[0] == 0);
+            Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
+            DEOPT_IF(index >= PyList_GET_SIZE(list), BINARY_SUBSCR);
+
+            STAT_INC(BINARY_SUBSCR, hit);
+            PyObject *res = PyList_GET_ITEM(list, index);
+            assert(res != NULL);
+            Py_INCREF(res);
+            STACK_SHRINK(1);
+            Py_DECREF(sub);
+            SET_TOP(res);
+            Py_DECREF(list);
+            DISPATCH();
+        }
+
+        case TARGET(BINARY_SUBSCR_TUPLE_INT): {
+            PyObject *sub = TOP();
+            PyObject *tuple = SECOND();
+            DEOPT_IF(!PyLong_CheckExact(sub), BINARY_SUBSCR);
+            DEOPT_IF(!PyTuple_CheckExact(tuple), BINARY_SUBSCR);
+
+            // Deopt unless 0 <= sub < PyTuple_Size(list)
+            Py_ssize_t signed_magnitude = Py_SIZE(sub);
+            DEOPT_IF(((size_t)signed_magnitude) > 1, BINARY_SUBSCR);
+            assert(((PyLongObject *)_PyLong_GetZero())->ob_digit[0] == 0);
+            Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
+            DEOPT_IF(index >= PyTuple_GET_SIZE(tuple), BINARY_SUBSCR);
+
+            STAT_INC(BINARY_SUBSCR, hit);
+            PyObject *res = PyTuple_GET_ITEM(tuple, index);
+            assert(res != NULL);
+            Py_INCREF(res);
+            STACK_SHRINK(1);
+            Py_DECREF(sub);
+            SET_TOP(res);
+            Py_DECREF(tuple);
+            DISPATCH();
+        }
+
+        case TARGET(BINARY_SUBSCR_DICT): {
+            PyObject *dict = SECOND();
+            DEOPT_IF(!PyDict_CheckExact(SECOND()), BINARY_SUBSCR);
+            STAT_INC(BINARY_SUBSCR, hit);
+            PyObject *sub = TOP();
+            PyObject *res = PyDict_GetItemWithError(dict, sub);
+            if (res == NULL) {
+                goto binary_subscr_dict_error;
+            }
+            Py_INCREF(res);
+            STACK_SHRINK(1);
+            Py_DECREF(sub);
+            SET_TOP(res);
+            Py_DECREF(dict);
+            DISPATCH();
+        }
+
         case TARGET(BINARY_LSHIFT): {
             PyObject *right = POP();
             PyObject *left = TOP();
@@ -4327,8 +4417,34 @@ opname ## _miss: \
         JUMP_TO_INSTRUCTION(opname); \
     }
 
+#define MISS_WITH_OPARG_COUNTER(opname) \
+opname ## _miss: \
+    { \
+        STAT_INC(opname, miss); \
+        uint8_t oparg = saturating_decrement(_Py_OPARG(next_instr[-1])); \
+        UPDATE_PREV_INSTR_OPARG(next_instr, oparg); \
+        assert(_Py_OPARG(next_instr[-1]) == oparg); \
+        if (oparg == saturating_zero()) /* too many cache misses */ { \
+            oparg = ADAPTIVE_CACHE_BACKOFF; \
+            next_instr[-1] = _Py_MAKECODEUNIT(opname ## _ADAPTIVE, oparg); \
+            STAT_INC(opname, deopt); \
+        } \
+        JUMP_TO_INSTRUCTION(opname); \
+    }
+
 MISS_WITH_CACHE(LOAD_ATTR)
 MISS_WITH_CACHE(LOAD_GLOBAL)
+MISS_WITH_OPARG_COUNTER(BINARY_SUBSCR)
+
+binary_subscr_dict_error:
+        {
+            PyObject *sub = POP();
+            if (!_PyErr_Occurred(tstate)) {
+                _PyErr_SetKeyError(sub);
+            }
+            Py_DECREF(sub);
+            goto error;
+        }
 
 error:
         /* Double-check exception status. */
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index ecc95da..d88c766 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -6,21 +6,21 @@ static void *opcode_targets[256] = {
     &&TARGET_DUP_TOP,
     &&TARGET_DUP_TOP_TWO,
     &&TARGET_ROT_FOUR,
-    &&TARGET_JUMP_ABSOLUTE_QUICK,
-    &&TARGET_LOAD_ATTR_ADAPTIVE,
+    &&TARGET_BINARY_SUBSCR_ADAPTIVE,
+    &&TARGET_BINARY_SUBSCR_LIST_INT,
     &&TARGET_NOP,
     &&TARGET_UNARY_POSITIVE,
     &&TARGET_UNARY_NEGATIVE,
     &&TARGET_UNARY_NOT,
-    &&TARGET_LOAD_ATTR_SPLIT_KEYS,
-    &&TARGET_LOAD_ATTR_WITH_HINT,
+    &&TARGET_BINARY_SUBSCR_TUPLE_INT,
+    &&TARGET_BINARY_SUBSCR_DICT,
     &&TARGET_UNARY_INVERT,
     &&TARGET_BINARY_MATRIX_MULTIPLY,
     &&TARGET_INPLACE_MATRIX_MULTIPLY,
-    &&TARGET_LOAD_ATTR_SLOT,
+    &&TARGET_JUMP_ABSOLUTE_QUICK,
     &&TARGET_BINARY_POWER,
     &&TARGET_BINARY_MULTIPLY,
-    &&TARGET_LOAD_ATTR_MODULE,
+    &&TARGET_LOAD_ATTR_ADAPTIVE,
     &&TARGET_BINARY_MODULO,
     &&TARGET_BINARY_ADD,
     &&TARGET_BINARY_SUBTRACT,
@@ -35,8 +35,12 @@ static void *opcode_targets[256] = {
     &&TARGET_MATCH_KEYS,
     &&TARGET_COPY_DICT_WITHOUT_KEYS,
     &&TARGET_PUSH_EXC_INFO,
-    &&TARGET_LOAD_GLOBAL_ADAPTIVE,
+    &&TARGET_LOAD_ATTR_SPLIT_KEYS,
     &&TARGET_POP_EXCEPT_AND_RERAISE,
+    &&TARGET_LOAD_ATTR_WITH_HINT,
+    &&TARGET_LOAD_ATTR_SLOT,
+    &&TARGET_LOAD_ATTR_MODULE,
+    &&TARGET_LOAD_GLOBAL_ADAPTIVE,
     &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
     &&_unknown_opcode,
@@ -44,10 +48,6 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_WITH_EXCEPT_START,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
diff --git a/Python/specialize.c b/Python/specialize.c
index 3277c6b..5ebe596 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -78,6 +78,7 @@ _Py_PrintSpecializationStats(void)
     printf("Specialization stats:\n");
     print_stats(&_specialization_stats[LOAD_ATTR], "load_attr");
     print_stats(&_specialization_stats[LOAD_GLOBAL], "load_global");
+    print_stats(&_specialization_stats[BINARY_SUBSCR], "binary_subscr");
 }
 
 #if SPECIALIZATION_STATS_DETAILED
@@ -162,12 +163,14 @@ get_cache_count(SpecializedCacheOrInstruction *quickened) {
 static uint8_t adaptive_opcodes[256] = {
     [LOAD_ATTR] = LOAD_ATTR_ADAPTIVE,
     [LOAD_GLOBAL] = LOAD_GLOBAL_ADAPTIVE,
+    [BINARY_SUBSCR] = BINARY_SUBSCR_ADAPTIVE,
 };
 
 /* The number of cache entries required for a "family" of instructions. */
 static uint8_t cache_requirements[256] = {
     [LOAD_ATTR] = 2, /* _PyAdaptiveEntry and _PyLoadAttrCache */
     [LOAD_GLOBAL] = 2, /* _PyAdaptiveEntry and _PyLoadGlobalCache */
+    [BINARY_SUBSCR] = 0,
 };
 
 /* Return the oparg for the cache_offset and instruction index.
@@ -251,7 +254,6 @@ optimize(SpecializedCacheOrInstruction *quickened, int len)
                 previous_opcode = opcode;
                 continue;
             }
-            instructions[i] = _Py_MAKECODEUNIT(adaptive_opcode, new_oparg);
             previous_opcode = adaptive_opcode;
             int entries_needed = cache_requirements[opcode];
             if (entries_needed) {
@@ -261,7 +263,11 @@ optimize(SpecializedCacheOrInstruction *quickened, int len)
                     _GetSpecializedCacheEntry(instructions, cache0_offset);
                 cache->adaptive.original_oparg = oparg;
                 cache->adaptive.counter = 0;
+            } else {
+                // oparg is the adaptive cache counter
+                new_oparg = 0;
             }
+            instructions[i] = _Py_MAKECODEUNIT(adaptive_opcode, new_oparg);
         }
         else {
             /* Super instructions don't use the cache,
@@ -637,3 +643,43 @@ success:
     cache0->counter = saturating_start();
     return 0;
 }
+
+int
+_Py_Specialize_BinarySubscr(
+     PyObject *container, PyObject *sub, _Py_CODEUNIT *instr)
+{
+    PyTypeObject *container_type = Py_TYPE(container);
+    if (container_type == &PyList_Type) {
+        if (PyLong_CheckExact(sub)) {
+            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_LIST_INT, saturating_start());
+            goto success;
+        } else {
+            SPECIALIZATION_FAIL(BINARY_SUBSCR, Py_TYPE(container), sub, "list; non-integer subscr");
+        }
+    }
+    if (container_type == &PyTuple_Type) {
+        if (PyLong_CheckExact(sub)) {
+            *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_TUPLE_INT, saturating_start());
+            goto success;
+        } else {
+            SPECIALIZATION_FAIL(BINARY_SUBSCR, Py_TYPE(container), sub, "tuple; non-integer subscr");
+        }
+    }
+    if (container_type == &PyDict_Type) {
+        *instr = _Py_MAKECODEUNIT(BINARY_SUBSCR_DICT, saturating_start());
+        goto success;
+    }
+
+    SPECIALIZATION_FAIL(BINARY_SUBSCR, Py_TYPE(container), sub, "not list|tuple|dict");
+    goto fail;
+fail:
+    STAT_INC(BINARY_SUBSCR, specialization_failure);
+    assert(!PyErr_Occurred());
+    *instr = _Py_MAKECODEUNIT(_Py_OPCODE(*instr), ADAPTIVE_CACHE_BACKOFF);
+    return 0;
+success:
+    STAT_INC(BINARY_SUBSCR, specialization_success);
+    assert(!PyErr_Occurred());
+    return 0;
+}
+