gh-126835: Move optimization of constant sequence creation from codegen to CFG (#129426)

Eclips4 · iritkatriel · WolframAlph · web-flow · commit 7d0521d5fc32 · 2025-02-01T11:39:44.000Z
Codegen phase has an optimization that transforms ``` LOAD_CONST x LOAD_CONST y LOAD_CONXT z BUILD_LIST/BUILD_SET (3) ``` -> ``` BUILD_LIST/BUILD_SET (0) LOAD_CONST (x, y, z) LIST_EXTEND/SET_UPDATE 1 ``` This optimization has now been moved to CFG phase to make #128802 work. Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com> Co-authored-by: Yan Yanchii <yyanchiy@gmail.com>
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -892,7 +892,7 @@ def loop_test():
 %3d           RESUME_CHECK             0
 
 %3d           BUILD_LIST               0
-              LOAD_CONST_MORTAL        0 ((1, 2, 3))
+              LOAD_CONST_MORTAL        1 ((1, 2, 3))
               LIST_EXTEND              1
               LOAD_SMALL_INT           3
               BINARY_OP                5 (*)
@@ -908,7 +908,7 @@ def loop_test():
 
 %3d   L2:     END_FOR
               POP_ITER
-              LOAD_CONST_IMMORTAL      1 (None)
+              LOAD_CONST_IMMORTAL      0 (None)
               RETURN_VALUE
 """ % (loop_test.__code__.co_firstlineno,
        loop_test.__code__.co_firstlineno + 1,
diff --git a/Python/codegen.c b/Python/codegen.c
@@ -201,9 +201,6 @@ static int codegen_subscript(compiler *, expr_ty);
 static int codegen_slice_two_parts(compiler *, expr_ty);
 static int codegen_slice(compiler *, expr_ty);
 
-static bool are_all_items_const(asdl_expr_seq *, Py_ssize_t, Py_ssize_t);
-
-
 static int codegen_with(compiler *, stmt_ty, int);
 static int codegen_async_with(compiler *, stmt_ty, int);
 static int codegen_async_for(compiler *, stmt_ty);
@@ -3210,34 +3207,6 @@ starunpack_helper_impl(compiler *c, location loc,
                        int build, int add, int extend, int tuple)
 {
     Py_ssize_t n = asdl_seq_LEN(elts);
-    if (!injected_arg && n > 2 && are_all_items_const(elts, 0, n)) {
-        PyObject *folded = PyTuple_New(n);
-        if (folded == NULL) {
-            return ERROR;
-        }
-        for (Py_ssize_t i = 0; i < n; i++) {
-            PyObject *val = ((expr_ty)asdl_seq_GET(elts, i))->v.Constant.value;
-            PyTuple_SET_ITEM(folded, i, Py_NewRef(val));
-        }
-        if (tuple && !pushed) {
-            ADDOP_LOAD_CONST_NEW(c, loc, folded);
-        } else {
-            if (add == SET_ADD) {
-                Py_SETREF(folded, PyFrozenSet_New(folded));
-                if (folded == NULL) {
-                    return ERROR;
-                }
-            }
-            ADDOP_I(c, loc, build, pushed);
-            ADDOP_LOAD_CONST_NEW(c, loc, folded);
-            ADDOP_I(c, loc, extend, 1);
-            if (tuple) {
-                ADDOP_I(c, loc, CALL_INTRINSIC_1, INTRINSIC_LIST_TO_TUPLE);
-            }
-        }
-        return SUCCESS;
-    }
-
     int big = n + pushed + (injected_arg ? 1 : 0) > STACK_USE_GUIDELINE;
     int seen_star = 0;
     for (Py_ssize_t i = 0; i < n; i++) {
@@ -3389,18 +3358,6 @@ codegen_set(compiler *c, expr_ty e)
                              BUILD_SET, SET_ADD, SET_UPDATE, 0);
 }
 
-static bool
-are_all_items_const(asdl_expr_seq *seq, Py_ssize_t begin, Py_ssize_t end)
-{
-    for (Py_ssize_t i = begin; i < end; i++) {
-        expr_ty key = (expr_ty)asdl_seq_GET(seq, i);
-        if (key == NULL || key->kind != Constant_kind) {
-            return false;
-        }
-    }
-    return true;
-}
-
 static int
 codegen_subdict(compiler *c, expr_ty e, Py_ssize_t begin, Py_ssize_t end)
 {
diff --git a/Python/flowgraph.c b/Python/flowgraph.c
@@ -1336,6 +1336,17 @@ add_const(PyObject *newconst, PyObject *consts, PyObject *const_cache)
     return (int)index;
 }
 
+static bool
+is_constant_sequence(cfg_instr *inst, int n)
+{
+    for (int i = 0; i < n; i++) {
+        if(!loads_const(inst[i].i_opcode)) {
+            return false;
+        }
+    }
+    return true;
+}
+
 /* Replace LOAD_CONST c1, LOAD_CONST c2 ... LOAD_CONST cn, BUILD_TUPLE n
    with    LOAD_CONST (c1, c2, ... cn).
    The consts table must still be in list form so that the
@@ -1353,10 +1364,8 @@ fold_tuple_on_constants(PyObject *const_cache,
     assert(inst[n].i_opcode == BUILD_TUPLE);
     assert(inst[n].i_oparg == n);
 
-    for (int i = 0; i < n; i++) {
-        if (!loads_const(inst[i].i_opcode)) {
-            return SUCCESS;
-        }
+    if (!is_constant_sequence(inst, n)) {
+        return SUCCESS;
     }
 
     /* Buildup new tuple of constants */
@@ -1384,6 +1393,56 @@ fold_tuple_on_constants(PyObject *const_cache,
     return SUCCESS;
 }
 
+#define MIN_CONST_SEQUENCE_SIZE 3
+/* Replace LOAD_CONST c1, LOAD_CONST c2 ... LOAD_CONST cN, BUILD_LIST N
+   with BUILD_LIST 0, LOAD_CONST (c1, c2, ... cN), LIST_EXTEND 1,
+   or BUILD_SET & SET_UPDATE respectively.
+*/
+static int
+optimize_if_const_list_or_set(PyObject *const_cache, cfg_instr* inst, int n, PyObject *consts)
+{
+    assert(PyDict_CheckExact(const_cache));
+    assert(PyList_CheckExact(consts));
+    assert(inst[n].i_oparg == n);
+
+    int build = inst[n].i_opcode;
+    assert(build == BUILD_LIST || build == BUILD_SET);
+    int extend = build == BUILD_LIST ? LIST_EXTEND : SET_UPDATE;
+
+    if (n < MIN_CONST_SEQUENCE_SIZE || !is_constant_sequence(inst, n)) {
+        return SUCCESS;
+    }
+    PyObject *newconst = PyTuple_New(n);
+    if (newconst == NULL) {
+        return ERROR;
+    }
+    for (int i = 0; i < n; i++) {
+        int op = inst[i].i_opcode;
+        int arg = inst[i].i_oparg;
+        PyObject *constant = get_const_value(op, arg, consts);
+        if (constant == NULL) {
+            return ERROR;
+        }
+        PyTuple_SET_ITEM(newconst, i, constant);
+    }
+    if (build == BUILD_SET) {
+        PyObject *frozenset = PyFrozenSet_New(newconst);
+        if (frozenset == NULL) {
+            return ERROR;
+        }
+        Py_SETREF(newconst, frozenset);
+    }
+    int index = add_const(newconst, consts, const_cache);
+    RETURN_IF_ERROR(index);
+    INSTR_SET_OP1(&inst[0], build, 0);
+    for (int i = 1; i < n - 1; i++) {
+        INSTR_SET_OP0(&inst[i], NOP);
+    }
+    INSTR_SET_OP1(&inst[n-1], LOAD_CONST, index);
+    INSTR_SET_OP1(&inst[n], extend, 1);
+    return SUCCESS;
+}
+
 #define VISITED (-1)
 
 // Replace an arbitrary run of SWAPs and NOPs with an optimal one that has the
@@ -1751,6 +1810,14 @@ optimize_basic_block(PyObject *const_cache, basicblock *bb, PyObject *consts)
                     }
                 }
                 break;
+            case BUILD_LIST:
+            case BUILD_SET:
+                if (i >= oparg) {
+                    if (optimize_if_const_list_or_set(const_cache, inst-oparg, oparg, consts) < 0) {
+                        goto error;
+                    }
+                }
+                break;
             case POP_JUMP_IF_NOT_NONE:
             case POP_JUMP_IF_NONE:
                 switch (target->i_opcode) {
