Fix match_bits on bfloat and drop strip_widening_cast integer guard

src/CodeGen_LLVM.cpp

@@ -179,6 +179,7 @@ CodeGen_LLVM::CodeGen_LLVM(const Target &t)
       wild_u64x_(Variable::make(UInt(64, 0), "*")),
       wild_f32x_(Variable::make(Float(32, 0), "*")),
       wild_f64x_(Variable::make(Float(64, 0), "*")),
+      wild_bf16x_(Variable::make(BFloat(16, 0), "*")),
 
       wild_u1_(Variable::make(UInt(1), "*")),
       wild_i8_(Variable::make(Int(8), "*")),

src/CodeGen_LLVM.h

@@ -252,7 +252,7 @@ class CodeGen_LLVM : public IRVisitor {
     // @{
     Expr wild_u1x_, wild_i8x_, wild_u8x_, wild_i16x_, wild_u16x_;
     Expr wild_i32x_, wild_u32x_, wild_i64x_, wild_u64x_;
-    Expr wild_f32x_, wild_f64x_;
+    Expr wild_f32x_, wild_f64x_, wild_bf16x_;
 
     // Wildcards for scalars.
     Expr wild_u1_, wild_i8_, wild_u8_, wild_i16_, wild_u16_;

src/CodeGen_X86.cpp

@@ -825,6 +825,7 @@ void CodeGen_X86::codegen_vector_reduce(const VectorReduce *op, const Expr &init
         {VectorReduce::Add, 2, i32(widening_mul(wild_i16x_, wild_i16x_)), "dot_product", Int(16)},
         {VectorReduce::SaturatingAdd, 2, i32(widening_mul(wild_i16x_, wild_i16x_)), "saturating_dot_product", {}, Pattern::CombineInit},
 
+        {VectorReduce::Add, 2, widening_mul(wild_bf16x_, wild_bf16x_), "dot_product", {}, Pattern::CombineInit},
         {VectorReduce::Add, 2, wild_f32x_ * wild_f32x_, "dot_product", BFloat(16), Pattern::CombineInit},
 
         // Horizontal widening addition using a dot_product against a vector of ones.

src/FindIntrinsics.cpp

@@ -114,27 +114,21 @@ class FindIntrinsics : public IRMutator {
             return e.type();
         }(x);
 
-        Expr best;
+        Expr best = x;
+
         auto consider = [&](Type t) {
-            if (!best.defined() || t.bits() < best.type().bits()) {
+            if (t.bits() < best.type().bits()) {
                 Expr e = lossless_cast(t, x);
                 best = e.defined() ? e : best;
             }
         };
 
-        // Only strip a genuine widening cast, i.e. one whose inner type is no
-        // wider than x. If the outer cast is a narrowing cast, inner_t is wider
-        // than x, and considering it would return an expression wider than x,
-        // which can send the surrounding narrowing rewrites into infinite
-        // recursion.
-        if (inner_t.is_int_or_uint() && inner_t != x.type() &&
-            inner_t.bits() <= x.type().bits()) {
-            consider(inner_t);
-        }
+        consider(inner_t);
         consider(x.type().narrow());
         consider(x.type().narrow().with_code(halide_type_uint));
 
-        return best;
+        // Only report success if we actually stripped a cast.
+        return best.same_as(x) ? Expr() : best;
     }
 
     Expr to_rounding_shift(const Call *c) {

src/IROperator.cpp

@@ -675,11 +675,20 @@ void match_lanes(Expr &a, Expr &b) {
 // Cast to the wider type of the two. Already guaranteed to leave
 // signed/unsigned on number of lanes unchanged.
 void match_bits(Expr &x, Expr &y) {
+    // Widen the narrower operand to match the wider one's bit count. A bfloat
+    // has no wider counterpart (there is no bfloat32), so widening one yields a
+    // regular float, consistent with Type::widen().
+    auto widened = [](Type t, int bits) {
+        if (t.is_bfloat()) {
+            t = t.with_code(halide_type_float);
+        }
+        return t.with_bits(bits);
+    };
     // The signedness doesn't match, so just match the bits.
     if (x.type().bits() < y.type().bits()) {
-        x = cast(x.type().with_bits(y.type().bits()), x);
+        x = cast(widened(x.type(), y.type().bits()), x);
     } else if (y.type().bits() < x.type().bits()) {
-        y = cast(y.type().with_bits(x.type().bits()), y);
+        y = cast(widened(y.type(), x.type().bits()), y);
     }
 }
 }  // namespace

test/correctness/intrinsics.cpp

@@ -141,6 +141,8 @@ int main(int argc, char **argv) {
     Expr i32y = make_leaf(Int(32, 4), "i32y");
     Expr f16x = make_leaf(Float(16, 4), "f16x");
     Expr f16y = make_leaf(Float(16, 4), "f16y");
+    Expr bf16x = make_leaf(BFloat(16, 4), "bf16x");
+    Expr bf16y = make_leaf(BFloat(16, 4), "bf16y");
     Expr f32x = make_leaf(Float(32, 4), "f32x");
     Expr f32y = make_leaf(Float(32, 4), "f32y");
 
@@ -172,6 +174,12 @@ int main(int argc, char **argv) {
     check(i32(i8x) * i8y, i32(widening_mul(i8x, i8y)));
     check(u32(u8x) * u8y, u32(widening_mul(u8x, u8y)));
     check(f32(f16x) * f32(f16y), widening_mul(f16x, f16y));
+    // bfloat16 widens to float32 (Type::widen special-cases it), so
+    // f32(bf16) * f32(bf16) is a genuine widening multiply, just like the
+    // float16 case above. This exercises match_bits with a bfloat operand:
+    // it must widen bfloat16 to a regular float, never to a (non-codegen-able)
+    // bfloat32. See test/fuzz/lossless_cast.cpp for the width-mismatched case.
+    check(f32(bf16x) * f32(bf16y), widening_mul(bf16x, bf16y));
     check(f32(i8x) * f32(i8y), f32(widening_mul(i8x, i8y)));
     check(f32(u8x) * f32(u8y), f32(widening_mul(u8x, u8y)));
 
@@ -412,6 +420,19 @@ int main(int argc, char **argv) {
         f.compile_jit();
     }
 
+    // A widening op with a bfloat16 operand narrower than the other operand
+    // used to crash codegen: match_bits widened the bfloat16 with with_bits,
+    // producing a bfloat32 (there is no such type). It must widen to a regular
+    // float32 instead.
+    {
+        Func f, q, s;
+        q(x) = cast(BFloat(16), 1.0f);
+        s(x) = cast(Float(32), 1.0f);
+        f(x) = widening_mul(q(x), s(x));
+
+        f.compile_jit();
+    }
+
     printf("Success!\n");
     return 0;
 }

test/fuzz/lossless_cast.cpp

@@ -9,17 +9,34 @@ namespace {
 constexpr int size = 1024;
 Buffer<uint8_t> buf_u8(size, "buf_u8");
 Buffer<int8_t> buf_i8(size, "buf_i8");
+Buffer<float> buf_f32(size, "buf_f32");
 Var x{"x"};
 
 Expr random_expr(FuzzingContext &fuzz) {
+    // Decide up front whether to build an integer or a floating-point
+    // expression. The two must not mix: casting an out-of-range float to an
+    // integer is undefined behavior, which would break the integer bounds and
+    // overflow checks below. Integer expressions are the main event; floats
+    // are a minority that exist only to smoke-test the float paths of
+    // find_intrinsics (in particular strip_widening_cast on bfloat16).
+    const bool make_float = fuzz.ConsumeIntegralInRange(0, 3) == 0;
+
     std::vector<Expr> exprs;
-    // Add some atoms
-    exprs.push_back(cast<uint8_t>(fuzz.ConsumeIntegral<uint8_t>()));
-    exprs.push_back(cast<int8_t>(fuzz.ConsumeIntegral<int8_t>()));
-    exprs.push_back(cast<uint8_t>(fuzz.ConsumeIntegral<uint8_t>()));
-    exprs.push_back(cast<int8_t>(fuzz.ConsumeIntegral<int8_t>()));
-    exprs.push_back(buf_u8(x));
-    exprs.push_back(buf_i8(x));
+    if (make_float) {
+        exprs.push_back(buf_f32(x));
+        // The narrow floats are only ever used as the source of a widening
+        // cast to float32 (see below).
+        exprs.push_back(cast(BFloat(16), buf_f32(x)));
+        exprs.push_back(cast(Float(16), buf_f32(x)));
+    } else {
+        // Add some atoms
+        exprs.push_back(cast<uint8_t>(fuzz.ConsumeIntegral<uint8_t>()));
+        exprs.push_back(cast<int8_t>(fuzz.ConsumeIntegral<int8_t>()));
+        exprs.push_back(cast<uint8_t>(fuzz.ConsumeIntegral<uint8_t>()));
+        exprs.push_back(cast<int8_t>(fuzz.ConsumeIntegral<int8_t>()));
+        exprs.push_back(buf_u8(x));
+        exprs.push_back(buf_i8(x));
+    }
 
     // Make random combinations of them
     while (true) {
@@ -30,6 +47,48 @@ Expr random_expr(FuzzingContext &fuzz) {
         Expr e2_narrow = e2;
         e2 = cast(e1.type(), e2);
 
+        if (make_float) {
+            // Arithmetic on float16/bfloat16 isn't widely supported, so keep
+            // narrow floats as widening-cast sources only and do plain
+            // arithmetic in float32/float64.
+            if (e1.type().bits() < 32) {
+                // e.g. float32(bfloat16(...)) - the pattern that made
+                // strip_widening_cast peel back to a bfloat and (via match_bits
+                // against a wider operand) synthesize an uncodegen-able
+                // bfloat32.
+                e = cast(Float(32), e1);
+            } else {
+                switch (fuzz.ConsumeIntegralInRange(0, 4)) {
+                case 0:
+                    e = e1 + e2;
+                    break;
+                case 1:
+                    e = e1 - e2;
+                    break;
+                case 2:
+                    e = e1 * e2;
+                    break;
+                case 3:
+                    e = e1 / e2;
+                    break;
+                case 4:
+                    if (e1.type().bits() < 64) {
+                        e = cast(e1.type().widen(), e1);
+                    }
+                    break;
+                }
+            }
+
+            if (!e.defined()) {
+                continue;
+            }
+            if (e.type().bits() == 64 && (e.as<Cast>() == nullptr || fuzz.ConsumeIntegralInRange(0, 7) == 0)) {
+                return e;
+            }
+            exprs.push_back(e);
+            continue;
+        }
+
         Expr e3 = cast(e1.type().with_code(halide_type_uint), fuzz.PickValueInVector(exprs));
         bool may_widen = e1.type().bits() < 64;
         bool may_widen_right = e2_narrow.type() == e1.type().narrow();
@@ -266,6 +325,7 @@ bool might_have_ub(Expr e) {
 FUZZ_TEST(lossless_cast, FuzzingContext &fuzz) {
     buf_u8.fill(fuzz);
     buf_i8.fill(fuzz);
+    buf_f32.fill(fuzz);
 
     Expr e1 = random_expr(fuzz);
     Expr simplified = simplify(e1);
@@ -276,6 +336,21 @@ FUZZ_TEST(lossless_cast, FuzzingContext &fuzz) {
         return 0;
     }
 
+    if (e1.type().is_float()) {
+        // We don't have an integer lossless_cast to compare against for float
+        // expressions. They exist purely to smoke-test that find_intrinsics
+        // and the rest of lowering can handle them without crashing or
+        // producing invalid IR - in particular, strip_widening_cast peeling a
+        // float32(bfloat16) back to a bfloat and forming a widening_mul must
+        // stay well-typed (bfloat16 widens to float32, not bfloat32).
+        // Vectorizing is required for find_intrinsics to run on the multiply.
+        Func f;
+        f(x) = e1;
+        f.vectorize(x, 4, TailStrategy::RoundUp);
+        f.realize({size});
+        return 0;
+    }
+
     // We're also going to test constant_integer_bounds here.
     ConstantInterval bounds = constant_integer_bounds(e1);