gh-115480: Type and constant propagation for int BINARY_OPs (GH-115478)

python · Feb 15, 2024 · 4ebf8fb · 4ebf8fb
1 parent ed23839
commit 4ebf8fb
Show file tree

Hide file tree

Showing 3 changed files with 126 additions and 16 deletions.
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
@@ -341,6 +341,18 @@ sym_new_const(_Py_UOpsAbstractInterpContext *ctx, PyObject *const_val)
     return temp;
 }
 
+static inline bool
+is_const(_Py_UOpsSymType *sym)
+{
+    return sym->const_val != NULL;
+}
+
+static inline PyObject *
+get_const(_Py_UOpsSymType *sym)
+{
+    return sym->const_val;
+}
+
 static _Py_UOpsSymType*
 sym_new_null(_Py_UOpsAbstractInterpContext *ctx)
 {

diff --git a/Python/tier2_redundancy_eliminator_bytecodes.c b/Python/tier2_redundancy_eliminator_bytecodes.c
@@ -81,12 +81,62 @@ dummy_func(void) {
 
 
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {
-        // TODO constant propagation
-        (void)left;
-        (void)right;
-        res = sym_new_known_type(ctx, &PyLong_Type);
-        if (res == NULL) {
-            goto out_of_space;
+        if (is_const(left) && is_const(right)) {
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Add((PyLongObject *)get_const(left),
+                                         (PyLongObject *)get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
+        }
+    }
+
+    op(_BINARY_OP_SUBTRACT_INT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Subtract((PyLongObject *)get_const(left),
+                                              (PyLongObject *)get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
+        }
+    }
+
+    op(_BINARY_OP_MULTIPLY_INT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Multiply((PyLongObject *)get_const(left),
+                                              (PyLongObject *)get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
         }
     }
 

diff --git a/Python/tier2_redundancy_eliminator_cases.c.h b/Python/tier2_redundancy_eliminator_cases.c.h
@@ -180,9 +180,28 @@
         }
 
         case _BINARY_OP_MULTIPLY_INT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Multiply((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
@@ -194,22 +213,51 @@
             _Py_UOpsSymType *res;
             right = stack_pointer[-1];
             left = stack_pointer[-2];
-            // TODO constant propagation
-            (void)left;
-            (void)right;
-            res = sym_new_known_type(ctx, &PyLong_Type);
-            if (res == NULL) {
-                goto out_of_space;
+            if (is_const(left) && is_const(right)) {
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Add((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
             }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
         }
 
         case _BINARY_OP_SUBTRACT_INT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Subtract((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;