[mypyc] Allow specifying primitives as pure (#17263)

Pure primitives have no side effects, take only immutable arguments,
and never fail. These properties will enable additional
optimizations. For example, it doesn't matter in which order
these primitives are evaluated, and we can perform common
subexpression elimination on them.

Only mark a few primitives as pure for now, but we can generalize
this later.

Author: JukkaL

mypyc/ir/ops.py

@@ -600,6 +600,7 @@ def __init__(
         ordering: list[int] | None,
         extra_int_constants: list[tuple[int, RType]],
         priority: int,
+        is_pure: bool,
     ) -> None:
         # Each primitive much have a distinct name, but otherwise they are arbitrary.
         self.name: Final = name
@@ -617,6 +618,11 @@ def __init__(
         self.ordering: Final = ordering
         self.extra_int_constants: Final = extra_int_constants
         self.priority: Final = priority
+        # Pure primitives have no side effects, take immutable arguments, and
+        # never fail. They support additional optimizations.
+        self.is_pure: Final = is_pure
+        if is_pure:
+            assert error_kind == ERR_NEVER
 
     def __repr__(self) -> str:
         return f"<PrimitiveDescription {self.name}>"
@@ -1036,6 +1042,8 @@ def __init__(
         error_kind: int,
         line: int,
         var_arg_idx: int = -1,
+        *,
+        is_pure: bool = False,
     ) -> None:
         self.error_kind = error_kind
         super().__init__(line)
@@ -1046,6 +1054,12 @@ def __init__(
         self.is_borrowed = is_borrowed
         # The position of the first variable argument in args (if >= 0)
         self.var_arg_idx = var_arg_idx
+        # Is the function pure? Pure functions have no side effects
+        # and all the arguments are immutable. Pure functions support
+        # additional optimizations. Pure functions never fail.
+        self.is_pure = is_pure
+        if is_pure:
+            assert error_kind == ERR_NEVER
 
     def sources(self) -> list[Value]:
         return self.args

mypyc/irbuild/ll_builder.py

@@ -1821,6 +1821,7 @@ def call_c(
                 error_kind,
                 line,
                 var_arg_idx,
+                is_pure=desc.is_pure,
             )
         )
         if desc.is_borrowed:
@@ -1903,6 +1904,7 @@ def primitive_op(
                 desc.ordering,
                 desc.extra_int_constants,
                 desc.priority,
+                is_pure=desc.is_pure,
             )
             return self.call_c(c_desc, args, line, result_type)
 

mypyc/primitives/int_ops.py

@@ -199,6 +199,7 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
     return_type=bit_rprimitive,
     c_function_name="CPyTagged_IsEq_",
     error_kind=ERR_NEVER,
+    is_pure=True,
 )
 
 # Less than operation on two boxed tagged integers
@@ -207,6 +208,7 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
     return_type=bit_rprimitive,
     c_function_name="CPyTagged_IsLt_",
     error_kind=ERR_NEVER,
+    is_pure=True,
 )
 
 int64_divide_op = custom_op(

mypyc/primitives/registry.py

@@ -60,6 +60,7 @@ class CFunctionDescription(NamedTuple):
     ordering: list[int] | None
     extra_int_constants: list[tuple[int, RType]]
     priority: int
+    is_pure: bool
 
 
 # A description for C load operations including LoadGlobal and LoadAddress
@@ -97,6 +98,7 @@ def method_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     priority: int = 1,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Define a c function call op that replaces a method call.
 
@@ -121,6 +123,8 @@ def method_op(
         steals: description of arguments that this steals (ref count wise)
         is_borrowed: if True, returned value is borrowed (no need to decrease refcount)
         priority: if multiple ops match, the one with the highest priority is picked
+        is_pure: if True, declare that the C function has no side effects, takes immutable
+                 arguments, and never raises an exception
     """
     if extra_int_constants is None:
         extra_int_constants = []
@@ -138,6 +142,7 @@ def method_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=is_pure,
     )
     ops.append(desc)
     return desc
@@ -183,6 +188,7 @@ def function_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=False,
     )
     ops.append(desc)
     return desc
@@ -228,6 +234,7 @@ def binary_op(
         ordering=ordering,
         extra_int_constants=extra_int_constants,
         priority=priority,
+        is_pure=False,
     )
     ops.append(desc)
     return desc
@@ -244,6 +251,8 @@ def custom_op(
     extra_int_constants: list[tuple[int, RType]] | None = None,
     steals: StealsDescription = False,
     is_borrowed: bool = False,
+    *,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Create a one-off CallC op that can't be automatically generated from the AST.
 
@@ -264,6 +273,7 @@ def custom_op(
         ordering,
         extra_int_constants,
         0,
+        is_pure=is_pure,
     )
 
 
@@ -279,6 +289,7 @@ def custom_primitive_op(
     extra_int_constants: list[tuple[int, RType]] | None = None,
     steals: StealsDescription = False,
     is_borrowed: bool = False,
+    is_pure: bool = False,
 ) -> PrimitiveDescription:
     """Define a primitive op that can't be automatically generated based on the AST.
 
@@ -299,6 +310,7 @@ def custom_primitive_op(
         ordering=ordering,
         extra_int_constants=extra_int_constants,
         priority=0,
+        is_pure=is_pure,
     )
 
 
@@ -314,6 +326,7 @@ def unary_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     priority: int = 1,
+    is_pure: bool = False,
 ) -> CFunctionDescription:
     """Define a c function call op for an unary operation.
 
@@ -338,6 +351,7 @@ def unary_op(
         ordering,
         extra_int_constants,
         priority,
+        is_pure=is_pure,
     )
     ops.append(desc)
     return desc