[mypyc] feat: remove redundant box/unbox operations

Author: BobTheBuidler

mypyc/analysis/dataflow.py

@@ -152,8 +152,16 @@ def cleanup_cfg(blocks: list[BasicBlock]) -> None:
         cfg = get_cfg(blocks)
         orig_blocks = blocks.copy()
         blocks.clear()
+        
         for i, block in enumerate(orig_blocks):
-            if i == 0 or cfg.pred[block]:
+            # Keep the block if it's the entry block or has any predecessors
+            # OR if it is a branch target (i.e., appears in any terminator.targets())
+            is_branch_target = any(
+                block is tgt
+                for b in orig_blocks
+                for tgt in b.terminator.targets()
+            )
+            if i == 0 or cfg.pred[block] or is_branch_target:
                 blocks.append(block)
             else:
                 changed = True

mypyc/codegen/emitmodule.py

@@ -67,6 +67,7 @@
 from mypyc.transform.flag_elimination import do_flag_elimination
 from mypyc.transform.log_trace import insert_event_trace_logging
 from mypyc.transform.lower import lower_ir
+from mypyc.transform.redundant_box_elimination import do_box_unbox_elimination
 from mypyc.transform.refcount import insert_ref_count_opcodes
 from mypyc.transform.spill import insert_spills
 from mypyc.transform.uninit import insert_uninit_checks
@@ -227,7 +228,6 @@ def compile_scc_to_ir(
 
     Returns the IR of the modules.
     """
-
     if compiler_options.verbose:
         print("Compiling {}".format(", ".join(x.name for x in scc)))
 
@@ -260,6 +260,7 @@ def compile_scc_to_ir(
             # Switch to lower abstraction level IR.
             lower_ir(fn, compiler_options)
             # Perform optimizations.
+            do_box_unbox_elimination(fn, compiler_options)
             do_copy_propagation(fn, compiler_options)
             do_flag_elimination(fn, compiler_options)
 

mypyc/test-data/irbuild-dict.test

@@ -339,10 +339,8 @@ def union_of_dicts(d):
     r12 :: object[1]
     r13 :: object_ptr
     r14 :: object
-    r15 :: int
-    r16 :: object
-    r17 :: i32
-    r18, r19, r20 :: bit
+    r15 :: i32
+    r16, r17, r18 :: bit
 L0:
     r0 = PyDict_New()
     new = r0
@@ -367,15 +365,13 @@ L2:
     r13 = load_address r12
     r14 = PyObject_Vectorcall(r11, r13, 1, 0)
     keep_alive v
-    r15 = unbox(int, r14)
-    r16 = box(int, r15)
-    r17 = CPyDict_SetItem(new, k, r16)
-    r18 = r17 >= 0 :: signed
+    r15 = CPyDict_SetItem(new, k, r14)
+    r16 = r15 >= 0 :: signed
 L3:
-    r19 = CPyDict_CheckSize(d, r2)
+    r17 = CPyDict_CheckSize(d, r2)
     goto L1
 L4:
-    r20 = CPy_NoErrOccurred()
+    r18 = CPy_NoErrOccurred()
 L5:
     return 1
 def typeddict(d):

mypyc/test-data/irbuild-lists.test

@@ -724,13 +724,11 @@ def test():
     r0, source :: bytes
     r1 :: native_int
     r2 :: list
-    r3 :: native_int
-    r4, r5, r6 :: bit
-    r7, r8, r9, r10 :: int
-    r11 :: object
-    r12, x, r13 :: int
-    r14 :: object
-    r15 :: native_int
+    r3, r4 :: native_int
+    r5, r6, r7 :: bit
+    r8, r9, r10, r11, x, r12 :: int
+    r13 :: object
+    r14 :: native_int
     a :: list
 L0:
     r0 = b'abc'
@@ -755,16 +753,14 @@ L5:
     r9 = r3 << 1
     r8 = r9
 L6:
-    r10 = CPyBytes_GetItem(source, r8)
-    r11 = box(int, r10)
-    r12 = unbox(int, r11)
-    x = r12
-    r13 = f2(x)
-    r14 = box(int, r13)
-    CPyList_SetItemUnsafe(r2, r3, r14)
+    r11 = CPyBytes_GetItem(source, r9)
+    x = r11
+    r12 = f2(x)
+    r13 = box(int, r12)
+    CPyList_SetItemUnsafe(r2, r4, r13)
 L7:
-    r15 = r3 + 1
-    r3 = r15
+    r14 = r4 + 1
+    r4 = r14
     goto L1
 L8:
     a = r2
@@ -785,12 +781,11 @@ L0:
     return r0
 def test():
     r0 :: bytes
-    r1 :: list
-    r2 :: native_int
-    r3, r4, r5 :: bit
-    r6, r7, r8, r9 :: int
-    r10 :: object
-    r11, x, r12 :: int
+    r1 :: native_int
+    r2 :: list
+    r3, r4 :: native_int
+    r5, r6, r7 :: bit
+    r8, r9, r10, r11, x, r12 :: int
     r13 :: object
     r14 :: native_int
     a :: list
@@ -815,16 +810,14 @@ L5:
     r8 = r2 << 1
     r7 = r8
 L6:
-    r9 = CPyBytes_GetItem(r0, r7)
-    r10 = box(int, r9)
-    r11 = unbox(int, r10)
+    r11 = CPyBytes_GetItem(r0, r9)
     x = r11
     r12 = f2(x)
     r13 = box(int, r12)
-    CPyList_SetItemUnsafe(r1, r2, r13)
+    CPyList_SetItemUnsafe(r2, r4, r13)
 L7:
-    r14 = r2 + 1
-    r2 = r14
+    r14 = r4 + 1
+    r4 = r14
     goto L1
 L8:
     a = r1
@@ -852,13 +845,11 @@ def test():
     r1 :: bool
     r2 :: native_int
     r3 :: list
-    r4 :: native_int
-    r5, r6, r7 :: bit
-    r8, r9, r10, r11 :: int
-    r12 :: object
-    r13, x, r14 :: int
-    r15 :: object
-    r16 :: native_int
+    r4, r5 :: native_int
+    r6, r7, r8 :: bit
+    r9, r10, r11, r12, x, r13 :: int
+    r14 :: object
+    r15 :: native_int
     a :: list
 L0:
     r0 = __main__.source :: static
@@ -887,16 +878,14 @@ L7:
     r10 = r4 << 1
     r9 = r10
 L8:
-    r11 = CPyBytes_GetItem(r0, r9)
-    r12 = box(int, r11)
-    r13 = unbox(int, r12)
-    x = r13
-    r14 = f2(x)
-    r15 = box(int, r14)
-    CPyList_SetItemUnsafe(r3, r4, r15)
+    r12 = CPyBytes_GetItem(r0, r10)
+    x = r12
+    r13 = f2(x)
+    r14 = box(int, r13)
+    CPyList_SetItemUnsafe(r3, r5, r14)
 L9:
-    r16 = r4 + 1
-    r4 = r16
+    r15 = r5 + 1
+    r5 = r15
     goto L3
 L10:
     a = r3

mypyc/test-data/irbuild-set.test

@@ -649,29 +649,25 @@ def not_precomputed_non_final_name(i):
     r0 :: dict
     r1 :: str
     r2 :: object
-    r3 :: int
-    r4 :: set
-    r5 :: object
-    r6 :: i32
-    r7 :: bit
-    r8 :: object
-    r9 :: i32
-    r10 :: bit
-    r11 :: bool
+    r3 :: set
+    r4 :: i32
+    r5 :: bit
+    r6 :: object
+    r7 :: i32
+    r8 :: bit
+    r9 :: bool
 L0:
     r0 = __main__.globals :: static
     r1 = 'non_const'
     r2 = CPyDict_GetItem(r0, r1)
-    r3 = unbox(int, r2)
-    r4 = PySet_New(0)
-    r5 = box(int, r3)
-    r6 = PySet_Add(r4, r5)
-    r7 = r6 >= 0 :: signed
-    r8 = box(int, i)
-    r9 = PySet_Contains(r4, r8)
-    r10 = r9 >= 0 :: signed
-    r11 = truncate r9: i32 to builtins.bool
-    return r11
+    r3 = PySet_New(0)
+    r4 = PySet_Add(r3, r2)
+    r5 = r4 >= 0 :: signed
+    r6 = box(int, i)
+    r7 = PySet_Contains(r3, r6)
+    r8 = r7 >= 0 :: signed
+    r9 = truncate r7: i32 to builtins.bool
+    return r9
 def not_precomputed_nested_set(i):
     i :: int
     r0 :: set
@@ -769,33 +765,30 @@ def not_precomputed():
     r0 :: dict
     r1 :: str
     r2 :: object
-    r3 :: int
-    r4 :: set
-    r5 :: object
-    r6 :: i32
-    r7 :: bit
-    r8, r9 :: object
-    r10, not_optimized :: int
-    r11 :: bit
+    r3 :: set
+    r4 :: i32
+    r5 :: bit
+    r6, r7 :: object
+    r8, not_optimized :: int
+    r9 :: bit
 L0:
     r0 = __main__.globals :: static
     r1 = 'non_const'
     r2 = CPyDict_GetItem(r0, r1)
-    r3 = unbox(int, r2)
-    r4 = PySet_New(0)
-    r5 = box(int, r3)
-    r6 = PySet_Add(r4, r5)
-    r7 = r6 >= 0 :: signed
-    r8 = PyObject_GetIter(r4)
+    r3 = PySet_New(0)
+    r4 = PySet_Add(r3, r2)
+    r5 = r4 >= 0 :: signed
+    r6 = PyObject_GetIter(r3)
 L1:
-    r9 = PyIter_Next(r8)
-    if is_error(r9) goto L4 else goto L2
+    r7 = PyIter_Next(r6)
+    if is_error(r7) goto L4 else goto L2
 L2:
-    r10 = unbox(int, r9)
-    not_optimized = r10
+    r8 = unbox(int, r7)
+    not_optimized = r8
 L3:
     goto L1
 L4:
-    r11 = CPy_NoErrOccurred()
+    r9 = CPy_NoErrOccurred()
 L5:
     return 1
+

mypyc/test-data/irbuild-singledispatch.test

@@ -188,11 +188,10 @@ def f_obj.__call__(__mypyc_self__, x):
     r17 :: ptr
     r18 :: object
     r19 :: bit
-    r20 :: int
-    r21 :: object[1]
-    r22 :: object_ptr
-    r23 :: object
-    r24 :: None
+    r20 :: object[1]
+    r21 :: object_ptr
+    r22 :: object
+    r23 :: None
 L0:
     r0 = get_element_ptr x ob_type :: PyObject
     r1 = borrow load_mem r0 :: builtins.object*
@@ -225,15 +224,14 @@ L3:
     r19 = r18 == r16
     if r19 goto L4 else goto L5 :: bool
 L4:
-    r20 = unbox(int, r6)
     unreachable
 L5:
-    r21 = [x]
-    r22 = load_address r21
-    r23 = PyObject_Vectorcall(r6, r22, 1, 0)
+    r20 = [x]
+    r21 = load_address r20
+    r22 = PyObject_Vectorcall(r6, r21, 1, 0)
     keep_alive x
-    r24 = unbox(None, r23)
-    return r24
+    r23 = unbox(None, r22)
+    return r23
 def f_obj.__get__(__mypyc_self__, instance, owner):
     __mypyc_self__, instance, owner, r0 :: object
     r1 :: bit

mypyc/test-data/irbuild-statements.test

@@ -273,8 +273,7 @@ def f(d):
     r6 :: object
     r7, key :: int
     r8, r9 :: object
-    r10 :: int
-    r11, r12 :: bit
+    r10, r11 :: bit
 L0:
     r0 = 0
     r1 = PyDict_Size(d)
@@ -291,12 +290,11 @@ L2:
     key = r7
     r8 = box(int, key)
     r9 = CPyDict_GetItem(d, r8)
-    r10 = unbox(int, r9)
 L3:
-    r11 = CPyDict_CheckSize(d, r1)
+    r10 = CPyDict_CheckSize(d, r1)
     goto L1
 L4:
-    r12 = CPy_NoErrOccurred()
+    r11 = CPy_NoErrOccurred()
 L5:
     return 1