feat: granular type checks (#2)

* feat: More granular type checks * fix: update bytecode in tests, and make tests more informative
oraluben · May 28, 2023 · 9cba136 · 9cba136
1 parent 7458aff
commit 9cba136
Show file tree

Hide file tree

Showing 10 changed files with 487 additions and 458 deletions.
diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
diff --git a/Include/opcode.h b/Include/opcode.h
diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -482,15 +482,8 @@ def pseudo_op(name, op, real_ops):
     # Common type checks
     # These instructions check that one operand is a certain type.
     # Their oparg is the offset from TOS to read.
-    # 'UNARY_CHECK_INT',
-    # 'UNARY_CHECK_FLOAT',
-    # 'UNARY_CHECK_STR',
-
-    # These instructions check that both operands are a certain type.
-    # The benefit is that they save some dispatch overhead versus the
-    # single operand forms.
-    'BINARY_CHECK_INT',
-    'BINARY_CHECK_FLOAT',
+    'CHECK_INT',
+    'CHECK_FLOAT',
     'CHECK_LIST',
 
     # These are guardless instructions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -309,22 +309,17 @@ dummy_func(
             DECREF_INPUTS_AND_REUSE_FLOAT(left, right, dsum, sum);
         }
 
-        inst(BINARY_CHECK_FLOAT, (
-                left, right 
-                -- left_unboxed : {<<= PyFloat_Type, PyRawFloat_Type}, 
-                   right_unboxed: {<<= PyFloat_Type, PyRawFloat_Type})
-            ) {
+        inst(CHECK_FLOAT, (maybe_float, unused[oparg] -- unboxed_float : {<<= PyFloat_Type, PyRawFloat_Type}, unused[oparg])) {
             assert(cframe.use_tracing == 0);
-            char is_successor = PyFloat_CheckExact(left) && (Py_TYPE(left) == Py_TYPE(right));
+            char is_successor = PyFloat_CheckExact(maybe_float);
             bb_test = BB_TEST(is_successor, 0);
 
             if (is_successor) {
-                left_unboxed = *((PyObject **)(&(((PyFloatObject *)left)->ob_fval)));
-                right_unboxed = *((PyObject **)(&(((PyFloatObject *)right)->ob_fval)));
+                unboxed_float = *((PyObject **)(&(((PyFloatObject *)maybe_float)->ob_fval)));
                 DECREF_INPUTS();
-            } else {
-                left_unboxed = left;
-                right_unboxed = right;
+            }
+            else {
+                unboxed_float = maybe_float;
             }
         }
 
@@ -363,9 +358,9 @@ dummy_func(
             U_INST(BINARY_OP_ADD_INT_REST);
         }
 
-        inst(BINARY_CHECK_INT, (left, right -- left : <<= PyLong_Type, right : <<= PyLong_Type)) {
+        inst(CHECK_INT, (maybe_int, unused[oparg] -- maybe_int : <<= PyLong_Type, unused[oparg])) {
             assert(cframe.use_tracing == 0);
-            char is_successor = PyLong_CheckExact(left) && (Py_TYPE(left) == Py_TYPE(right));
+            char is_successor = PyLong_CheckExact(maybe_int);
             bb_test = BB_TEST(is_successor, 0);
         }
 

diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
@@ -75,8 +75,8 @@ _PyOpcode_num_popped(int opcode, int oparg, bool jump) {
             return 2;
         case BINARY_OP_ADD_FLOAT:
             return 2;
-        case BINARY_CHECK_FLOAT:
-            return 2;
+        case CHECK_FLOAT:
+            return oparg + 1;
         case BINARY_OP_ADD_FLOAT_UNBOXED:
             return 2;
         case BINARY_OP_SUBTRACT_FLOAT_UNBOXED:
@@ -89,8 +89,8 @@ _PyOpcode_num_popped(int opcode, int oparg, bool jump) {
             return oparg + 1;
         case BINARY_OP_ADD_INT:
             return 2;
-        case BINARY_CHECK_INT:
-            return 2;
+        case CHECK_INT:
+            return oparg + 1;
         case BINARY_OP_ADD_INT_REST:
             return 2;
         case BINARY_SUBSCR:
@@ -491,8 +491,8 @@ _PyOpcode_num_pushed(int opcode, int oparg, bool jump) {
             return 0;
         case BINARY_OP_ADD_FLOAT:
             return 1;
-        case BINARY_CHECK_FLOAT:
-            return 2;
+        case CHECK_FLOAT:
+            return oparg + 1;
         case BINARY_OP_ADD_FLOAT_UNBOXED:
             return 1;
         case BINARY_OP_SUBTRACT_FLOAT_UNBOXED:
@@ -505,8 +505,8 @@ _PyOpcode_num_pushed(int opcode, int oparg, bool jump) {
             return oparg + 1;
         case BINARY_OP_ADD_INT:
             return 1;
-        case BINARY_CHECK_INT:
-            return 2;
+        case CHECK_INT:
+            return oparg + 1;
         case BINARY_OP_ADD_INT_REST:
             return 1;
         case BINARY_SUBSCR:
@@ -878,14 +878,14 @@ const struct opcode_metadata _PyOpcode_opcode_metadata[256] = {
     [BINARY_OP_ADD_UNICODE] = { true, INSTR_FMT_IXC },
     [BINARY_OP_INPLACE_ADD_UNICODE] = { true, INSTR_FMT_IX },
     [BINARY_OP_ADD_FLOAT] = { true, INSTR_FMT_IXC },
-    [BINARY_CHECK_FLOAT] = { true, INSTR_FMT_IX },
+    [CHECK_FLOAT] = { true, INSTR_FMT_IB },
     [BINARY_OP_ADD_FLOAT_UNBOXED] = { true, INSTR_FMT_IX },
     [BINARY_OP_SUBTRACT_FLOAT_UNBOXED] = { true, INSTR_FMT_IX },
     [BINARY_OP_MULTIPLY_FLOAT_UNBOXED] = { true, INSTR_FMT_IX },
     [UNBOX_FLOAT] = { true, INSTR_FMT_IB },
     [BOX_FLOAT] = { true, INSTR_FMT_IB },
     [BINARY_OP_ADD_INT] = { true, INSTR_FMT_IXC },
-    [BINARY_CHECK_INT] = { true, INSTR_FMT_IX },
+    [CHECK_INT] = { true, INSTR_FMT_IB },
     [BINARY_OP_ADD_INT_REST] = { true, INSTR_FMT_IX },
     [BINARY_SUBSCR] = { true, INSTR_FMT_IXC000 },
     [BINARY_SLICE] = { true, INSTR_FMT_IX },

diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
diff --git a/Python/tier2.c b/Python/tier2.c
@@ -14,7 +14,7 @@
 // Max typed version basic blocks per basic block
 #define MAX_BB_VERSIONS 10
 
-#define OVERALLOCATE_FACTOR 6
+#define OVERALLOCATE_FACTOR 7
 
 
 /* Dummy types used by the types propagator */
@@ -1090,8 +1090,8 @@ write_bb_id(_PyBBBranchCache *cache, int bb_id, bool is_type_guard) {
 */
 static int type_guard_ladder[256] = {
     -1,
-    BINARY_CHECK_FLOAT,
-    BINARY_CHECK_INT,
+    CHECK_FLOAT,
+    CHECK_INT,
     -1,
     CHECK_LIST,
     -1,
@@ -1103,8 +1103,8 @@ static int type_guard_ladder[256] = {
  * KEEP IN SYNC WITH INDEX IN type_guard_ladder
 */
 static int type_guard_to_index[256] = {
-    [BINARY_CHECK_FLOAT] = 1,
-    [BINARY_CHECK_INT] = 2,
+    [CHECK_FLOAT] = 1,
+    [CHECK_INT] = 2,
     [CHECK_LIST] = 4,
 };
 
@@ -1480,8 +1480,8 @@ infer_BINARY_OP(
     *needs_guard = false;
     PyTypeObject *right = typenode_get_type(type_context->type_stack_ptr[-1]);
     PyTypeObject *left = typenode_get_type(type_context->type_stack_ptr[-2]);
-    if (left == &PyLong_Type) {
-        if (right == &PyLong_Type) {
+    if (left == &PyLong_Type || right == &PyLong_Type) {
+        if (left == &PyLong_Type && right == &PyLong_Type) {
             int opcode = oparg == NB_ADD
                 ? BINARY_OP_ADD_INT_REST
                 : oparg == NB_SUBTRACT
@@ -1494,40 +1494,63 @@ infer_BINARY_OP(
             type_propagate(opcode, 0, type_context, NULL);
             return write_curr;
         }
+        // One is unknown
+        int other_oparg = right == NULL ? 0 : 1;
+        if (prev_type_guard != NULL &&
+            prev_type_guard->op.code == CHECK_INT &&
+            prev_type_guard->op.arg != other_oparg &&
+            (right == NULL || left == NULL)) {
+            *needs_guard = true;
+            return emit_type_guard(write_curr, CHECK_INT, other_oparg, bb_id);
+        }
     }
-    if ((left == &PyRawFloat_Type || left == &PyFloat_Type) &&
+    // One of them are floats
+    if ((left == &PyRawFloat_Type || left == &PyFloat_Type) ||
         (right == &PyRawFloat_Type || right == &PyFloat_Type)) {
-        int opcode = oparg == NB_ADD
-            ? BINARY_OP_ADD_FLOAT_UNBOXED
-            : oparg == NB_SUBTRACT
-            ? BINARY_OP_SUBTRACT_FLOAT_UNBOXED
-            : oparg == NB_MULTIPLY
-            ? BINARY_OP_MULTIPLY_FLOAT_UNBOXED
-            : (Py_UNREACHABLE(), 1);
-        if (right == &PyFloat_Type) {
-            write_curr->op.code = UNBOX_FLOAT;
-            write_curr->op.arg = 0;
+        // Both side float, emit the optimised addition instruction
+        if ((left == &PyRawFloat_Type || left == &PyFloat_Type) &&
+            (right == &PyRawFloat_Type || right == &PyFloat_Type)) {
+            int opcode = oparg == NB_ADD
+                ? BINARY_OP_ADD_FLOAT_UNBOXED
+                : oparg == NB_SUBTRACT
+                ? BINARY_OP_SUBTRACT_FLOAT_UNBOXED
+                : oparg == NB_MULTIPLY
+                ? BINARY_OP_MULTIPLY_FLOAT_UNBOXED
+                : (Py_UNREACHABLE(), 1);
+            if (right == &PyFloat_Type) {
+                write_curr->op.code = UNBOX_FLOAT;
+                write_curr->op.arg = 0;
+                write_curr++;
+                type_propagate(UNBOX_FLOAT, 0, type_context, NULL);
+            }
+            if (left == &PyFloat_Type) {
+                write_curr->op.code = UNBOX_FLOAT;
+                write_curr->op.arg = 1;
+                write_curr++;
+                type_propagate(UNBOX_FLOAT, 1, type_context, NULL);
+            }
+            write_curr->op.code = opcode;
             write_curr++;
-            type_propagate(UNBOX_FLOAT, 0, type_context, NULL);
+            type_propagate(opcode, 0, type_context, NULL);
+            return write_curr;
         }
-        if (left == &PyFloat_Type) {
-            write_curr->op.code = UNBOX_FLOAT;
-            write_curr->op.arg = 1;
-            write_curr++;
-            type_propagate(UNBOX_FLOAT, 1, type_context, NULL);
+        // One is unknown
+        int other_oparg = right == NULL ? 0 : 1;
+        if (prev_type_guard != NULL &&
+            prev_type_guard->op.code == CHECK_FLOAT &&
+            prev_type_guard->op.arg != other_oparg &&
+            (right == NULL || left == NULL)) {
+            *needs_guard = true;
+            return emit_type_guard(write_curr, CHECK_FLOAT, right == NULL ? 0 : 1, bb_id);
         }
-        write_curr->op.code = opcode;
-        write_curr++;
-        type_propagate(opcode, 0, type_context, NULL);
-        return write_curr;
     }
-    // Unknown, time to emit the chain of guards.
+    // Both unknown, time to emit the chain of guards.
     // No type guard before this, or it's not the first in the new BB.
     // First in new BB usually indicates it's already part of a pre-existing ladder.
     if (prev_type_guard == NULL || !is_first_instr) {
         write_curr = rebox_stack(write_curr, type_context, 2);
         *needs_guard = true;
-        return emit_type_guard(write_curr, BINARY_CHECK_FLOAT, 0, bb_id);
+        return emit_type_guard(write_curr, CHECK_FLOAT, 0, bb_id);
     }
     else {
         int next_guard = type_guard_ladder[
@@ -1650,8 +1673,8 @@ _PyTier2_Code_DetectAndEmitBB(
 {
     assert(
         prev_type_guard == NULL ||
-        prev_type_guard->op.code == BINARY_CHECK_INT ||
-        prev_type_guard->op.code == BINARY_CHECK_FLOAT ||
+        prev_type_guard->op.code == CHECK_INT ||
+        prev_type_guard->op.code == CHECK_FLOAT ||
         prev_type_guard->op.code == CHECK_LIST
     );
 #define END() goto end;

diff --git a/Python/tier2_typepropagator.c.h b/Python/tier2_typepropagator.c.h