[mypyc] Support generic primitives and add some generic primitives for vec (#21656)

Generic primitives have `RTypeVar` types in parameter and/or return
types, and these get expanded away when the primitive is added to IR.

Generic primitives let us use lowering for various `vec` operations,
many of which are generic. Using higher-level operations in the IR helps
with various optimizations. It's also easier to verify that the
generated IR is correct when the IR is less verbose.

Add generic primitives for unsafe `vec` get item op as an initial use
case. We can later use these for other `vec` operations as well.

Used some coding agent assist (mostly for tests).

Author: JukkaL

mypyc/ir/ops.py

@@ -39,6 +39,7 @@ class to enable the new behavior. Sometimes adding a new abstract
     RStruct,
     RTuple,
     RType,
+    RTypeVar,
     RUnion,
     RVec,
     RVoid,
@@ -703,7 +704,7 @@ def __init__(
         self,
         name: str,
         arg_types: list[RType],
-        return_type: RType,  # TODO: What about generic?
+        return_type: RType,
         var_arg_type: RType | None,
         truncated_type: RType | None,
         c_function_name: str | None,
@@ -716,6 +717,7 @@ def __init__(
         is_pure: bool,
         experimental: bool,
         dependencies: list[Dependency] | None,
+        type_params: list[RTypeVar] | None,
     ) -> None:
         # Each primitive much have a distinct name, but otherwise they are arbitrary.
         self.name: Final = name
@@ -749,6 +751,7 @@ def __init__(
         # If this flag is set, the primitive has native integer types and must
         # be matched using more complex rules.
         self.is_ambiguous = any(has_fixed_width_int(t) for t in arg_types)
+        self.type_params = None if not type_params else type_params
 
     def __repr__(self) -> str:
         return f"<PrimitiveDescription {self.name!r}: {self.arg_types}>"
@@ -776,11 +779,23 @@ class PrimitiveOp(RegisterOp):
     code paths for short and long representations.
     """
 
-    def __init__(self, args: list[Value], desc: PrimitiveDescription, line: int = -1) -> None:
+    def __init__(
+        self,
+        args: list[Value],
+        desc: PrimitiveDescription,
+        line: int = -1,
+        *,
+        arg_types: list[RType] | None = None,
+        return_type: RType | None = None,
+        type_args: list[RType] | None = None,
+    ) -> None:
         self.error_kind = desc.error_kind
         super().__init__(line)
         self.args = args
-        self.type = desc.return_type
+        self.arg_types = arg_types if arg_types is not None else desc.arg_types
+        self.type = return_type if return_type is not None else desc.return_type
+        self.is_borrowed = desc.is_borrowed
+        self.type_args = type_args
         self.desc = desc
 
     def sources(self) -> list[Value]:

mypyc/ir/pprint.py

@@ -231,10 +231,15 @@ def visit_call_c(self, op: CallC) -> str:
 
     def visit_primitive_op(self, op: PrimitiveOp) -> str:
         args_str = ", ".join(self.format("%r", arg) for arg in op.args)
+        if op.type_args:
+            joined = ", ".join(str(arg) for arg in op.type_args)
+            type_args = f"[{joined}]"
+        else:
+            type_args = ""
         if op.is_void:
-            return self.format("%s %s", op.desc.name, args_str)
+            return self.format("%s%s %s", op.desc.name, type_args, args_str)
         else:
-            return self.format("%r = %s %s", op, op.desc.name, args_str)
+            return self.format("%r = %s%s %s", op, op.desc.name, type_args, args_str)
 
     def visit_truncate(self, op: Truncate) -> str:
         return self.format("%r = truncate %r: %t to %t", op, op.src, op.src_type, op.type)

mypyc/ir/rtypes.py

@@ -152,6 +152,9 @@ def visit_rprimitive(self, typ: RPrimitive, /) -> T:
     def visit_rinstance(self, typ: RInstance, /) -> T:
         raise NotImplementedError
 
+    def visit_rtypevar(self, typ: RTypeVar, /) -> T:
+        raise RuntimeError("RTypeVar should not be encountered here")
+
     @abstractmethod
     def visit_rvec(self, typ: RVec, /) -> T:
         raise NotImplementedError
@@ -747,6 +750,12 @@ def visit_rarray(self, t: RArray) -> str:
     def visit_rvoid(self, t: RVoid) -> str:
         assert False, "rvoid in tuple?"
 
+    def visit_rtypevar(self, typ: RTypeVar) -> str:
+        # We need to return something to support generic RTuples, etc. Make sure
+        # the return value is invalid C so that generic RTuples must be expanded
+        # before they can be used in IR.
+        return f"!RTypeVar {typ.id} invalid!"
+
 
 @final
 class RTuple(RType):
@@ -1013,6 +1022,50 @@ def serialize(self) -> str:
         return self.name
 
 
+@final
+class RTypeVar(RType):
+    """Type variable type used for generic primitive ops.
+
+    This allows having generic primitive operations like vec get item, which is
+    parametrized by the vec item type.
+
+    These types are not valid in any other context outside PrimitiveDescription,
+    and they will always be substituted during the construction of a PrimitiveOp.
+
+    NOTE: This is not related to mypy's TypeVarType!
+    """
+
+    def __init__(self, id: int) -> None:
+        self.id = id
+
+    @property
+    def may_be_immortal(self) -> bool:
+        # RTypeVar must always be substituted before use, so this should never matter.
+        return False
+
+    def accept(self, visitor: RTypeVisitor[T]) -> T:
+        return visitor.visit_rtypevar(self)
+
+    def __str__(self) -> str:
+        return f"<RTypeVar {self.id}>"
+
+    def __repr__(self) -> str:
+        return f"<RTypeVar {self.id}>"
+
+    def __eq__(self, other: object) -> TypeGuard[RTypeVar]:
+        return isinstance(other, RTypeVar) and other.id == self.id
+
+    def __hash__(self) -> int:
+        return self.id ^ 12345
+
+    def serialize(self) -> JsonDict:
+        return {".class": "RTypeVar", "id": self.id}
+
+    @classmethod
+    def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> RTypeVar:
+        return RTypeVar(data["id"])
+
+
 @final
 class RVec(RType):
     """librt.vecs.vec[T]"""

mypyc/irbuild/ll_builder.py

@@ -206,6 +206,7 @@
     new_tuple_with_length_op,
     sequence_tuple_op,
 )
+from mypyc.rt_expandtype import expand_rtype
 from mypyc.rt_subtype import is_runtime_subtype
 from mypyc.sametype import is_same_type
 from mypyc.subtype import is_subtype
@@ -2334,6 +2335,7 @@ def primitive_op(
         args: list[Value],
         line: int,
         result_type: RType | None = None,
+        type_args: list[RType] | None = None,
     ) -> Value:
         """Add a primitive op."""
         # Does this primitive map into calling a Python C API
@@ -2363,18 +2365,37 @@ def primitive_op(
         # This primitive gets transformed in a lowering pass to
         # lower-level IR ops using a custom transform function.
 
+        # Evaluate argument and return types for generic primitives
+        return_type = None
+        if desc.type_params is not None:
+            assert type_args is not None, "Generic primitive op requires explicit type arguments"
+            assert len(type_args) == len(desc.type_params)
+            arg_types = [expand_rtype(arg_type, type_args) for arg_type in desc.arg_types]
+            return_type = expand_rtype(desc.return_type, type_args)
+        else:
+            arg_types = desc.arg_types
+
         coerced = []
         # Coerce fixed number arguments
-        for i in range(min(len(args), len(desc.arg_types))):
-            formal_type = desc.arg_types[i]
+        for i in range(min(len(args), len(arg_types))):
+            formal_type = arg_types[i]
             arg = args[i]
             assert formal_type is not None  # TODO
             arg = self.coerce(arg, formal_type, line)
             coerced.append(arg)
         assert desc.ordering is None
         assert desc.var_arg_type is None
         assert not desc.extra_int_constants
-        target = self.add(PrimitiveOp(coerced, desc, line=line))
+        target = self.add(
+            PrimitiveOp(
+                coerced,
+                desc,
+                line=line,
+                arg_types=arg_types,
+                return_type=return_type,
+                type_args=type_args,
+            )
+        )
         if desc.is_borrowed:
             # If the result is borrowed, force the arguments to be
             # kept alive afterwards, as otherwise the result might be

mypyc/irbuild/vec.py

@@ -52,6 +52,7 @@
     vec_api_by_item_type,
     vec_item_type_tags,
 )
+from mypyc.primitives.librt_vecs_ops import vec_get_item_unsafe_borrow_op, vec_get_item_unsafe_op
 
 if TYPE_CHECKING:
     from mypyc.irbuild.ll_builder import LowLevelIRBuilder
@@ -316,9 +317,10 @@ def vec_get_item(
 ) -> Value:
     """Generate inlined vec __getitem__ call.
 
-    We inline this, since it's simple but performance-critical.
+    We inline the length and bounds check, since they are simple but
+    performance-critical. The actual item load is emitted as a generic primitive
+    op that is lowered later.
     """
-    # TODO: Support more item types
     # TODO: Support more index types
     len_val = vec_len(builder, base)
     index = vec_check_and_adjust_index(builder, len_val, index, line)
@@ -328,7 +330,22 @@ def vec_get_item(
 def vec_get_item_unsafe(
     builder: LowLevelIRBuilder, base: Value, index: Value, line: int, *, can_borrow: bool = False
 ) -> Value:
-    """Get vec item, assuming index is non-negative and within bounds."""
+    """Get vec item, assuming index is non-negative and within bounds.
+
+    This emits a generic primitive op that is inlined during lowering.
+    """
+    assert isinstance(base.type, RVec)
+    if can_borrow:
+        desc = vec_get_item_unsafe_borrow_op
+    else:
+        desc = vec_get_item_unsafe_op
+    return builder.primitive_op(desc, [base, index], line, type_args=[base.type.item_type])
+
+
+def vec_get_item_unsafe_lower(
+    builder: LowLevelIRBuilder, base: Value, index: Value, line: int, *, can_borrow: bool = False
+) -> Value:
+    """Generate the low-level IR for an unsafe vec item load."""
     assert isinstance(base.type, RVec)
     index = as_platform_int(builder, index, line)
     vtype = base.type

mypyc/lower/registry.py

@@ -26,4 +26,4 @@ def wrapper(f: LF) -> LF:
 
 
 # Import various modules that set up global state.
-from mypyc.lower import int_ops, list_ops, misc_ops  # noqa: F401
+from mypyc.lower import int_ops, list_ops, misc_ops, vec_ops  # noqa: F401

mypyc/lower/vec_ops.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from mypyc.ir.ops import Value
+from mypyc.irbuild.ll_builder import LowLevelIRBuilder
+from mypyc.irbuild.vec import vec_get_item_unsafe_lower
+from mypyc.lower.registry import lower_primitive_op
+
+
+@lower_primitive_op("vec_get_item_unsafe")
+def vec_get_item_unsafe(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
+    base, index = args
+    return vec_get_item_unsafe_lower(builder, base, index, line, can_borrow=False)
+
+
+@lower_primitive_op("vec_get_item_unsafe_borrow")
+def vec_get_item_unsafe_borrow(builder: LowLevelIRBuilder, args: list[Value], line: int) -> Value:
+    base, index = args
+    return vec_get_item_unsafe_lower(builder, base, index, line, can_borrow=True)

mypyc/primitives/librt_vecs_ops.py

@@ -1,13 +1,15 @@
 from mypyc.ir.deps import LIBRT_VECS, VECS_EXTRA_OPS
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
+    RTypeVar,
     RVec,
     bit_rprimitive,
     bytes_rprimitive,
+    int64_rprimitive,
     object_rprimitive,
     uint8_rprimitive,
 )
-from mypyc.primitives.registry import function_op
+from mypyc.primitives.registry import custom_primitive_op, function_op
 
 # isinstance(obj, vec)
 isinstance_vec = function_op(
@@ -28,3 +30,22 @@
     error_kind=ERR_MAGIC,
     dependencies=[LIBRT_VECS, VECS_EXTRA_OPS],
 )
+
+# Get vec item, assuming the index is valid (no bounds check)
+vec_get_item_unsafe_op = custom_primitive_op(
+    name="vec_get_item_unsafe",
+    arg_types=[RVec(RTypeVar(0)), int64_rprimitive],
+    return_type=RTypeVar(0),
+    error_kind=ERR_NEVER,
+    type_params=[RTypeVar(0)],
+)
+
+# Like vec_get_item_unsafe, but the result is a borrowed reference
+vec_get_item_unsafe_borrow_op = custom_primitive_op(
+    name="vec_get_item_unsafe_borrow",
+    arg_types=[RVec(RTypeVar(0)), int64_rprimitive],
+    is_borrowed=True,
+    return_type=RTypeVar(0),
+    error_kind=ERR_NEVER,
+    type_params=[RTypeVar(0)],
+)

mypyc/primitives/registry.py

@@ -41,7 +41,7 @@
 
 from mypyc.ir.deps import Dependency
 from mypyc.ir.ops import PrimitiveDescription, StealsDescription
-from mypyc.ir.rtypes import RType
+from mypyc.ir.rtypes import RType, RTypeVar
 
 # Error kind for functions that return negative integer on exception. This
 # is only used for primitives. We translate it away during IR building.
@@ -154,6 +154,7 @@ def method_op(
         is_pure=is_pure,
         experimental=experimental,
         dependencies=dependencies,
+        type_params=None,
     )
     ops.append(desc)
     return desc
@@ -204,6 +205,7 @@ def function_op(
         is_pure=False,
         experimental=experimental,
         dependencies=dependencies,
+        type_params=None,
     )
     ops.append(desc)
     return desc
@@ -253,6 +255,7 @@ def binary_op(
         is_pure=False,
         experimental=False,
         dependencies=dependencies,
+        type_params=None,
     )
     ops.append(desc)
     return desc
@@ -313,6 +316,7 @@ def custom_primitive_op(
     is_pure: bool = False,
     experimental: bool = False,
     dependencies: list[Dependency] | None = None,
+    type_params: list[RTypeVar] | None = None,
 ) -> PrimitiveDescription:
     """Define a primitive op that can't be automatically generated based on the AST.
 
@@ -336,6 +340,7 @@ def custom_primitive_op(
         is_pure=is_pure,
         experimental=experimental,
         dependencies=dependencies,
+        type_params=type_params,
     )
 
 
@@ -380,6 +385,7 @@ def unary_op(
         is_pure=is_pure,
         experimental=False,
         dependencies=dependencies,
+        type_params=None,
     )
     ops.append(desc)
     return desc