Fix SkipStages narrowing loaded by Select condition (#9152)

SkipStages::visit(Select) was combining each branch's `loaded` predicate
with the select condition (via make_select / make_and). That's wrong:
Halide's Select evaluates both branches and only picks one of the
results, so any load inside either branch fires unconditionally. The
`loaded` predicate must be the OR of both branches, ungated by the
condition.

The bug caused allocation bounds inference to size affected buffers down
to zero whenever the runtime condition was false, while the generated
code still emitted a vectorized load from them -- a heap OOB read that
showed up intermittently as a valgrind use-after-free on the
truncated_pyramid test.

Keep `used` gated by the condition as before (with the same select/and
collapse that fixed the exponential blow-up in #9147).

Add a regression test in skip_stages.cpp that records the minimum
producer allocation size through a custom malloc handler and verifies
it's non-zero when the producer's only consumer is inside a
runtime-false select branch.

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

Author: abadams

src/SkipStages.cpp

@@ -488,11 +488,6 @@ class SkipStages : public IRMutator {
         mutate(op->false_value);
         // func_info now holds the false-branch info.
 
-        // Ids touched on both branches: combine with select(cond, t, f),
-        // so make_select can collapse `select(c, X, X) -> X` when the
-        // two branches contributed the same Expr.
-        // Ids touched on only one branch: AND the predicate with the
-        // appropriate side of the condition.
         auto merge_into_old = [&](size_t id, const Expr &u, const Expr &l) {
             auto [q, inserted] = old.try_emplace(id, FuncInfo{u, l});
             if (!inserted) {
@@ -505,13 +500,16 @@ class SkipStages : public IRMutator {
             size_t id = p.first;
             auto it_t = true_info.find(id);
             Expr u, l;
+            // Select evaluates both sides but only uses one, so `loaded` is
+            // the unconditional OR of the two branches while `used` is gated
+            // by the condition.
             if (it_t != true_info.end()) {
                 u = make_select(op->condition, it_t->second.used, p.second.used);
-                l = make_select(op->condition, it_t->second.loaded, p.second.loaded);
+                l = make_or(it_t->second.loaded, p.second.loaded);
                 true_info.erase(it_t);
             } else {
                 u = make_and(p.second.used, !op->condition);
-                l = make_and(p.second.loaded, !op->condition);
+                l = p.second.loaded;
             }
             merge_into_old(id, u, l);
         }
@@ -521,7 +519,7 @@ class SkipStages : public IRMutator {
         for (auto &p : true_info) {
             size_t id = p.first;
             Expr u = make_and(p.second.used, op->condition);
-            Expr l = make_and(p.second.loaded, op->condition);
+            Expr l = p.second.loaded;
             merge_into_old(id, u, l);
         }
         func_info.clear();

test/correctness/skip_stages.cpp

@@ -10,6 +10,22 @@ extern "C" HALIDE_EXPORT_SYMBOL int call_counter(int x, int idx) {
 }
 HalideExtern_2(int, call_counter, int, int);
 
+size_t min_alloc_size = (size_t)-1;
+void *recording_malloc(JITUserContext *, size_t x) {
+    if (x < min_alloc_size) {
+        min_alloc_size = x;
+    }
+    // Over-allocate so we can return a 32-byte-aligned pointer; stash the
+    // original so we can free it later.
+    void *orig = malloc(x + 32);
+    void *ptr = (void *)((((size_t)orig + 32) >> 5) << 5);
+    ((void **)ptr)[-1] = orig;
+    return ptr;
+}
+void recording_free(JITUserContext *, void *p) {
+    free(((void **)p)[-1]);
+}
+
 void reset_counts() {
     for (int i = 0; i < 4; i++) {
         call_count[i] = 0;
@@ -245,7 +261,7 @@ int main(int argc, char **argv) {
     }
 
     {
-        // Check the iteration with storage hoisting
+        // Check the interaction with storage hoisting
 
         // This Func may or may not be loaded, depending on y
         Func maybe_loaded("maybe_loaded");
@@ -268,6 +284,45 @@ int main(int argc, char **argv) {
         output.realize({100, 100});
     }
 
+    {
+        // Regression test: a Func loaded only inside a select branch whose
+        // condition is false at runtime should still get a non-zero
+        // allocation, because Halide's Select evaluates both branches
+        // (the load fires regardless of the condition). The producer's
+        // compute body can still be skipped because `used` is correctly
+        // gated.
+        Func producer("producer");
+        producer(x) = call_counter(x, 0);
+
+        Func consumer("consumer");
+        consumer(x) = select(toggle1, producer(x) + producer(x + 1), 42);
+
+        producer.compute_root().vectorize(x, 8);
+        consumer.vectorize(x, 8);
+
+        consumer.jit_handlers().custom_malloc = recording_malloc;
+        consumer.jit_handlers().custom_free = recording_free;
+
+        reset_counts();
+        toggle1.set(false);
+        min_alloc_size = (size_t)-1;
+        Buffer<int> out = consumer.realize({64});
+
+        // Producer must be skipped when toggle is false.
+        check_counts(0);
+        // ...but the allocation must still happen, sized for the load.
+        if (min_alloc_size == 0) {
+            printf("Producer allocation was zero-sized; unconditional load would OOB\n");
+            exit(1);
+        }
+        for (int i = 0; i < 64; i++) {
+            if (out(i) != 42) {
+                printf("out(%d) = %d, expected 42\n", i, out(i));
+                exit(1);
+            }
+        }
+    }
+
     printf("Success!\n");
     return 0;
 }