Add 'branch(cond, a, b)' to have a materialized branch

[soufianekhiat]

Context

Today when we write select(cond, a, b) in a Halide kernel, both side a
and b are always evaluated, and then we pick one. This is fine most of the
time but when one side is really expensive it is a waste, we compute things
for nothing.

The compiler already have internally what we need for this: the if_then_else
intrinsic. Unlike Select, it "promise to evaluate exactly one of the
conditions" and it already lower to a real control flow branch (basic blocks +
conditional branch + phi in LLVM, real if/else in the C backend, blocks + phi
in Vulkan). The only thing missing was a way for the user to reach it.

What this PR do

Add a new method Expr::branch(). You call it on the result of a select and
it give you back the same thing but as an explicit branch that evaluate only
the taken side:

f(x) = select(cond, cheap, very_expensive(x)).branch();

It is just a thin rewrite from the Select node to the if_then_else
intrinsic, so all the existing machinery (lowering, simplification, bounds,
cost model, serialization, every backend) already handle it.

It also work with the multi way select, the full chain become a full
if / else-if / else chain of branches:

select(x < 10, 1, x < 20, 2, 3).branch();

CSE barrier

There was one real problem. The CSE pass don't know that the two arms of
if_then_else are lazy. So a subexpression used only inside one arm could be
lifted into a let above the branch and then be evaluated all the time, which
kill the whole point (and could even run an operation the branch was there to
protect).

So this PR also teach ComputeUseCounts in CSE.cpp to not count the uses
inside the value arms of if_then_else. The condition is still counted
normally because it is always evaluated. Like this an arm only subexpression is
never hoisted out of the branch. This stay conservative for the two callers
that use lift_all, it only make them a bit more careful for the arms. The
within arm duplicates are anyway cleaned later by LLVM.

Limitations (documented in the header)

• A per lane branch can not exist under vectorization or on GPU, so in those
  context .branch() degrade back to a select (both sides evaluated). This is
  expected, not a bug.
• .branch() only rewrite selects that are directly a select or directly the
  arm of a select. A select buried inside some arithmetic in an arm stay a
  normal data select.
• It is an error to call .branch() on something that is not a select.

Tests

New test test/correctness/branch.cpp (added to the correctness group):

• the frontend produce the if_then_else intrinsic and keep the operands and
  the type
• .branch() give the exact same result than select() on a realized Func
• the multi way select recurse into the tail
• calling .branch() on a non select throw
• the CSE barrier: an expensive extern call used twice in an arm stay inside
  the branch and is not hoisted out (checked on the lowered IR)

I also run mux, multi_way_select, likely and tuple_select to check the CSE
change do not break the select / if_then_else lowering, they all pass.

Files

• src/Expr.h: declare Expr::branch() with the doc
• src/IROperator.cpp: the recursive rewrite select -> if_then_else
• src/CSE.cpp: the barrier so arm subexpressions are not hoisted
• test/correctness/branch.cpp + test/correctness/CMakeLists.txt: the test

Disclosure: Co-author Claude Opus 4.8

[abadams]

Pervasive throughout the compiler is the implicit assumption that within a single expr there are no sequence points or side-effects. It's more than just CSE that this is likely to break. The current way to get a branch from a select is to make the true/false values of the select calls to a Func, and then schedule those Funcs somewhere such that the condition is a constant for the lifetime of the func. This will give you an actual branch around the computation of the Func.

[codecov]

Codecov Report

❌ Patch coverage is 79.36508% with 13 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@4e39fc9). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/IROperator.cpp	66.66%	5 Missing and 3 partials ⚠️
src/Lower.cpp	87.09%	1 Missing and 3 partials ⚠️
src/CSE.cpp	83.33%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #9194   +/-   ##
=======================================
  Coverage        ?   69.40%           
=======================================
  Files           ?      254           
  Lines           ?    78338           
  Branches        ?    18749           
=======================================
  Hits            ?    54370           
  Misses          ?    18460           
  Partials        ?     5508           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[mcourteaux]

I think exposing a branched variant of Select to users is fine, but I have doubts:

• It is an error to call .branch() on something that is not a select.

  It's weird to have a .branch() modifier. I'd prefer to see a .select_branched(...) or a .mux_branched(), or something similar.

• A per lane branch can not exist under vectorization or on GPU, so in those
  context .branch() degrade back to a select (both sides evaluated). This is
  expected, not a bug.

  This should user_assert instead IMO. If your schedule can't be fulfilled, that's something the user should know. For me having an explicitly branched version of the select is the guarantee that it doesn't evaluate both sides: this is not only useful for performance, but more importantly useful for more aggressive bounds pruning during bounds inference. If that guarantee is randomly not fulfilled, you don't gain anything from this, because then it'd be almost always better to do not compute values you don't need.

[soufianekhiat]

• it is now a free function branch(cond, a, b), same signature as select. No
  more .branch() modifier, make more sense.
• the condition must be scalar. If you pass a vector condition it is a hard
  error.
• when it can not become a real branch (condition became lane varying after
  vectorize, or inside a gpu kernel) it is now a user_assert.

About @abadams point: you are right, the no side effect assumption is broad and i
do not try to fix that everywhere. I keep branch very limited (scalar only,
and it becomes a plain if_then_else before codegen) so the number of passes
that must understand it stays small.

I'll try to do something to have it on GPU too.

[alexreinking]

I have reservations about the tractability of this proposal, too. This fundamentally changes the contract on expressions. I think that, at a minimum, you'd need to test the hypothetical of "what breaks if every select() that exists today were replaced by branch()?". It's likely more than just vectorization / CSE.

I'd say my biggest semantic concern is that bounds inference and multi-stage funcs each restrict the usefulness of this feature. For example:

f(x) = ... branch(..., g(x), h(x)) ...;

Only the loads from g and h are protected by the branch (unless they are inlined). Bounds inference will force the full extents of g and h to be realized. If g or h has an update definition, then they cannot be inlined and you would need to copy the branch predicate onto the RDom. Making this all "smarter" is probably a research project.

A lesser issue with the above: at a glance, the code looks like a C-like ternary (p ? g(x) : h(x)) whose semantics would gate the evaluation of g/h (not just the load). So the interaction with bounds inference makes the syntax potentially misleading/confusing. This is a problem with select today, too, but adding a second ternary-ish statement would make it much easier to assume that "oh one is eager and the other is lazy" or some other such inaccurate mental models. Again, this is a lesser issue, but I'm just thinking through the ergonomics.

I'm curious: what is the shape of the computation you're trying to branch? Maybe there are smaller features we could add that would make it easier to express. I could see adding predicates to update stages without a dummy RDom, or maybe even to pure definitions (I haven't thought all those consequences through).

[alexreinking]

@soufianekhiat — would it be sufficient to allow branching only as the outermost "expression"? In that case, it could become a part of the Stage's definition and branch() would return something like a ValueAlternatives type (instead of an Expr) that could be passed to FuncRef::operator=.

[soufianekhiat]

Let be sure I understand.
If we force a and b in branch(cond, a, b) are forced inlined it would help the boundary inferences.

The ideas ValueAlternatives would be to move everything in FuncRef to avoid:

f(x) = g(x) + branch(cond, a, b) * 2;   // branch buried in the middle -> contract problem

To:

f(x) = branch(cond, a, b);              // OK - branch is the ENTIRE right-hand side
f(x) = g(x) + branch(cond, a, b);       // COMPILE ERROR - can't add Expr + ValueAlternatives

And for the alternative:

f(x) = undef<int>();          // no default work
RDom ra(0, N); ra.where(cond(ra));
f(ra) = expensive_a(ra);      // -> for r: if(cond)  f[r] = expensive_a(r)   (a only runs where cond)
RDom rb(0, N); rb.where(!cond(rb));
f(rb) = expensive_b(rb);      // -> for r: if(!cond) f[r] = expensive_b(r)

I don't like the RDom make everything two stage vs branch which had simpler "expressiveness".

Solution:

1. Make everything in branch force inline
2. Implement ValueAlternatives (or alternative name)