Handle nested vectorization in store predicates

src/IR.cpp

@@ -574,6 +574,8 @@ Stmt IfThenElse::make(Expr condition, Stmt then_case, Stmt else_case) {
     internal_assert(condition.defined() && then_case.defined()) << "IfThenElse of undefined\n";
     // else_case may be null.
 
+    internal_assert(condition.type().is_scalar()) << "IfThenElse with vector condition\n";
+
     IfThenElse *node = new IfThenElse;
     node->condition = std::move(condition);
     node->then_case = std::move(then_case);

src/PartitionLoops.cpp

@@ -263,7 +263,7 @@ class FindSimplifications : public IRVisitor {
         }
         condition = remove_likelies(condition);
         Simplification s = {condition, std::move(old), std::move(likely_val), std::move(unlikely_val), true};
-        if (s.condition.type().is_vector()) {
+        while (s.condition.type().is_vector()) {
             s.condition = simplify(s.condition);
             if (const Broadcast *b = s.condition.as<Broadcast>()) {
                 s.condition = b->value;

src/Simplify_Stmts.cpp

@@ -328,6 +328,10 @@ Stmt Simplify::visit(const Store *op) {
 
     const Load *load = value.as<Load>();
     const Broadcast *scalar_pred = predicate.as<Broadcast>();
+    if (scalar_pred && !scalar_pred->value.type().is_scalar()) {
+        // Nested vectorization
+        scalar_pred = nullptr;
+    }
 
     ModulusRemainder align = ModulusRemainder::intersect(op->alignment, base_info.alignment);
 

test/correctness/CMakeLists.txt

@@ -115,6 +115,7 @@ tests(GROUPS correctness
       fuse_gpu_threads.cpp
       fused_where_inner_extent_is_zero.cpp
       fuzz_float_stores.cpp
+      fuzz_schedule.cpp
       gameoflife.cpp
       gather.cpp
       gpu_allocation_cache.cpp

test/correctness/fuzz_schedule.cpp

@@ -0,0 +1,60 @@
+#include "Halide.h"
+
+using namespace Halide;
+
+void check_blur_output(const Buffer<int> &out, const Buffer<int> &correct) {
+    for (int y = 0; y < out.height(); y++) {
+        for (int x = 0; x < out.width(); x++) {
+            if (out(x, y) != correct(x, y)) {
+                printf("out(%d, %d) = %d instead of %d\n",
+                       x, y, out(x, y), correct(x, y));
+                exit(1);
+            }
+        }
+    }
+}
+
+int main(int argc, char **argv) {
+    // This test is for schedules that crash the compiler found via fuzzing that
+    // are hard to otherwise reproduce. We don't need to check the output.
+
+    Buffer<int> correct;
+    {
+        // An unscheduled instance to act as a reference output
+        Func input("input");
+        Func local_sum("local_sum");
+        Func blurry("blurry");
+        Var x("x"), y("y");
+        input(x, y) = 2 * x + 5 * y;
+        RDom r(-2, 5, -2, 5);
+        local_sum(x, y) = 0;
+        local_sum(x, y) += input(x + r.x, y + r.y);
+        blurry(x, y) = cast<int32_t>(local_sum(x, y) / 25);
+        correct = blurry.realize({32, 32});
+    }
+
+    // https://github.com/halide/Halide/issues/7851
+    {
+        Func input("input");
+        Func local_sum("local_sum");
+        Func blurry("blurry");
+        Var x("x"), y("y");
+        input(x, y) = 2 * x + 5 * y;
+        RDom r(-2, 5, -2, 5);
+        local_sum(x, y) = 0;
+        local_sum(x, y) += input(x + r.x, y + r.y);
+        blurry(x, y) = cast<int32_t>(local_sum(x, y) / 25);
+        Var yo("yo"), yi("yi"), xo("xo"), xi("xi"), yo_x_f("yo_x_f"), yo_x_fo("yo_x_fo"), yo_x_fi("yo_x_fi");
+        blurry.split(y, yo, yi, 2, TailStrategy::RoundUp).fuse(yo, x, yo_x_f).vectorize(yi).split(yo_x_f, yo_x_fo, yo_x_fi, 2, TailStrategy::Predicate).reorder(yo_x_fo, yo_x_fi, yi);
+        input.split(y, yo, yi, 2, TailStrategy::PredicateStores).fuse(yo, x, yo_x_f).vectorize(yi).split(yo_x_f, yo_x_fo, yo_x_fi, 2, TailStrategy::Predicate).reorder(yo_x_fo, yo_x_fi, yi);
+        blurry.store_root();
+        input.compute_at(blurry, yi);
+        Pipeline p({blurry});
+        Buffer<int> buf = p.realize({32, 32});
+        check_blur_output(buf, correct);
+    }
+
+    printf("Success!\n");
+
+    return 0;
+}