[mypyc] Use a native unboxed representation for floats (#14880)

Instead of each float value being a heap-allocated Python object, use
unboxed C doubles to represent floats. This makes float operations much
faster, and this also significantly reduces memory use of floats (when
not stored in Python containers, which always use a boxed
representation).

Update IR to support float arithmetic and comparison ops, and float
literals.

Also add a few primitives corresponding to common math functions, such
as `math.sqrt`. These don't require any boxing or unboxing. (I will add
more of these in follow-up PRs.)

Use -113.0 as an overlapping error value for floats. This is similar to
native ints. Reuse much of the infrastructure we have to support
overlapping error values with native ints (e.g. various bitmaps).

Also improve support for negative float literals.

There are two backward compatibility breaks worth highlighting. First,
assigning an int value to a float variable is disallowed within mypyc,
since narrowing down to a different value representation is inefficient
and can lose precision. Second, information about float subclasses is
lost during unboxing.

This makes the bm_float benchmark about 5x faster and the raytrace
benchmark about 3x faster.

Closes mypyc/mypyc#966 (I'll create separate
issues for remaining open issues).

Author: JukkaL

mypy/constant_fold.py

@@ -64,6 +64,8 @@ def constant_fold_expr(expr: Expression, cur_mod_id: str) -> ConstantValue | Non
         value = constant_fold_expr(expr.expr, cur_mod_id)
         if isinstance(value, int):
             return constant_fold_unary_int_op(expr.op, value)
+        if isinstance(value, float):
+            return constant_fold_unary_float_op(expr.op, value)
     return None
 
 
@@ -110,6 +112,14 @@ def constant_fold_unary_int_op(op: str, value: int) -> int | None:
     return None
 
 
+def constant_fold_unary_float_op(op: str, value: float) -> float | None:
+    if op == "-":
+        return -value
+    elif op == "+":
+        return value
+    return None
+
+
 def constant_fold_binary_str_op(op: str, left: str, right: str) -> str | None:
     if op == "+":
         return left + right

mypyc/analysis/dataflow.py

@@ -18,6 +18,10 @@
     ComparisonOp,
     ControlOp,
     Extend,
+    Float,
+    FloatComparisonOp,
+    FloatNeg,
+    FloatOp,
     GetAttr,
     GetElementPtr,
     Goto,
@@ -245,9 +249,18 @@ def visit_load_global(self, op: LoadGlobal) -> GenAndKill[T]:
     def visit_int_op(self, op: IntOp) -> GenAndKill[T]:
         return self.visit_register_op(op)
 
+    def visit_float_op(self, op: FloatOp) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
+    def visit_float_neg(self, op: FloatNeg) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
     def visit_comparison_op(self, op: ComparisonOp) -> GenAndKill[T]:
         return self.visit_register_op(op)
 
+    def visit_float_comparison_op(self, op: FloatComparisonOp) -> GenAndKill[T]:
+        return self.visit_register_op(op)
+
     def visit_load_mem(self, op: LoadMem) -> GenAndKill[T]:
         return self.visit_register_op(op)
 
@@ -444,7 +457,7 @@ def analyze_undefined_regs(
 def non_trivial_sources(op: Op) -> set[Value]:
     result = set()
     for source in op.sources():
-        if not isinstance(source, Integer):
+        if not isinstance(source, (Integer, Float)):
             result.add(source)
     return result
 
@@ -454,7 +467,7 @@ def visit_branch(self, op: Branch) -> GenAndKill[Value]:
         return non_trivial_sources(op), set()
 
     def visit_return(self, op: Return) -> GenAndKill[Value]:
-        if not isinstance(op.value, Integer):
+        if not isinstance(op.value, (Integer, Float)):
             return {op.value}, set()
         else:
             return set(), set()

mypyc/analysis/ircheck.py

@@ -16,6 +16,9 @@
     ControlOp,
     DecRef,
     Extend,
+    FloatComparisonOp,
+    FloatNeg,
+    FloatOp,
     GetAttr,
     GetElementPtr,
     Goto,
@@ -43,6 +46,7 @@
     TupleSet,
     Unbox,
     Unreachable,
+    Value,
 )
 from mypyc.ir.pprint import format_func
 from mypyc.ir.rtypes import (
@@ -54,6 +58,7 @@
     bytes_rprimitive,
     dict_rprimitive,
     int_rprimitive,
+    is_float_rprimitive,
     is_object_rprimitive,
     list_rprimitive,
     range_rprimitive,
@@ -221,6 +226,14 @@ def check_compatibility(self, op: Op, t: RType, s: RType) -> None:
         if not can_coerce_to(t, s) or not can_coerce_to(s, t):
             self.fail(source=op, desc=f"{t.name} and {s.name} are not compatible")
 
+    def expect_float(self, op: Op, v: Value) -> None:
+        if not is_float_rprimitive(v.type):
+            self.fail(op, f"Float expected (actual type is {v.type})")
+
+    def expect_non_float(self, op: Op, v: Value) -> None:
+        if is_float_rprimitive(v.type):
+            self.fail(op, "Float not expected")
+
     def visit_goto(self, op: Goto) -> None:
         self.check_control_op_targets(op)
 
@@ -376,10 +389,24 @@ def visit_load_global(self, op: LoadGlobal) -> None:
         pass
 
     def visit_int_op(self, op: IntOp) -> None:
-        pass
+        self.expect_non_float(op, op.lhs)
+        self.expect_non_float(op, op.rhs)
 
     def visit_comparison_op(self, op: ComparisonOp) -> None:
         self.check_compatibility(op, op.lhs.type, op.rhs.type)
+        self.expect_non_float(op, op.lhs)
+        self.expect_non_float(op, op.rhs)
+
+    def visit_float_op(self, op: FloatOp) -> None:
+        self.expect_float(op, op.lhs)
+        self.expect_float(op, op.rhs)
+
+    def visit_float_neg(self, op: FloatNeg) -> None:
+        self.expect_float(op, op.src)
+
+    def visit_float_comparison_op(self, op: FloatComparisonOp) -> None:
+        self.expect_float(op, op.lhs)
+        self.expect_float(op, op.rhs)
 
     def visit_load_mem(self, op: LoadMem) -> None:
         pass

mypyc/analysis/selfleaks.py

@@ -14,6 +14,9 @@
     Cast,
     ComparisonOp,
     Extend,
+    FloatComparisonOp,
+    FloatNeg,
+    FloatOp,
     GetAttr,
     GetElementPtr,
     Goto,
@@ -160,6 +163,15 @@ def visit_int_op(self, op: IntOp) -> GenAndKill:
     def visit_comparison_op(self, op: ComparisonOp) -> GenAndKill:
         return CLEAN
 
+    def visit_float_op(self, op: FloatOp) -> GenAndKill:
+        return CLEAN
+
+    def visit_float_neg(self, op: FloatNeg) -> GenAndKill:
+        return CLEAN
+
+    def visit_float_comparison_op(self, op: FloatComparisonOp) -> GenAndKill:
+        return CLEAN
+
     def visit_load_mem(self, op: LoadMem) -> GenAndKill:
         return CLEAN
 

mypyc/codegen/emit.py

@@ -895,6 +895,16 @@ def emit_unbox(
             self.emit_line(f"{dest} = CPyLong_AsInt32({src});")
             # TODO: Handle 'optional'
             # TODO: Handle 'failure'
+        elif is_float_rprimitive(typ):
+            if declare_dest:
+                self.emit_line("double {};".format(dest))
+            # TODO: Don't use __float__ and __index__
+            self.emit_line(f"{dest} = PyFloat_AsDouble({src});")
+            self.emit_lines(
+                f"if ({dest} == -1.0 && PyErr_Occurred()) {{", f"{dest} = -113.0;", "}"
+            )
+            # TODO: Handle 'optional'
+            # TODO: Handle 'failure'
         elif isinstance(typ, RTuple):
             self.declare_tuple_struct(typ)
             if declare_dest:
@@ -983,6 +993,8 @@ def emit_box(
             self.emit_line(f"{declaration}{dest} = PyLong_FromLong({src});")
         elif is_int64_rprimitive(typ):
             self.emit_line(f"{declaration}{dest} = PyLong_FromLongLong({src});")
+        elif is_float_rprimitive(typ):
+            self.emit_line(f"{declaration}{dest} = PyFloat_FromDouble({src});")
         elif isinstance(typ, RTuple):
             self.declare_tuple_struct(typ)
             self.emit_line(f"{declaration}{dest} = PyTuple_New({len(typ.types)});")

mypyc/codegen/emitclass.py

@@ -1004,6 +1004,7 @@ def generate_readonly_getter(
                 emitter.ctype_spaced(rtype), NATIVE_PREFIX, func_ir.cname(emitter.names)
             )
         )
+        emitter.emit_error_check("retval", rtype, "return NULL;")
         emitter.emit_box("retval", "retbox", rtype, declare_dest=True)
         emitter.emit_line("return retbox;")
     else:

mypyc/codegen/emitfunc.py

@@ -25,6 +25,10 @@
     ComparisonOp,
     DecRef,
     Extend,
+    Float,
+    FloatComparisonOp,
+    FloatNeg,
+    FloatOp,
     GetAttr,
     GetElementPtr,
     Goto,
@@ -671,6 +675,27 @@ def visit_comparison_op(self, op: ComparisonOp) -> None:
             lhs_cast = self.emit_signed_int_cast(op.lhs.type)
         self.emit_line(f"{dest} = {lhs_cast}{lhs} {op.op_str[op.op]} {rhs_cast}{rhs};")
 
+    def visit_float_op(self, op: FloatOp) -> None:
+        dest = self.reg(op)
+        lhs = self.reg(op.lhs)
+        rhs = self.reg(op.rhs)
+        if op.op != FloatOp.MOD:
+            self.emit_line("%s = %s %s %s;" % (dest, lhs, op.op_str[op.op], rhs))
+        else:
+            # TODO: This may set errno as a side effect, that is a little sketchy.
+            self.emit_line("%s = fmod(%s, %s);" % (dest, lhs, rhs))
+
+    def visit_float_neg(self, op: FloatNeg) -> None:
+        dest = self.reg(op)
+        src = self.reg(op.src)
+        self.emit_line(f"{dest} = -{src};")
+
+    def visit_float_comparison_op(self, op: FloatComparisonOp) -> None:
+        dest = self.reg(op)
+        lhs = self.reg(op.lhs)
+        rhs = self.reg(op.rhs)
+        self.emit_line("%s = %s %s %s;" % (dest, lhs, op.op_str[op.op], rhs))
+
     def visit_load_mem(self, op: LoadMem) -> None:
         dest = self.reg(op)
         src = self.reg(op.src)
@@ -732,6 +757,13 @@ def reg(self, reg: Value) -> str:
             elif val <= -(1 << 31):
                 s += "LL"
             return s
+        elif isinstance(reg, Float):
+            r = repr(reg.value)
+            if r == "inf":
+                return "INFINITY"
+            elif r == "-inf":
+                return "-INFINITY"
+            return r
         else:
             return self.emitter.reg(reg)
 

mypyc/common.py

@@ -69,6 +69,7 @@
     "getargs.c",
     "getargsfast.c",
     "int_ops.c",
+    "float_ops.c",
     "str_ops.c",
     "bytes_ops.c",
     "list_ops.c",