Add new flow graph optimization to fold to constants

Summary:
3.16 added a new flow graph optimization. This catches our compiler up.

  PythonLib/cinderx/compiler/flow_graph_optimizer.py
  - Generalized fold_constant_intrinsic_list_to_tuple → fold_constant_seq_into_load_const, mirroring upstream: the target may be the LIST_TO_TUPLE intrinsic or a trailing LIST_APPEND/SET_ADD, with sets folding to a frozenset. Behavior for 3.14/3.15 is unchanged (their
  only caller still passes the intrinsic).
  - Added FlowGraphOptimizer316 with the new LIST_APPEND/SET_ADD handler and the reordered CALL_INTRINSIC_1 handler.

  PythonLib/cinderx/compiler/pyassem.py — PyFlowGraph316.flow_graph_optimizer = FlowGraphOptimizer316.

  PythonLib/test_cinderx/test_compiler/test_optimizer.py — three version-gated regression tests (big const list iteration, big const set membership, and a negative case ensuring a list that escapes is still built as a list). Two of the three fail without the fix; 3.12
  is skipped since it folds big constant sequences in codegen instead.

Reviewed By: yoney

Differential Revision: D116650708

fbshipit-source-id: 6b2626b6aaa5ba002d18ee8c986068d5731f3cb4

Author: DinoV

cinderx/PythonLib/cinderx/compiler/flow_graph_optimizer.py

@@ -1048,43 +1048,65 @@ def optimize_unary_not(
         assert isinstance(self, FlowGraphOptimizer314)
         self.optimize_one_unary(instr_index, instr, block, operator.not_)
 
-    def fold_constant_intrinsic_list_to_tuple(
-        self, block: Block, instr_index: int
-    ) -> None:
-        consts_found = 0
-        expect_append = True
+    def fold_constant_seq_into_load_const(self, block: Block, instr_index: int) -> None:
+        """Replace:
+            BUILD_LIST/BUILD_SET 0
+            LOAD_CONST c1
+            LIST_APPEND/SET_ADD 1
+            ...
+            LOAD_CONST cN
+            LIST_APPEND/SET_ADD 1
+            [CALL_INTRINSIC_1 INTRINSIC_LIST_TO_TUPLE]   <-- optional
+        with:
+            LOAD_CONST (c1, c2, ... cN)
+
+        The instruction at `instr_index` is either the LIST_TO_TUPLE intrinsic
+        (so only the BUILD_LIST/LIST_APPEND form is considered), or the trailing
+        LIST_APPEND or SET_ADD itself, in which case the matching
+        BUILD_LIST/BUILD_SET start is selected from its opcode and sets are
+        folded into a frozenset.
+        """
+        target = block.insts[instr_index]
+        expected_append = target.opname == "CALL_INTRINSIC_1"
+        append_op = "LIST_APPEND" if expected_append else target.opname
+        build_op = "BUILD_LIST" if append_op == "LIST_APPEND" else "BUILD_SET"
+        # Walking backwards, appends and the constants they push alternate. The
+        # intrinsic is preceded by an append, a trailing append by a constant.
+        expect_append = expected_append
         for i in range(instr_index - 1, -1, -1):
             instr = block.insts[i]
             opcode = instr.opname
             oparg = instr.oparg
             if opcode == "NOP":
                 continue
 
-            if opcode == "BUILD_LIST" and oparg == 0:
+            if opcode == build_op and oparg == 0:
                 if not expect_append:
                     # Not a start sequence
                     return
 
                 # Sequence start, we are done.
                 consts = []
-                if opcode == "BUILD_LIST" and oparg == 0:
-                    for newpos in range(instr_index - 1, i - 1, -1):
-                        instr = block.insts[newpos]
-                        if instr.opname in LOAD_CONST_INSTRS:
-                            const = instr.oparg
-                            consts.append(const)
-                        instr.set_to_nop_no_loc()
+                start = instr_index - 1 if expected_append else instr_index
+                for newpos in range(start, i - 1, -1):
+                    instr = block.insts[newpos]
+                    if instr.opname in LOAD_CONST_INSTRS:
+                        const = instr.oparg
+                        consts.append(const)
+                    instr.set_to_nop_no_loc()
 
                 consts.reverse()
-                self.make_load_const(block.insts[instr_index], tuple(consts))
+                newconst: object = (
+                    frozenset(consts) if build_op == "BUILD_SET" else tuple(consts)
+                )
+                self.make_load_const(target, newconst)
                 return
 
             if expect_append:
-                if opcode != "LIST_APPEND" or oparg != 1:
+                if opcode != append_op or oparg != 1:
                     return
             elif opcode not in LOAD_CONST_INSTRS:
                 return
-            consts_found += 1
             expect_append = not expect_append
 
     def optimize_call_intrinsic_1(
@@ -1101,7 +1123,7 @@ def optimize_call_intrinsic_1(
             if next_instr is not None and next_instr.opname == "GET_ITER":
                 instr.set_to_nop()
             else:
-                self.fold_constant_intrinsic_list_to_tuple(block, instr_index)
+                self.fold_constant_seq_into_load_const(block, instr_index)
         if intrins == "INTRINSIC_UNARY_POSITIVE":
             # pyrefly: ignore [bad-argument-type]
             self.optimize_one_unary(instr_index, instr, block, operator.pos)
@@ -1154,6 +1176,62 @@ def optimize_basic_block(self, block: Block) -> None:
             i += 1
 
 
+class FlowGraphOptimizer316(FlowGraphOptimizer314):
+    """Python 3.16-specific optimizations."""
+
+    def optimize_call_intrinsic_1(
+        self: FlowGraphOptimizer,
+        instr_index: int,
+        instr: Instruction,
+        next_instr: Instruction | None,
+        target: Instruction | None,
+        block: Block,
+    ) -> int | None:
+        assert isinstance(self, FlowGraphOptimizer316)
+        intrins = INTRINSIC_1[instr.ioparg]
+        if intrins == "INTRINSIC_LIST_TO_TUPLE":
+            # Unlike 3.14/3.15, folding is attempted even when iterating, so a
+            # big constant tuple becomes a single LOAD_CONST instead of staying
+            # a list build. The intrinsic is only dropped if folding didn't
+            # already rewrite it.
+            self.fold_constant_seq_into_load_const(block, instr_index)
+            if (
+                instr.opname == "CALL_INTRINSIC_1"
+                and next_instr is not None
+                and next_instr.opname == "GET_ITER"
+            ):
+                instr.set_to_nop()
+        if intrins == "INTRINSIC_UNARY_POSITIVE":
+            # pyrefly: ignore [bad-argument-type]
+            self.optimize_one_unary(instr_index, instr, block, operator.pos)
+
+    def optimize_list_append_set_add(
+        self: FlowGraphOptimizer,
+        instr_index: int,
+        instr: Instruction,
+        next_instr: Instruction | None,
+        target: Instruction | None,
+        block: Block,
+    ) -> int | None:
+        assert isinstance(self, FlowGraphOptimizer316)
+        # Sequences too big for optimize_lists_and_sets are built with repeated
+        # appends. When they're only iterated over or tested for membership a
+        # constant tuple/frozenset is a suitable replacement.
+        if (
+            instr.oparg == 1
+            and next_instr is not None
+            and next_instr.opname in ("GET_ITER", "CONTAINS_OP")
+        ):
+            self.fold_constant_seq_into_load_const(block, instr_index)
+
+    handlers: dict[str, Handler] = {
+        **FlowGraphOptimizer314.handlers,
+        "CALL_INTRINSIC_1": optimize_call_intrinsic_1,
+        "LIST_APPEND": optimize_list_append_set_add,
+        "SET_ADD": optimize_list_append_set_add,
+    }
+
+
 class FlowGraphConstOptimizer314(BaseFlowGraphOptimizer314):
     def opt_load_const(
         self: FlowGraphOptimizer,

cinderx/PythonLib/cinderx/compiler/pyassem.py

@@ -55,6 +55,7 @@
     FlowGraphOptimizer,
     FlowGraphOptimizer312,
     FlowGraphOptimizer314,
+    FlowGraphOptimizer316,
 )
 from .opcode_cinder import opcode as cinder_opcode
 from .opcodebase import Opcode
@@ -3305,6 +3306,7 @@ def propagate_line_numbers(self) -> None:
 
 
 class PyFlowGraph316(PyFlowGraph315):
+    flow_graph_optimizer = FlowGraphOptimizer316
     flow_graph_const_optimizer = FlowGraphConstOptimizer316
 
     def convert_load_const_to_load_common_constant(self) -> None:

cinderx/PythonLib/test_cinderx/test_compiler/test_optimizer.py

@@ -30,6 +30,11 @@
 
 from .common import CompilerTest
 
+# A constant sequence long enough to exceed STACK_USE_GUIDELINE, so it is built
+# with repeated appends rather than a single BUILD_LIST/BUILD_SET.
+BIG_CONST_TUPLE: tuple[int, ...] = tuple(range(100, 140))
+BIG_CONST_ELTS: str = ", ".join(str(i) for i in BIG_CONST_TUPLE)
+
 
 class AstOptimizerTests(CompilerTest):
     class _Comparer:
@@ -483,6 +488,38 @@ def test_frozenset_call_optimization(self) -> None:
         else:
             self.assertNotInGraph(graph, "CALL_INTRINSIC_1")
 
+    @passIf(sys.version_info < (3, 14), "3.12 folds big constant sequences in codegen")
+    def test_big_const_list_iteration(self) -> None:
+        # Sequences over STACK_USE_GUIDELINE elements are built with repeated
+        # appends. 3.16 folds such a chain into a constant tuple when the
+        # result is only iterated over; earlier versions keep the build.
+        graph = self.to_graph(f"for x in [{BIG_CONST_ELTS}]: pass")
+        if sys.version_info >= (3, 16):
+            self.assertInGraph(graph, "LOAD_CONST", BIG_CONST_TUPLE)
+            self.assertNotInGraph(graph, "BUILD_LIST")
+        else:
+            self.assertInGraph(graph, "BUILD_LIST", 0)
+            self.assertNotInGraph(graph, "LOAD_CONST", BIG_CONST_TUPLE)
+
+    @passIf(sys.version_info < (3, 14), "3.12 folds big constant sequences in codegen")
+    def test_big_const_set_membership(self) -> None:
+        # Same as above for sets, which fold into a frozenset.
+        graph = self.to_graph(f"y = x in {{{BIG_CONST_ELTS}}}")
+        if sys.version_info >= (3, 16):
+            self.assertInGraph(graph, "LOAD_CONST", frozenset(BIG_CONST_TUPLE))
+            self.assertNotInGraph(graph, "BUILD_SET")
+        else:
+            self.assertInGraph(graph, "BUILD_SET", 0)
+            self.assertNotInGraph(graph, "LOAD_CONST", frozenset(BIG_CONST_TUPLE))
+
+    @passIf(sys.version_info < (3, 14), "3.12 folds big constant sequences in codegen")
+    def test_big_const_list_not_folded_when_used_as_a_list(self) -> None:
+        # The fold is only valid when the sequence is consumed by GET_ITER or
+        # CONTAINS_OP; a list that escapes must still be built as a list.
+        graph = self.to_graph(f"x = [{BIG_CONST_ELTS}]")
+        self.assertInGraph(graph, "BUILD_LIST", 0)
+        self.assertNotInGraph(graph, "LOAD_CONST", BIG_CONST_TUPLE)
+
 
 class _FakeInstr:
     __slots__ = ("opname", "oparg", "ioparg")