Promote narrow ints to floats in widening_mul (#8846)

Author: alexreinking

src/IROperator.cpp

@@ -1162,6 +1162,20 @@ Expr widening_add(Expr a, Expr b) {
 
 Expr widening_mul(Expr a, Expr b) {
     user_assert(a.defined() && b.defined()) << "widening_mul of undefined Expr\n";
+    // Promote float to int if lossless
+    if (a.type().is_float() && b.type().is_int_or_uint()) {
+        Expr float_b = lossless_cast(a.type(), b);
+        user_assert(float_b.defined())
+            << "widening_mul: cannot promote RHS of type " << b.type() << " to " << a.type() << ".\n"
+            << "Please use an explicit cast.\n";
+        b = float_b;
+    } else if (b.type().is_float() && a.type().is_int_or_uint()) {
+        Expr float_a = lossless_cast(b.type(), a);
+        user_assert(float_a.defined())
+            << "widening_mul: cannot promote LHS of type " << a.type() << " to " << b.type() << ".\n"
+            << "Please use an explicit cast.\n";
+        a = float_a;
+    }
     // Widening multiplies can have different signs.
     match_bits(a, b);
     match_lanes(a, b);

src/Type.cpp

@@ -118,9 +118,10 @@ bool Type::can_represent(Type other) const {
         if (other.is_bfloat()) {
             return bits() > other.bits();
         } else {
-            return ((other.is_float() && other.bits() <= bits()) ||
-                    (bits() == 64 && other.bits() <= 32) ||
-                    (bits() == 32 && other.bits() <= 16));
+            return (other.is_float() && other.bits() <= bits()) ||
+                   (bits() == 64 && other.bits() <= 32) ||
+                   (bits() == 32 && other.bits() <= 16) ||
+                   (bits() == 16 && other.bits() <= 8);
         }
     } else {
         return false;

test/correctness/intrinsics.cpp

@@ -15,7 +15,7 @@ void check(Expr test, Expr expected, Type required_type) {
 
     // Make sure the pattern is robust to CSE. We only enforce this
     // for types with well-defined overflow for now.
-    if (test.type().bits() < 32 || test.type().is_uint()) {
+    if (test.type().can_overflow()) {
         auto bundle = [](const Expr &e) {
             return Call::make(e.type(), Call::bundle, {e, e}, Call::PureIntrinsic);
         };
@@ -131,7 +131,10 @@ int main(int argc, char **argv) {
     Expr u8y = make_leaf(UInt(8, 4), "u8y");
     Expr u8z = make_leaf(UInt(8, 4), "u8w");
     Expr u8w = make_leaf(UInt(8, 4), "u8z");
+    Expr i16x = make_leaf(Int(16, 4), "i16x");
     Expr u16x = make_leaf(UInt(16, 4), "u16x");
+    Expr i16y = make_leaf(Int(16, 4), "i16y");
+    Expr u16y = make_leaf(UInt(16, 4), "u16y");
     Expr u32x = make_leaf(UInt(32, 4), "u32x");
     Expr u32y = make_leaf(UInt(32, 4), "u32y");
     Expr i32x = make_leaf(Int(32, 4), "i32x");
@@ -170,6 +173,21 @@ int main(int argc, char **argv) {
     check(u32(u8x) * u8y, u32(widening_mul(u8x, u8y)));
     check(f32(f16x) * f32(f16y), widening_mul(f16x, f16y));
 
+    // Check mixed float/int promotion
+    check(widening_mul(i8x, f16y), widening_mul(f16(i8x), f16y));
+    check(widening_mul(i8x, f32y), widening_mul(f32(i8x), f32y));
+    check(widening_mul(u8x, f16y), widening_mul(f16(u8x), f16y));
+    check(widening_mul(u8x, f32y), widening_mul(f32(u8x), f32y));
+    check(widening_mul(i16x, f32y), widening_mul(f32(i16x), f32y));
+    check(widening_mul(u16x, f32y), widening_mul(f32(u16x), f32y));
+
+    check(widening_mul(f16x, i8y), widening_mul(f16x, f16(i8y)));
+    check(widening_mul(f32x, i8y), widening_mul(f32x, f32(i8y)));
+    check(widening_mul(f16x, u8y), widening_mul(f16x, f16(u8y)));
+    check(widening_mul(f32x, u8y), widening_mul(f32x, f32(u8y)));
+    check(widening_mul(f32x, i16y), widening_mul(f32x, f32(i16y)));
+    check(widening_mul(f32x, u16y), widening_mul(f32x, f32(u16y)));
+
     // Widening mul allows mixed signs
     check(i16(i8x) * u8y, widening_mul(i8x, u8y), Int(16, 4));
     check(i32(i8x) * u8y, i32(widening_mul(i8x, u8y)));
@@ -377,6 +395,16 @@ int main(int argc, char **argv) {
         g.compile_jit();
     }
 
+    // Mixed integer-floating-point used to crash codegen
+    {
+        Func f, q, s;
+        q(x) = cast<int8_t>(0);
+        s(x) = cast(Float(16), 1.0f);
+        f(x) = widening_mul(q(x), s(x));
+
+        f.compile_jit();
+    }
+
     printf("Success!\n");
     return 0;
 }