GH-115480: Reduce guard strength for binary ops when type of one operand is known already (GH-118050)

6 files changed, 171 insertions, 19 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index b7511b9..4541eb6 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -426,6 +426,14 @@ dummy_func(
             EXIT_IF(!PyLong_CheckExact(right));
         }
 
+        op(_GUARD_NOS_INT, (left, unused -- left, unused)) {
+            EXIT_IF(!PyLong_CheckExact(left));
+        }
+
+        op(_GUARD_TOS_INT, (value -- value)) {
+            EXIT_IF(!PyLong_CheckExact(value));
+        }
+
         pure op(_BINARY_OP_MULTIPLY_INT, (left, right -- res)) {
             STAT_INC(BINARY_OP, hit);
             res = _PyLong_Multiply((PyLongObject *)left, (PyLongObject *)right);
@@ -462,6 +470,14 @@ dummy_func(
             EXIT_IF(!PyFloat_CheckExact(right));
         }
 
+        op(_GUARD_NOS_FLOAT, (left, unused -- left, unused)) {
+            EXIT_IF(!PyFloat_CheckExact(left));
+        }
+
+        op(_GUARD_TOS_FLOAT, (value -- value)) {
+            EXIT_IF(!PyFloat_CheckExact(value));
+        }
+
         pure op(_BINARY_OP_MULTIPLY_FLOAT, (left, right -- res)) {
             STAT_INC(BINARY_OP, hit);
             double dres =
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 841ce8c..43b0221 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -447,6 +447,26 @@
             break;
         }
 
+        case _GUARD_NOS_INT: {
+            PyObject *left;
+            left = stack_pointer[-2];
+            if (!PyLong_CheckExact(left)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            break;
+        }
+
+        case _GUARD_TOS_INT: {
+            PyObject *value;
+            value = stack_pointer[-1];
+            if (!PyLong_CheckExact(value)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            break;
+        }
+
         case _BINARY_OP_MULTIPLY_INT: {
             PyObject *right;
             PyObject *left;
@@ -511,6 +531,26 @@
             break;
         }
 
+        case _GUARD_NOS_FLOAT: {
+            PyObject *left;
+            left = stack_pointer[-2];
+            if (!PyFloat_CheckExact(left)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            break;
+        }
+
+        case _GUARD_TOS_FLOAT: {
+            PyObject *value;
+            value = stack_pointer[-1];
+            if (!PyFloat_CheckExact(value)) {
+                UOP_STAT_INC(uopcode, miss);
+                JUMP_TO_JUMP_TARGET();
+            }
+            break;
+        }
+
         case _BINARY_OP_MULTIPLY_FLOAT: {
             PyObject *right;
             PyObject *left;
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 155f702..76de6e5 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -320,6 +320,7 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
 #define sym_new_const _Py_uop_sym_new_const
 #define sym_new_null _Py_uop_sym_new_null
 #define sym_has_type _Py_uop_sym_has_type
+#define sym_get_type _Py_uop_sym_get_type
 #define sym_matches_type _Py_uop_sym_matches_type
 #define sym_set_null _Py_uop_sym_set_null
 #define sym_set_non_null _Py_uop_sym_set_non_null
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index f119b8e..481fb83 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -21,6 +21,7 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_new_const _Py_uop_sym_new_const
 #define sym_new_null _Py_uop_sym_new_null
 #define sym_matches_type _Py_uop_sym_matches_type
+#define sym_get_type _Py_uop_sym_get_type
 #define sym_has_type _Py_uop_sym_has_type
 #define sym_set_null _Py_uop_sym_set_null
 #define sym_set_non_null _Py_uop_sym_set_non_null
@@ -99,9 +100,18 @@ dummy_func(void) {
     }
 
     op(_GUARD_BOTH_INT, (left, right -- left, right)) {
-        if (sym_matches_type(left, &PyLong_Type) &&
-            sym_matches_type(right, &PyLong_Type)) {
-            REPLACE_OP(this_instr, _NOP, 0, 0);
+        if (sym_matches_type(left, &PyLong_Type)) {
+            if (sym_matches_type(right, &PyLong_Type)) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+            else {
+                REPLACE_OP(this_instr, _GUARD_TOS_INT, 0, 0);
+            }
+        }
+        else {
+            if (sym_matches_type(right, &PyLong_Type)) {
+                REPLACE_OP(this_instr, _GUARD_NOS_INT, 0, 0);
+            }
         }
         if (!sym_set_type(left, &PyLong_Type)) {
             goto hit_bottom;
@@ -112,9 +122,18 @@ dummy_func(void) {
     }
 
     op(_GUARD_BOTH_FLOAT, (left, right -- left, right)) {
-        if (sym_matches_type(left, &PyFloat_Type) &&
-            sym_matches_type(right, &PyFloat_Type)) {
-            REPLACE_OP(this_instr, _NOP, 0 ,0);
+        if (sym_matches_type(left, &PyFloat_Type)) {
+            if (sym_matches_type(right, &PyFloat_Type)) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+            else {
+                REPLACE_OP(this_instr, _GUARD_TOS_FLOAT, 0, 0);
+            }
+        }
+        else {
+            if (sym_matches_type(right, &PyFloat_Type)) {
+                REPLACE_OP(this_instr, _GUARD_NOS_FLOAT, 0, 0);
+            }
         }
         if (!sym_set_type(left, &PyFloat_Type)) {
             goto hit_bottom;
@@ -137,6 +156,25 @@ dummy_func(void) {
         }
     }
 
+    op(_BINARY_OP, (left, right -- res)) {
+        PyTypeObject *ltype = sym_get_type(left);
+        PyTypeObject *rtype = sym_get_type(right);
+        if (ltype != NULL && (ltype == &PyLong_Type || ltype == &PyFloat_Type) &&
+            rtype != NULL && (rtype == &PyLong_Type || rtype == &PyFloat_Type))
+        {
+            if (oparg != NB_TRUE_DIVIDE && oparg != NB_INPLACE_TRUE_DIVIDE &&
+                ltype == &PyLong_Type && rtype == &PyLong_Type) {
+                /* If both inputs are ints and the op is not division the result is an int */
+                OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyLong_Type));
+            }
+            else {
+                /* For any other op combining ints/floats the result is a float */
+                OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyFloat_Type));
+            }
+        }
+        OUT_OF_SPACE_IF_NULL(res = sym_new_unknown(ctx));
+    }
+
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {
         if (sym_is_const(left) && sym_is_const(right) &&
             sym_matches_type(left, &PyLong_Type) && sym_matches_type(right, &PyLong_Type))
@@ -424,7 +462,6 @@ dummy_func(void) {
         OUT_OF_SPACE_IF_NULL(null = sym_new_null(ctx));
     }
 
-
     op(_COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
         assert(oparg > 0);
         top = bottom;
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 50f335e..0a7d96d 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -225,9 +225,18 @@
             _Py_UopsSymbol *left;
             right = stack_pointer[-1];
             left = stack_pointer[-2];
-            if (sym_matches_type(left, &PyLong_Type) &&
-                sym_matches_type(right, &PyLong_Type)) {
-                REPLACE_OP(this_instr, _NOP, 0, 0);
+            if (sym_matches_type(left, &PyLong_Type)) {
+                if (sym_matches_type(right, &PyLong_Type)) {
+                    REPLACE_OP(this_instr, _NOP, 0, 0);
+                }
+                else {
+                    REPLACE_OP(this_instr, _GUARD_TOS_INT, 0, 0);
+                }
+            }
+            else {
+                if (sym_matches_type(right, &PyLong_Type)) {
+                    REPLACE_OP(this_instr, _GUARD_NOS_INT, 0, 0);
+                }
             }
             if (!sym_set_type(left, &PyLong_Type)) {
                 goto hit_bottom;
@@ -238,6 +247,14 @@
             break;
         }
 
+        case _GUARD_NOS_INT: {
+            break;
+        }
+
+        case _GUARD_TOS_INT: {
+            break;
+        }
+
         case _BINARY_OP_MULTIPLY_INT: {
             _Py_UopsSymbol *right;
             _Py_UopsSymbol *left;
@@ -333,9 +350,18 @@
             _Py_UopsSymbol *left;
             right = stack_pointer[-1];
             left = stack_pointer[-2];
-            if (sym_matches_type(left, &PyFloat_Type) &&
-                sym_matches_type(right, &PyFloat_Type)) {
-                REPLACE_OP(this_instr, _NOP, 0 ,0);
+            if (sym_matches_type(left, &PyFloat_Type)) {
+                if (sym_matches_type(right, &PyFloat_Type)) {
+                    REPLACE_OP(this_instr, _NOP, 0, 0);
+                }
+                else {
+                    REPLACE_OP(this_instr, _GUARD_TOS_FLOAT, 0, 0);
+                }
+            }
+            else {
+                if (sym_matches_type(right, &PyFloat_Type)) {
+                    REPLACE_OP(this_instr, _GUARD_NOS_FLOAT, 0, 0);
+                }
             }
             if (!sym_set_type(left, &PyFloat_Type)) {
                 goto hit_bottom;
@@ -346,6 +372,14 @@
             break;
         }
 
+        case _GUARD_NOS_FLOAT: {
+            break;
+        }
+
+        case _GUARD_TOS_FLOAT: {
+            break;
+        }
+
         case _BINARY_OP_MULTIPLY_FLOAT: {
             _Py_UopsSymbol *right;
             _Py_UopsSymbol *left;
@@ -1852,9 +1886,27 @@
         }
 
         case _BINARY_OP: {
+            _Py_UopsSymbol *right;
+            _Py_UopsSymbol *left;
             _Py_UopsSymbol *res;
-            res = sym_new_not_null(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            PyTypeObject *ltype = sym_get_type(left);
+            PyTypeObject *rtype = sym_get_type(right);
+            if (ltype != NULL && (ltype == &PyLong_Type || ltype == &PyFloat_Type) &&
+                rtype != NULL && (rtype == &PyLong_Type || rtype == &PyFloat_Type))
+            {
+                if (oparg != NB_TRUE_DIVIDE && oparg != NB_INPLACE_TRUE_DIVIDE &&
+                    ltype == &PyLong_Type && rtype == &PyLong_Type) {
+                    /* If both inputs are ints and the op is not division the result is an int */
+                    OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyLong_Type));
+                }
+                else {
+                    /* For any other op combining ints/floats the result is a float */
+                    OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyFloat_Type));
+                }
+            }
+            OUT_OF_SPACE_IF_NULL(res = sym_new_unknown(ctx));
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
diff --git a/Python/optimizer_symbols.c b/Python/optimizer_symbols.c
index 86b0d4d..204599b 100644
--- a/Python/optimizer_symbols.c
+++ b/Python/optimizer_symbols.c
@@ -231,6 +231,15 @@ _Py_uop_sym_new_null(_Py_UOpsContext *ctx)
     return null_sym;
 }
 
+PyTypeObject *
+_Py_uop_sym_get_type(_Py_UopsSymbol *sym)
+{
+    if (_Py_uop_sym_is_bottom(sym)) {
+        return NULL;
+    }
+    return sym->typ;
+}
+
 bool
 _Py_uop_sym_has_type(_Py_UopsSymbol *sym)
 {
@@ -244,10 +253,7 @@ bool
 _Py_uop_sym_matches_type(_Py_UopsSymbol *sym, PyTypeObject *typ)
 {
     assert(typ != NULL && PyType_Check(typ));
-    if (_Py_uop_sym_is_bottom(sym)) {
-        return false;
-    }
-    return sym->typ == typ;
+    return _Py_uop_sym_get_type(sym) == typ;
 }
 
 int