[mypyc] Basic test-only support for i32 and i64 (#13018)

Basic support for arithmetic, bitwise and comparison ops using native
int types i32 and i64. Also support some implicit coercions between
int and native int types. There are still various gaps in the functionality,
and that's why the feature is test-only so far.

Division and modulus ops require jumping through some hoops to make
the behavior with negative operands match Python semantics.

This can already show some improvements in benchmarks (using hacks
these can be used outside tests). I'll post figures once the implementation
is more complete.

Work on mypyc/mypyc#837.

Author: JukkaL

mypy/meet.py

@@ -131,6 +131,14 @@ def narrow_declared_type(declared: Type, narrowed: Type) -> Type:
         if declared.type.alt_promote:
             # Special case: low-level integer type can't be narrowed
             return original_declared
+        if (
+            isinstance(narrowed, Instance)
+            and narrowed.type.alt_promote
+            and narrowed.type.alt_promote is declared.type
+        ):
+            # Special case: 'int' can't be narrowed down to a native int type such as
+            # i64, since they have different runtime representations.
+            return original_declared
         return meet_types(original_declared, original_narrowed)
     elif isinstance(declared, (TupleType, TypeType, LiteralType)):
         return meet_types(original_declared, original_narrowed)

mypyc/common.py

@@ -43,6 +43,8 @@
 
 # Maximum value for a short tagged integer.
 MAX_SHORT_INT: Final = sys.maxsize >> 1
+# Minimum value for a short tagged integer.
+MIN_SHORT_INT: Final = -(sys.maxsize >> 1) - 1
 
 # Maximum value for a short tagged integer represented as a C integer literal.
 #

mypyc/irbuild/ast_helpers.py

@@ -21,7 +21,7 @@
     Var,
 )
 from mypyc.ir.ops import BasicBlock
-from mypyc.ir.rtypes import is_tagged
+from mypyc.ir.rtypes import is_fixed_width_rtype, is_tagged
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.constant_fold import constant_fold_expr
 
@@ -70,7 +70,10 @@ def maybe_process_conditional_comparison(
         return False
     ltype = self.node_type(e.operands[0])
     rtype = self.node_type(e.operands[1])
-    if not is_tagged(ltype) or not is_tagged(rtype):
+    if not (
+        (is_tagged(ltype) or is_fixed_width_rtype(ltype))
+        and (is_tagged(rtype) or is_fixed_width_rtype(rtype))
+    ):
         return False
     op = e.operators[0]
     if op not in ("==", "!=", "<", "<=", ">", ">="):
@@ -80,8 +83,17 @@ def maybe_process_conditional_comparison(
     borrow_left = is_borrow_friendly_expr(self, right_expr)
     left = self.accept(left_expr, can_borrow=borrow_left)
     right = self.accept(right_expr, can_borrow=True)
-    # "left op right" for two tagged integers
-    self.builder.compare_tagged_condition(left, right, op, true, false, e.line)
+    if is_fixed_width_rtype(ltype) or is_fixed_width_rtype(rtype):
+        if not is_fixed_width_rtype(ltype):
+            left = self.coerce(left, rtype, e.line)
+        elif not is_fixed_width_rtype(rtype):
+            right = self.coerce(right, ltype, e.line)
+        reg = self.binary_op(left, right, op, e.line)
+        self.builder.flush_keep_alives()
+        self.add_bool_branch(reg, true, false)
+    else:
+        # "left op right" for two tagged integers
+        self.builder.compare_tagged_condition(left, right, op, true, false, e.line)
     return True
 
 

mypyc/irbuild/expression.py

@@ -52,6 +52,8 @@
 from mypyc.ir.ops import (
     Assign,
     BasicBlock,
+    ComparisonOp,
+    Integer,
     LoadAddress,
     RaiseStandardError,
     Register,
@@ -62,6 +64,7 @@
 from mypyc.ir.rtypes import (
     RTuple,
     int_rprimitive,
+    is_fixed_width_rtype,
     is_int_rprimitive,
     is_list_rprimitive,
     is_none_rprimitive,
@@ -472,21 +475,26 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     if folded:
         return folded
 
+    borrow_left = False
+    borrow_right = False
+
+    ltype = builder.node_type(expr.left)
+    rtype = builder.node_type(expr.right)
+
     # Special case some int ops to allow borrowing operands.
-    if is_int_rprimitive(builder.node_type(expr.left)) and is_int_rprimitive(
-        builder.node_type(expr.right)
-    ):
+    if is_int_rprimitive(ltype) and is_int_rprimitive(rtype):
         if expr.op == "//":
             expr = try_optimize_int_floor_divide(expr)
         if expr.op in int_borrow_friendly_op:
             borrow_left = is_borrow_friendly_expr(builder, expr.right)
-            left = builder.accept(expr.left, can_borrow=borrow_left)
-            right = builder.accept(expr.right, can_borrow=True)
-            return builder.binary_op(left, right, expr.op, expr.line)
+            borrow_right = True
+    elif is_fixed_width_rtype(ltype) and is_fixed_width_rtype(rtype):
+        borrow_left = is_borrow_friendly_expr(builder, expr.right)
+        borrow_right = True
 
-    return builder.binary_op(
-        builder.accept(expr.left), builder.accept(expr.right), expr.op, expr.line
-    )
+    left = builder.accept(expr.left, can_borrow=borrow_left)
+    right = builder.accept(expr.right, can_borrow=borrow_right)
+    return builder.binary_op(left, right, expr.op, expr.line)
 
 
 def try_optimize_int_floor_divide(expr: OpExpr) -> OpExpr:
@@ -708,6 +716,25 @@ def transform_basic_comparison(
         and op in int_comparison_op_mapping.keys()
     ):
         return builder.compare_tagged(left, right, op, line)
+    if is_fixed_width_rtype(left.type) and op in int_comparison_op_mapping.keys():
+        if right.type == left.type:
+            op_id = ComparisonOp.signed_ops[op]
+            return builder.builder.comparison_op(left, right, op_id, line)
+        elif isinstance(right, Integer):
+            op_id = ComparisonOp.signed_ops[op]
+            return builder.builder.comparison_op(
+                left, Integer(right.value >> 1, left.type), op_id, line
+            )
+    elif (
+        is_fixed_width_rtype(right.type)
+        and op in int_comparison_op_mapping.keys()
+        and isinstance(left, Integer)
+    ):
+        op_id = ComparisonOp.signed_ops[op]
+        return builder.builder.comparison_op(
+            Integer(left.value >> 1, right.type), right, op_id, line
+        )
+
     negate = False
     if op == "is not":
         op, negate = "is", True

mypyc/irbuild/for_helpers.py

@@ -38,6 +38,7 @@
     bool_rprimitive,
     int_rprimitive,
     is_dict_rprimitive,
+    is_fixed_width_rtype,
     is_list_rprimitive,
     is_sequence_rprimitive,
     is_short_int_rprimitive,
@@ -887,7 +888,9 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         self.step = step
         self.end_target = builder.maybe_spill(end_reg)
         if is_short_int_rprimitive(start_reg.type) and is_short_int_rprimitive(end_reg.type):
-            index_type = short_int_rprimitive
+            index_type: RType = short_int_rprimitive
+        elif is_fixed_width_rtype(end_reg.type):
+            index_type = end_reg.type
         else:
             index_type = int_rprimitive
         index_reg = Register(index_type)