gh-100239: specialize BINARY_OP/SUBSCR for list-slice (#132626)

6 files changed, 162 insertions, 2 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 496a53d..9ddc3f3 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -519,6 +519,11 @@ dummy_func(
             EXIT_IF(!PyList_CheckExact(o));
         }
 
+        op(_GUARD_TOS_SLICE, (tos -- tos)) {
+            PyObject *o = PyStackRef_AsPyObjectBorrow(tos);
+            EXIT_IF(!PySlice_Check(o));
+        }
+
         macro(TO_BOOL_LIST) = _GUARD_TOS_LIST + unused/1 + unused/2 + _TO_BOOL_LIST;
 
         op(_TO_BOOL_LIST, (value -- res)) {
@@ -591,6 +596,7 @@ dummy_func(
             BINARY_OP_SUBTRACT_FLOAT,
             BINARY_OP_ADD_UNICODE,
             BINARY_OP_SUBSCR_LIST_INT,
+            BINARY_OP_SUBSCR_LIST_SLICE,
             BINARY_OP_SUBSCR_TUPLE_INT,
             BINARY_OP_SUBSCR_STR_INT,
             BINARY_OP_SUBSCR_DICT,
@@ -900,6 +906,24 @@ dummy_func(
             DECREF_INPUTS();
         }
 
+        macro(BINARY_OP_SUBSCR_LIST_SLICE) =
+            _GUARD_TOS_SLICE + _GUARD_NOS_LIST + unused/5 + _BINARY_OP_SUBSCR_LIST_SLICE;
+
+        op(_BINARY_OP_SUBSCR_LIST_SLICE, (list_st, sub_st -- res)) {
+            PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
+            PyObject *list = PyStackRef_AsPyObjectBorrow(list_st);
+
+            assert(PySlice_Check(sub));
+            assert(PyList_CheckExact(list));
+
+            PyObject *res_o = _PyList_SliceSubscript(list, sub);
+            DEOPT_IF(res_o == NULL);
+            STAT_INC(BINARY_OP, hit);
+            res = PyStackRef_FromPyObjectSteal(res_o);
+            STAT_INC(BINARY_OP, hit);
+            DECREF_INPUTS();
+        }
+
         macro(BINARY_OP_SUBSCR_STR_INT) =
             _GUARD_TOS_INT + _GUARD_NOS_UNICODE + unused/5 + _BINARY_OP_SUBSCR_STR_INT;
 
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index b37a875..82c37bc 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -740,6 +740,17 @@
             break;
         }
 
+        case _GUARD_TOS_SLICE: {
+            _PyStackRef tos;
+            tos = stack_pointer[-1];
+            PyObject *o = PyStackRef_AsPyObjectBorrow(tos);
+            if (!PySlice_Check(o)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            break;
+        }
+
         case _TO_BOOL_LIST: {
             _PyStackRef value;
             _PyStackRef res;
@@ -1324,6 +1335,41 @@
             break;
         }
 
+        case _BINARY_OP_SUBSCR_LIST_SLICE: {
+            _PyStackRef sub_st;
+            _PyStackRef list_st;
+            _PyStackRef res;
+            sub_st = stack_pointer[-1];
+            list_st = stack_pointer[-2];
+            PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
+            PyObject *list = PyStackRef_AsPyObjectBorrow(list_st);
+            assert(PySlice_Check(sub));
+            assert(PyList_CheckExact(list));
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            PyObject *res_o = _PyList_SliceSubscript(list, sub);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            if (res_o == NULL) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            STAT_INC(BINARY_OP, hit);
+            res = PyStackRef_FromPyObjectSteal(res_o);
+            STAT_INC(BINARY_OP, hit);
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            _PyStackRef tmp = list_st;
+            list_st = res;
+            stack_pointer[-2] = list_st;
+            PyStackRef_CLOSE(tmp);
+            tmp = sub_st;
+            sub_st = PyStackRef_NULL;
+            stack_pointer[-1] = sub_st;
+            PyStackRef_CLOSE(tmp);
+            stack_pointer = _PyFrame_GetStackPointer(frame);
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _BINARY_OP_SUBSCR_STR_INT: {
             _PyStackRef sub_st;
             _PyStackRef str_st;
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 6841fbc..a432603 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -763,6 +763,78 @@
             DISPATCH();
         }
 
+        TARGET(BINARY_OP_SUBSCR_LIST_SLICE) {
+            #if Py_TAIL_CALL_INTERP
+            int opcode = BINARY_OP_SUBSCR_LIST_SLICE;
+            (void)(opcode);
+            #endif
+            _Py_CODEUNIT* const this_instr = next_instr;
+            (void)this_instr;
+            frame->instr_ptr = next_instr;
+            next_instr += 6;
+            INSTRUCTION_STATS(BINARY_OP_SUBSCR_LIST_SLICE);
+            static_assert(INLINE_CACHE_ENTRIES_BINARY_OP == 5, "incorrect cache size");
+            _PyStackRef tos;
+            _PyStackRef nos;
+            _PyStackRef list_st;
+            _PyStackRef sub_st;
+            _PyStackRef res;
+            // _GUARD_TOS_SLICE
+            {
+                tos = stack_pointer[-1];
+                PyObject *o = PyStackRef_AsPyObjectBorrow(tos);
+                if (!PySlice_Check(o)) {
+                    UPDATE_MISS_STATS(BINARY_OP);
+                    assert(_PyOpcode_Deopt[opcode] == (BINARY_OP));
+                    JUMP_TO_PREDICTED(BINARY_OP);
+                }
+            }
+            // _GUARD_NOS_LIST
+            {
+                nos = stack_pointer[-2];
+                PyObject *o = PyStackRef_AsPyObjectBorrow(nos);
+                if (!PyList_CheckExact(o)) {
+                    UPDATE_MISS_STATS(BINARY_OP);
+                    assert(_PyOpcode_Deopt[opcode] == (BINARY_OP));
+                    JUMP_TO_PREDICTED(BINARY_OP);
+                }
+            }
+            /* Skip 5 cache entries */
+            // _BINARY_OP_SUBSCR_LIST_SLICE
+            {
+                sub_st = tos;
+                list_st = nos;
+                PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
+                PyObject *list = PyStackRef_AsPyObjectBorrow(list_st);
+                assert(PySlice_Check(sub));
+                assert(PyList_CheckExact(list));
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                PyObject *res_o = _PyList_SliceSubscript(list, sub);
+                stack_pointer = _PyFrame_GetStackPointer(frame);
+                if (res_o == NULL) {
+                    UPDATE_MISS_STATS(BINARY_OP);
+                    assert(_PyOpcode_Deopt[opcode] == (BINARY_OP));
+                    JUMP_TO_PREDICTED(BINARY_OP);
+                }
+                STAT_INC(BINARY_OP, hit);
+                res = PyStackRef_FromPyObjectSteal(res_o);
+                STAT_INC(BINARY_OP, hit);
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                _PyStackRef tmp = list_st;
+                list_st = res;
+                stack_pointer[-2] = list_st;
+                PyStackRef_CLOSE(tmp);
+                tmp = sub_st;
+                sub_st = PyStackRef_NULL;
+                stack_pointer[-1] = sub_st;
+                PyStackRef_CLOSE(tmp);
+                stack_pointer = _PyFrame_GetStackPointer(frame);
+                stack_pointer += -1;
+                assert(WITHIN_STACK_BOUNDS());
+            }
+            DISPATCH();
+        }
+
         TARGET(BINARY_OP_SUBSCR_STR_INT) {
             #if Py_TAIL_CALL_INTERP
             int opcode = BINARY_OP_SUBSCR_STR_INT;
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 8c2dfb4..8af445d 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -138,6 +138,7 @@ static void *opcode_targets[256] = {
     &&TARGET_BINARY_OP_SUBSCR_DICT,
     &&TARGET_BINARY_OP_SUBSCR_GETITEM,
     &&TARGET_BINARY_OP_SUBSCR_LIST_INT,
+    &&TARGET_BINARY_OP_SUBSCR_LIST_SLICE,
     &&TARGET_BINARY_OP_SUBSCR_STR_INT,
     &&TARGET_BINARY_OP_SUBSCR_TUPLE_INT,
     &&TARGET_BINARY_OP_SUBTRACT_FLOAT,
@@ -233,7 +234,6 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_INSTRUMENTED_END_FOR,
     &&TARGET_INSTRUMENTED_POP_ITER,
     &&TARGET_INSTRUMENTED_END_SEND,
@@ -278,6 +278,7 @@ Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_MULTIPLY_INT(TAIL_CALL
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_DICT(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_GETITEM(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_LIST_INT(TAIL_CALL_PARAMS);
+Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_LIST_SLICE(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_STR_INT(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBSCR_TUPLE_INT(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_BINARY_OP_SUBTRACT_FLOAT(TAIL_CALL_PARAMS);
@@ -516,6 +517,7 @@ static py_tail_call_funcptr INSTRUCTION_TABLE[256] = {
     [BINARY_OP_SUBSCR_DICT] = _TAIL_CALL_BINARY_OP_SUBSCR_DICT,
     [BINARY_OP_SUBSCR_GETITEM] = _TAIL_CALL_BINARY_OP_SUBSCR_GETITEM,
     [BINARY_OP_SUBSCR_LIST_INT] = _TAIL_CALL_BINARY_OP_SUBSCR_LIST_INT,
+    [BINARY_OP_SUBSCR_LIST_SLICE] = _TAIL_CALL_BINARY_OP_SUBSCR_LIST_SLICE,
     [BINARY_OP_SUBSCR_STR_INT] = _TAIL_CALL_BINARY_OP_SUBSCR_STR_INT,
     [BINARY_OP_SUBSCR_TUPLE_INT] = _TAIL_CALL_BINARY_OP_SUBSCR_TUPLE_INT,
     [BINARY_OP_SUBTRACT_FLOAT] = _TAIL_CALL_BINARY_OP_SUBTRACT_FLOAT,
@@ -738,7 +740,6 @@ static py_tail_call_funcptr INSTRUCTION_TABLE[256] = {
     [125] = _TAIL_CALL_UNKNOWN_OPCODE,
     [126] = _TAIL_CALL_UNKNOWN_OPCODE,
     [127] = _TAIL_CALL_UNKNOWN_OPCODE,
-    [211] = _TAIL_CALL_UNKNOWN_OPCODE,
     [212] = _TAIL_CALL_UNKNOWN_OPCODE,
     [213] = _TAIL_CALL_UNKNOWN_OPCODE,
     [214] = _TAIL_CALL_UNKNOWN_OPCODE,
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index b8391c5..622c0b5 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -215,6 +215,10 @@
             break;
         }
 
+        case _GUARD_TOS_SLICE: {
+            break;
+        }
+
         case _TO_BOOL_LIST: {
             JitOptSymbol *value;
             JitOptSymbol *res;
@@ -587,6 +591,15 @@
             break;
         }
 
+        case _BINARY_OP_SUBSCR_LIST_SLICE: {
+            JitOptSymbol *res;
+            res = sym_new_not_null(ctx);
+            stack_pointer[-2] = res;
+            stack_pointer += -1;
+            assert(WITHIN_STACK_BOUNDS());
+            break;
+        }
+
         case _BINARY_OP_SUBSCR_STR_INT: {
             JitOptSymbol *res;
             res = sym_new_type(ctx, &PyUnicode_Type);
diff --git a/Python/specialize.c b/Python/specialize.c
index f4b0691..59ec9a4 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -2657,6 +2657,10 @@ _Py_Specialize_BinaryOp(_PyStackRef lhs_st, _PyStackRef rhs_st, _Py_CODEUNIT *in
                 specialize(instr, BINARY_OP_SUBSCR_DICT);
                 return;
             }
+            if (PyList_CheckExact(lhs) && PySlice_Check(rhs)) {
+                specialize(instr, BINARY_OP_SUBSCR_LIST_SLICE);
+                return;
+            }
             unsigned int tp_version;
             PyTypeObject *container_type = Py_TYPE(lhs);
             PyObject *descriptor = _PyType_LookupRefAndVersion(container_type, &_Py_ID(__getitem__), &tp_version);