Fix double-rounding bug in double -> (b)float16 casts

src/CodeGen_X86.cpp

@@ -524,7 +524,8 @@ void CodeGen_X86::visit(const Cast *op) {
     if (target.has_feature(Target::F16C) &&
         dst.code() == Type::Float &&
         src.code() == Type::Float &&
-        (dst.bits() == 16 || src.bits() == 16)) {
+        (dst.bits() == 16 || src.bits() == 16) &&
+        src.bits() <= 32) {  // Don't use for narrowing casts from double - it results in a libm call
         // Node we use code() == Type::Float instead of is_float(), because we
         // don't want to catch bfloat casts.
 

src/EmulateFloat16Math.cpp

@@ -9,27 +9,46 @@ namespace Halide {
 namespace Internal {
 
 Expr bfloat16_to_float32(Expr e) {
+    const int lanes = e.type().lanes();
     if (e.type().is_bfloat()) {
         e = reinterpret(e.type().with_code(Type::UInt), e);
     }
-    e = cast(UInt(32, e.type().lanes()), e);
+    e = cast(UInt(32, lanes), e);
     e = e << 16;
-    e = reinterpret(Float(32, e.type().lanes()), e);
+    e = reinterpret(Float(32, lanes), e);
     e = strict_float(e);
     return e;
 }
 
 Expr float32_to_bfloat16(Expr e) {
     internal_assert(e.type().bits() == 32);
+    const int lanes = e.type().lanes();
     e = strict_float(e);
-    e = reinterpret(UInt(32, e.type().lanes()), e);
+    e = reinterpret(UInt(32, lanes), e);
     // We want to round ties to even, so before truncating either
     // add 0x8000 (0.5) to odd numbers or 0x7fff (0.499999) to
     // even numbers.
     e += 0x7fff + ((e >> 16) & 1);
     e = (e >> 16);
-    e = cast(UInt(16, e.type().lanes()), e);
-    e = reinterpret(BFloat(16, e.type().lanes()), e);
+    e = cast(UInt(16, lanes), e);
+    e = reinterpret(BFloat(16, lanes), e);
+    return e;
+}
+
+Expr float64_to_bfloat16(Expr e) {
+    internal_assert(e.type().bits() == 64);
+    const int lanes = e.type().lanes();
+    e = strict_float(e);
+
+    // First round to float and record any gain of loss of magnitude
+    Expr f = cast(Float(32, lanes), e);
+    Expr err = abs(e) - abs(f);
+    e = reinterpret(UInt(32, lanes), f);
+    // As above, but break ties using err, if non-zero
+    e += 0x7fff + (((err >= 0) & ((e >> 16) & 1)) | (err > 0));
+    e = (e >> 16);
+    e = cast(UInt(16, lanes), e);
+    e = reinterpret(BFloat(16, lanes), e);
     return e;
 }
 
@@ -96,10 +115,11 @@ Expr float32_to_float16(Expr value) {
     // 0.5 if the integer part is odd, or 0.4999999 if the
     // integer part is even, then truncate.
     bits += (bits >> 13) & 1;
-    bits += 0xfff;
-    bits = bits >> 13;
+    bits += make_const(UInt(32), ((uint32_t)1 << (13 - 1)) - 1);
+    bits = cast(u16_t, bits >> 13);
+
     // Rebias the exponent
-    bits -= 0x1c000;
+    bits -= 0x4000;
     // Truncate the top bits of the exponent
     bits = bits & 0x7fff;
     bits = select(is_denorm, denorm_bits,
@@ -111,6 +131,55 @@ Expr float32_to_float16(Expr value) {
     return common_subexpression_elimination(reinterpret(f16_t, bits));
 }
 
+Expr float64_to_float16(Expr value) {
+    value = strict_float(value);
+
+    Type f64_t = Float(64, value.type().lanes());
+    Type f16_t = Float(16, value.type().lanes());
+    Type u64_t = UInt(64, value.type().lanes());
+    Type u16_t = UInt(16, value.type().lanes());
+
+    Expr bits = reinterpret(u64_t, value);
+
+    // Extract the sign bit
+    Expr sign = bits & make_const(u64_t, (uint64_t)(0x8000000000000000ULL));
+    bits = bits ^ sign;
+
+    // Test the endpoints
+    Expr is_denorm = (bits < make_const(u64_t, (uint64_t)(0x3f10000000000000ULL)));
+    Expr is_inf = (bits >= make_const(u64_t, (uint64_t)(0x40f0000000000000ULL)));
+    Expr is_nan = (bits > make_const(u64_t, (uint64_t)(0x7ff0000000000000ULL)));
+
+    // Denorms are linearly spaced, so we can handle them by scaling up the
+    // input as a float or double by 2^24 and using the existing int-conversion
+    // rounding instructions. We can scale up by adding 24 to the exponent.
+    Expr denorm_bits = cast(u16_t, strict_float(round(strict_float(reinterpret(f64_t, bits + make_const(u64_t, (uint64_t)(0x0180000000000000ULL)))))));
+    Expr inf_bits = make_const(u16_t, 0x7c00);
+    Expr nan_bits = make_const(u16_t, 0x7fff);
+
+    // We want to round to nearest even, so we add either 0.5 if after
+    // truncation the last bit would be 1, or 0.4999999 if after truncation the
+    // last bit would be zero, then truncate.
+    bits += (bits >> 42) & 1;
+    bits += make_const(UInt(64), ((uint64_t)1 << (42 - 1)) - 1);
+    bits = bits >> 42;
+
+    // We no longer need the high bits
+    bits = cast(u16_t, bits);
+
+    // Rebias the exponent
+    bits -= 0x4000;
+    // Truncate the top bits of the exponent
+    bits = bits & 0x7fff;
+    bits = select(is_denorm, denorm_bits,
+                  is_inf, inf_bits,
+                  is_nan, nan_bits,
+                  cast(u16_t, bits));
+    // Recover the sign bit
+    bits = bits | cast(u16_t, sign >> 48);
+    return common_subexpression_elimination(reinterpret(f16_t, bits));
+}
+
 namespace {
 
 const std::map<std::string, std::string> transcendental_remapping =
@@ -171,6 +240,7 @@ Expr lower_float16_cast(const Cast *op) {
     Type src = op->value.type();
     Type dst = op->type;
     Type f32 = Float(32, dst.lanes());
+    Type f64 = Float(64, dst.lanes());
     Expr val = op->value;
 
     if (src.is_bfloat()) {
@@ -183,10 +253,18 @@ Expr lower_float16_cast(const Cast *op) {
 
     if (dst.is_bfloat()) {
         internal_assert(dst.bits() == 16);
-        val = float32_to_bfloat16(cast(f32, val));
+        if (src.bits() > 32) {
+            val = float64_to_bfloat16(cast(f64, val));
+        } else {
+            val = float32_to_bfloat16(cast(f32, val));
+        }
     } else if (dst.is_float() && dst.bits() < 32) {
         internal_assert(dst.bits() == 16);
-        val = float32_to_float16(cast(f32, val));
+        if (src.bits() > 32) {
+            val = float64_to_float16(cast(f64, val));
+        } else {
+            val = float32_to_float16(cast(f32, val));
+        }
     }
 
     return cast(dst, val);

src/Float16.cpp

@@ -9,7 +9,10 @@ namespace Internal {
 
 // Conversion routines to and from float cribbed from Christian Rau's
 // half library (half.sourceforge.net)
-uint16_t float_to_float16(float value) {
+template<typename T>
+uint16_t float_to_float16(T value) {
+    static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
+                  "float_to_float16 only supports float and double types");
     // Start by copying over the sign bit
     uint16_t bits = std::signbit(value) << 15;
 
@@ -40,14 +43,14 @@ uint16_t float_to_float16(float value) {
 
     // We've normalized value as much as possible. Put the integer
     // portion of it into the mantissa.
-    float ival;
-    float frac = std::modf(value, &ival);
+    T ival;
+    T frac = std::modf(value, &ival);
     bits += (uint16_t)(std::abs((int)ival));
 
     // Now consider the fractional part. We round to nearest with ties
     // going to even.
     frac = std::abs(frac);
-    bits += (frac > 0.5f) | ((frac == 0.5f) & bits);
+    bits += (frac > T(0.5)) | ((frac == T(0.5)) & bits);
 
     return bits;
 }
@@ -341,6 +344,19 @@ uint16_t float_to_bfloat16(float f) {
     return ret >> 16;
 }
 
+uint16_t float_to_bfloat16(double f) {
+    // Coming from double is a little tricker. We first narrow to float and
+    // record if any magnitude was lost of gained in the process. If so we'll
+    // use that to break ties instead of testing whether or not truncation would
+    // return odd.
+    float f32 = (float)f;
+    const double err = std::abs(f) - (double)std::abs(f32);
+    uint32_t ret;
+    memcpy(&ret, &f32, sizeof(float));
+    ret += 0x7fff + (((err >= 0) & ((ret >> 16) & 1)) | (err > 0));
+    return ret >> 16;
+}
+
 float bfloat16_to_float(uint16_t b) {
     // Assume little-endian floats
     uint16_t bits[2] = {0, b};
@@ -362,7 +378,17 @@ float16_t::float16_t(double value)
 }
 
 float16_t::float16_t(int value)
-    : data(float_to_float16(value)) {
+    : data(float_to_float16((float)value)) {
+    // integers of any size that map to finite float16s are all representable as
+    // float, so we can go via the float conversion method.
+}
+
+float16_t::float16_t(int64_t value)
+    : data(float_to_float16((float)value)) {
+}
+
+float16_t::float16_t(uint64_t value)
+    : data(float_to_float16((float)value)) {
 }
 
 float16_t::operator float() const {
@@ -464,7 +490,15 @@ bfloat16_t::bfloat16_t(double value)
 }
 
 bfloat16_t::bfloat16_t(int value)
-    : data(float_to_bfloat16(value)) {
+    : data(float_to_bfloat16((double)value)) {
+}
+
+bfloat16_t::bfloat16_t(int64_t value)
+    : data(float_to_bfloat16((double)value)) {
+}
+
+bfloat16_t::bfloat16_t(uint64_t value)
+    : data(float_to_bfloat16((double)value)) {
 }
 
 bfloat16_t::operator float() const {

src/Float16.h

@@ -32,6 +32,8 @@ struct float16_t {
     explicit float16_t(float value);
     explicit float16_t(double value);
     explicit float16_t(int value);
+    explicit float16_t(int64_t value);
+    explicit float16_t(uint64_t value);
     // @}
 
     /** Construct a float16_t with the bits initialised to 0. This represents
@@ -175,6 +177,8 @@ struct bfloat16_t {
     explicit bfloat16_t(float value);
     explicit bfloat16_t(double value);
     explicit bfloat16_t(int value);
+    explicit bfloat16_t(int64_t value);
+    explicit bfloat16_t(uint64_t value);
     // @}
 
     /** Construct a bfloat16_t with the bits initialised to 0. This represents

src/IR.cpp

@@ -678,6 +678,7 @@ const char *const intrinsic_op_names[] = {
     "sliding_window_marker",
     "sorted_avg",
     "strict_add",
+    "strict_cast",
     "strict_div",
     "strict_eq",
     "strict_le",

src/IR.h

@@ -626,6 +626,7 @@ struct Call : public ExprNode<Call> {
         // them as reals and ignoring the existence of nan and inf. Using these
         // intrinsics instead prevents any such optimizations.
         strict_add,
+        strict_cast,
         strict_div,
         strict_eq,
         strict_le,
@@ -792,6 +793,7 @@ struct Call : public ExprNode<Call> {
     bool is_strict_float_intrinsic() const {
         return is_intrinsic(
             {Call::strict_add,
+             Call::strict_cast,
              Call::strict_div,
              Call::strict_max,
              Call::strict_min,

src/StrictifyFloat.cpp

@@ -83,6 +83,16 @@ class Strictify : public IRMutator {
             return IRMutator::visit(op);
         }
     }
+
+    Expr visit(const Cast *op) override {
+        if (op->value.type().is_float() &&
+            op->type.is_float()) {
+            return Call::make(op->type, Call::strict_cast,
+                              {mutate(op->value)}, Call::PureIntrinsic);
+        } else {
+            return IRMutator::visit(op);
+        }
+    }
 };
 
 const std::set<std::string> strict_externs = {
@@ -142,6 +152,8 @@ Expr unstrictify_float(const Call *op) {
         return op->args[0] <= op->args[1];
     } else if (op->is_intrinsic(Call::strict_eq)) {
         return op->args[0] == op->args[1];
+    } else if (op->is_intrinsic(Call::strict_cast)) {
+        return cast(op->type, op->args[0]);
     } else {
         internal_error << "Missing lowering of strict float intrinsic: "
                        << Expr(op) << "\n";