[mypyc] Add lowered primitive for unsafe list get item op

This inlines the list get item op in loops like `for x in <list>`.

I estimated the impact using two microbenchmarks that iterate over
`list[int]` objects. One of them was 1.3x faster, while the other
was 1.09x faster.

Since we now generate detailed IR for the op, instead of using a
C primitive function, this also opens up further IR optimization
opportunities in the future.

Author: JukkaL

mypyc/irbuild/builder.py

@@ -84,6 +84,7 @@
     IntOp,
     LoadStatic,
     Op,
+    PrimitiveDescription,
     RaiseStandardError,
     Register,
     SetAttr,
@@ -381,6 +382,15 @@ def load_module(self, name: str) -> Value:
     def call_c(self, desc: CFunctionDescription, args: list[Value], line: int) -> Value:
         return self.builder.call_c(desc, args, line)
 
+    def primitive_op(
+        self,
+        desc: PrimitiveDescription,
+        args: list[Value],
+        line: int,
+        result_type: RType | None = None,
+    ) -> Value:
+        return self.builder.primitive_op(desc, args, line, result_type)
+
     def int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int) -> Value:
         return self.builder.int_op(type, lhs, rhs, op, line)
 
@@ -739,7 +749,7 @@ def process_sequence_assignment(
             item = target.items[i]
             index = self.builder.load_int(i)
             if is_list_rprimitive(rvalue.type):
-                item_value = self.call_c(list_get_item_unsafe_op, [rvalue, index], line)
+                item_value = self.primitive_op(list_get_item_unsafe_op, [rvalue, index], line)
             else:
                 item_value = self.builder.gen_method_call(
                     rvalue, "__getitem__", [index], item.type, line

mypyc/irbuild/for_helpers.py

@@ -693,7 +693,7 @@ def unsafe_index(builder: IRBuilder, target: Value, index: Value, line: int) ->
     # since we want to use __getitem__ if we don't have an unsafe version,
     # so we just check manually.
     if is_list_rprimitive(target.type):
-        return builder.call_c(list_get_item_unsafe_op, [target, index], line)
+        return builder.primitive_op(list_get_item_unsafe_op, [target, index], line)
     else:
         return builder.gen_method_call(target, "__getitem__", [index], None, line)
 

mypyc/lower/list_ops.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from mypyc.common import PLATFORM_SIZE
-from mypyc.ir.ops import GetElementPtr, Integer, IntOp, LoadMem, SetMem, Value
+from mypyc.ir.ops import GetElementPtr, IncRef, Integer, IntOp, LoadMem, SetMem, Value
 from mypyc.ir.rtypes import (
     PyListObject,
     c_pyssize_t_rprimitive,
@@ -43,3 +43,24 @@ def buf_init_item(builder: LowLevelIRBuilder, args: list[Value], line: int) -> V
 def list_items(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
     ob_item_ptr = builder.add(GetElementPtr(args[0], PyListObject, "ob_item", line))
     return builder.add(LoadMem(pointer_rprimitive, ob_item_ptr, line))
+
+
+def list_item_ptr(builder: LowLevelIRBuilder, obj: Value, index: Value, line: int) -> Value:
+    """Get a pointer to a list item (index must be valid and non-negative).
+
+    Type of index must be c_pyssize_t_rprimitive.
+    """
+    items = list_items(builder, [obj], line)
+    delta = builder.add(
+        IntOp(c_pyssize_t_rprimitive, index, Integer(8, c_pyssize_t_rprimitive), IntOp.MUL)
+    )
+    return builder.add(IntOp(pointer_rprimitive, items, delta, IntOp.ADD))
+
+
+@lower_primitive_op("list_get_item_unsafe")
+def list_get_item_unsafe(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
+    index = builder.coerce(args[1], c_pyssize_t_rprimitive, line)
+    item_ptr = list_item_ptr(builder, args[0], index, line)
+    value = builder.add(LoadMem(object_rprimitive, item_ptr, line))
+    builder.add(IncRef(value))
+    return value

mypyc/primitives/list_ops.py

@@ -134,10 +134,10 @@
 
 # This is unsafe because it assumes that the index is a non-negative short integer
 # that is in-bounds for the list.
-list_get_item_unsafe_op = custom_op(
+list_get_item_unsafe_op = custom_primitive_op(
+    name="list_get_item_unsafe",
     arg_types=[list_rprimitive, short_int_rprimitive],
     return_type=object_rprimitive,
-    c_function_name="CPyList_GetItemUnsafe",
     error_kind=ERR_NEVER,
 )
 

mypyc/test-data/irbuild-basic.test

@@ -1874,7 +1874,7 @@ L1:
     r9 = int_lt r6, r8
     if r9 goto L2 else goto L8 :: bool
 L2:
-    r10 = CPyList_GetItemUnsafe(r1, r6)
+    r10 = list_get_item_unsafe r1, r6
     r11 = unbox(int, r10)
     x = r11
     r12 = int_ne x, 4
@@ -1938,7 +1938,7 @@ L1:
     r9 = int_lt r6, r8
     if r9 goto L2 else goto L8 :: bool
 L2:
-    r10 = CPyList_GetItemUnsafe(r1, r6)
+    r10 = list_get_item_unsafe r1, r6
     r11 = unbox(int, r10)
     x = r11
     r12 = int_ne x, 4
@@ -2000,7 +2000,7 @@ L1:
     r3 = int_lt r0, r2
     if r3 goto L2 else goto L4 :: bool
 L2:
-    r4 = CPyList_GetItemUnsafe(l, r0)
+    r4 = list_get_item_unsafe l, r0
     r5 = unbox(tuple[int, int, int], r4)
     r6 = r5[0]
     x = r6
@@ -2022,7 +2022,7 @@ L5:
     r15 = int_lt r12, r14
     if r15 goto L6 else goto L8 :: bool
 L6:
-    r16 = CPyList_GetItemUnsafe(l, r12)
+    r16 = list_get_item_unsafe l, r12
     r17 = unbox(tuple[int, int, int], r16)
     r18 = r17[0]
     x_2 = r18

mypyc/test-data/irbuild-lists.test

@@ -341,7 +341,7 @@ L1:
     r5 = int_lt r2, r4
     if r5 goto L2 else goto L4 :: bool
 L2:
-    r6 = CPyList_GetItemUnsafe(source, r2)
+    r6 = list_get_item_unsafe source, r2
     r7 = unbox(int, r6)
     x = r7
     r8 = CPyTagged_Add(x, 2)
@@ -362,7 +362,7 @@ L5:
     r17 = int_lt r14, r16
     if r17 goto L6 else goto L8 :: bool
 L6:
-    r18 = CPyList_GetItemUnsafe(source, r14)
+    r18 = list_get_item_unsafe source, r14
     r19 = unbox(int, r18)
     x_2 = r19
     r20 = CPyTagged_Add(x_2, 2)
@@ -403,7 +403,7 @@ L1:
     r3 = int_lt r0, r2
     if r3 goto L2 else goto L4 :: bool
 L2:
-    r4 = CPyList_GetItemUnsafe(x, r0)
+    r4 = list_get_item_unsafe x, r0
     r5 = unbox(int, r4)
     i = r5
     r6 = box(int, i)
@@ -476,7 +476,7 @@ L1:
     r3 = int_lt r0, r2
     if r3 goto L2 else goto L4 :: bool
 L2:
-    r4 = CPyList_GetItemUnsafe(a, r0)
+    r4 = list_get_item_unsafe a, r0
     r5 = cast(union[str, bytes], r4)
     x = r5
 L3:
@@ -502,7 +502,7 @@ L1:
     r3 = int_lt r0, r2
     if r3 goto L2 else goto L4 :: bool
 L2:
-    r4 = CPyList_GetItemUnsafe(a, r0)
+    r4 = list_get_item_unsafe a, r0
     r5 = cast(union[str, None], r4)
     x = r5
 L3:

mypyc/test-data/irbuild-set.test

@@ -115,7 +115,7 @@ L1:
     r9 = int_lt r6, r8
     if r9 goto L2 else goto L4 :: bool
 L2:
-    r10 = CPyList_GetItemUnsafe(tmp_list, r6)
+    r10 = list_get_item_unsafe tmp_list, r6
     r11 = unbox(int, r10)
     x = r11
     r12 = f(x)
@@ -361,7 +361,7 @@ L1:
     r13 = int_lt r10, r12
     if r13 goto L2 else goto L6 :: bool
 L2:
-    r14 = CPyList_GetItemUnsafe(tmp_list, r10)
+    r14 = list_get_item_unsafe tmp_list, r10
     r15 = unbox(int, r14)
     z = r15
     r16 = int_lt z, 8

mypyc/test-data/irbuild-statements.test

@@ -246,7 +246,7 @@ L1:
     r3 = int_lt r0, r2
     if r3 goto L2 else goto L4 :: bool
 L2:
-    r4 = CPyList_GetItemUnsafe(ls, r0)
+    r4 = list_get_item_unsafe ls, r0
     r5 = unbox(int, r4)
     x = r5
     r6 = CPyTagged_Add(y, x)
@@ -594,8 +594,8 @@ def f(l, t):
 L0:
     r0 = CPySequence_CheckUnpackCount(l, 2)
     r1 = r0 >= 0 :: signed
-    r2 = CPyList_GetItemUnsafe(l, 0)
-    r3 = CPyList_GetItemUnsafe(l, 2)
+    r2 = list_get_item_unsafe l, 0
+    r3 = list_get_item_unsafe l, 2
     x = r2
     r4 = unbox(int, r3)
     y = r4
@@ -883,7 +883,7 @@ L1:
     r4 = int_lt r1, r3
     if r4 goto L2 else goto L4 :: bool
 L2:
-    r5 = CPyList_GetItemUnsafe(a, r1)
+    r5 = list_get_item_unsafe a, r1
     r6 = unbox(int, r5)
     x = r6
     r7 = CPyTagged_Add(i, x)
@@ -965,7 +965,7 @@ L2:
     r5 = PyIter_Next(r1)
     if is_error(r5) goto L7 else goto L3
 L3:
-    r6 = CPyList_GetItemUnsafe(a, r0)
+    r6 = list_get_item_unsafe a, r0
     r7 = unbox(int, r6)
     x = r7
     r8 = unbox(bool, r5)
@@ -1019,7 +1019,7 @@ L3:
 L4:
     r8 = unbox(bool, r3)
     x = r8
-    r9 = CPyList_GetItemUnsafe(b, r1)
+    r9 = list_get_item_unsafe b, r1
     r10 = unbox(int, r9)
     y = r10
     x = 0

mypyc/test-data/irbuild-tuple.test

@@ -265,7 +265,7 @@ L1:
     r10 = int_lt r7, r9
     if r10 goto L2 else goto L4 :: bool
 L2:
-    r11 = CPyList_GetItemUnsafe(source, r7)
+    r11 = list_get_item_unsafe source, r7
     r12 = unbox(int, r11)
     x = r12
     r13 = f(x)

mypyc/test-data/lowering-int.test

@@ -346,10 +346,14 @@ def f(l):
     r2 :: native_int
     r3 :: short_int
     r4 :: bit
-    r5 :: object
-    r6, x :: int
-    r7 :: short_int
-    r8 :: None
+    r5 :: native_int
+    r6, r7 :: ptr
+    r8 :: native_int
+    r9 :: ptr
+    r10 :: object
+    r11, x :: int
+    r12 :: short_int
+    r13 :: None
 L0:
     r0 = 0
 L1:
@@ -359,19 +363,25 @@ L1:
     r4 = r0 < r3 :: signed
     if r4 goto L2 else goto L5 :: bool
 L2:
-    r5 = CPyList_GetItemUnsafe(l, r0)
-    r6 = unbox(int, r5)
-    dec_ref r5
-    if is_error(r6) goto L6 (error at f:4) else goto L3
+    r5 = r0 >> 1
+    r6 = get_element_ptr l ob_item :: PyListObject
+    r7 = load_mem r6 :: ptr*
+    r8 = r5 * 8
+    r9 = r7 + r8
+    r10 = load_mem r9 :: builtins.object*
+    inc_ref r10
+    r11 = unbox(int, r10)
+    dec_ref r10
+    if is_error(r11) goto L6 (error at f:4) else goto L3
 L3:
-    x = r6
+    x = r11
     dec_ref x :: int
 L4:
-    r7 = r0 + 2
-    r0 = r7
+    r12 = r0 + 2
+    r0 = r12
     goto L1
 L5:
     return 1
 L6:
-    r8 = <error> :: None
-    return r8
+    r13 = <error> :: None
+    return r13