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 80.22599% with 35 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/ScheduleFunctions.cpp	78.94%	13 Missing and 11 partials ⚠️
src/IROperator.cpp	60.86%	6 Missing and 3 partials ⚠️
src/CSE.cpp	83.33%	0 Missing and 1 partial ⚠️
src/DerivativeUtils.cpp	80.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #9194   +/-   ##
=======================================
  Coverage        ?   70.18%           
=======================================
  Files           ?      255           
  Lines           ?    79127           
  Branches        ?    18930           
=======================================
  Hits            ?    55538           
  Misses          ?    17933           
  Partials        ?     5656           

☔ 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)

[alexreinking]

@soufianekhiat -- yes, that's the gist of my suggestion. But I don't want to over-promise either. I don't know if that's the best workable solution for your problem because I don't know what your problem is.

[mcourteaux]

Have you tried to just hack in Halide::Internal::if_then_else() in your Func definition and see if it satisfies your situation?

[soufianekhiat]

Here my interpretation of the ideas. Move to FuncRef, force_inline each branch (error if we can't).
I encouter this limitation multiple times I just decided to create a PR this time. For instance on GPU it's totally ok to have wavefront coherent branches, but we never benefit of that optimization. Or on CPU we have most of the time good branch predictor which discarted most of the time, we have to do some mental gymnastic to force Halide to produce real branch. Here it make everything cleaner and explicit and if we can if produces an error to the users.

[alexreinking]

Why not do this in ScheduleFuncs? The implementation would look very similar as for specialize, just point-wise rather than func-wise.

[alexreinking]

I don't think this is a good idea. What if the condition is loop-invariant at a middle level. Shouldn't you be able to ultimately hoist the condition and compute_at something at a loop level inside the hoisted condition?

[soufianekhiat]

ScheduleFuncs: yes, need sometime for the refacto.

Force inline: convinced! i will remove it. It kills the scheduling.
My understanding is that case:

f(x, y) = branch(cond(y), a(x, y), b(x, y));   // cond depends only on y
cond(y) is invariant in the x loop. If the branch is real control flow, you can hoist it to the y level and compute_at a producer inside it:
for y:
  if (cond(y)):
    for x: <compute a's tile HERE>   f[x, y] = a(x, y)
  else:
    for x:                            f[x, y] = b(x, y)

[soufianekhiat]

Bump on a new version

[soufianekhiat]

Another usecase, the AutoDiff, I have a branch with VJP/JPV auto diff side by side. But exploded for a given Func, the SDF with lot of branches:

// ============================================================================
// Ellipse (iquilez sdEllipse)
// Analytic exact SDF using cubic solver.
// ============================================================================
Halide::Expr
compute_sdf_ellipse( Halide::Type t, Halide::Expr px, Halide::Expr py,
                     Halide::Expr semi_a, Halide::Expr semi_b )
{
	auto C = [t](double v) { return Halide::Internal::make_const(t, v); };
	semi_a = Halide::cast( t, semi_a );
	semi_b = Halide::cast( t, semi_b );
	Halide::Expr qx = Halide::abs( px );
	Halide::Expr qy = Halide::abs( py );
	// Swap so ab.x <= ab.y (swap if p.x > p.y)
	Halide::Expr cond_swap = qx > qy;
	Halide::Expr ax = Halide::select( cond_swap, qy, qx );
	Halide::Expr ay = Halide::select( cond_swap, qx, qy );
	Halide::Expr ab_x = Halide::select( cond_swap, semi_b, semi_a );
	Halide::Expr ab_y = Halide::select( cond_swap, semi_a, semi_b );
	Halide::Expr l = ab_y * ab_y - ab_x * ab_x;
	Halide::Expr safe_l = Halide::select( Halide::abs( l ) < C(1e-8), C(1e-8), l );
	Halide::Expr m_val = ab_x * ax / safe_l;
	Halide::Expr m2 = m_val * m_val;
	Halide::Expr n_val = ab_y * ay / safe_l;
	Halide::Expr n2 = n_val * n_val;
	Halide::Expr c = ( m2 + n2 - C(1) ) / C(3);
	Halide::Expr c3 = c * c * c;
	Halide::Expr q_val = c3 + m2 * n2 * C(2);
	Halide::Expr d_val = c3 + m2 * n2;
	Halide::Expr g = m_val + m_val * n2;
	// Case 1: d < 0 (three roots)
	Halide::Expr safe_c3 = Halide::select( Halide::abs( c3 ) < C(1e-10), C(-1e-10), c3 );
	Halide::Expr h1 = Halide::acos( Halide::clamp( q_val / safe_c3, C(-1), C(1) ) ) / C(3);
	Halide::Expr s1 = Halide::cos( h1 );
	Halide::Expr t1 = Halide::sin( h1 ) * C(1.73205080757);
	Halide::Expr rx1 = Halide::sqrt( Halide::max( -c * ( s1 + t1 + C(2) ) + m2, C(0) ) );
	Halide::Expr ry1 = Halide::sqrt( Halide::max( -c * ( s1 - t1 + C(2) ) + m2, C(0) ) );
	Halide::Expr co1 = ( ry1 + Halide::select( l > C(0), C(1), C(-1) ) * rx1
	                    + Halide::abs( g ) / Halide::max( rx1 * ry1, C(1e-8) )
	                    - m_val ) * C(0.5);
	// Case 2: d >= 0 (one root)
	Halide::Expr h2_inner = C(2) * m_val * n_val * Halide::sqrt( Halide::max( d_val, C(0) ) );
	Halide::Expr s2_x = q_val + h2_inner;
	Halide::Expr s2 = Halide::select( s2_x >= C(0), C(1), C(-1) )
	                * Halide::pow( Halide::abs( s2_x ), C(1.0 / 3.0) );
	Halide::Expr u2_x = q_val - h2_inner;
	Halide::Expr u2 = Halide::select( u2_x >= C(0), C(1), C(-1) )
	                * Halide::pow( Halide::abs( u2_x ), C(1.0 / 3.0) );
	Halide::Expr rx2 = -s2 - u2 - c * C(4) + C(2) * m2;
	Halide::Expr ry2 = ( s2 - u2 ) * C(1.73205080757);
	Halide::Expr rm2 = Halide::sqrt( Halide::max( rx2 * rx2 + ry2 * ry2, C(0) ) );
	Halide::Expr co2 = ( ry2 / Halide::max( Halide::sqrt( Halide::max( rm2 - rx2, C(0) ) ), C(1e-8) )
	                    + C(2) * g / Halide::max( rm2, C(1e-8) ) - m_val ) * C(0.5);
	// Select case
	Halide::Expr co = Halide::select( d_val < C(0), co1, co2 );
	co = Halide::clamp( co, C(0), C(1) );
	Halide::Expr r_x = ab_x * co;
	Halide::Expr r_y = ab_y * Halide::sqrt( Halide::max( C(1) - co * co, C(0) ) );
	return sdf_length2( r_x - ax, r_y - ay )
	     * Halide::select( ay - r_y > C(0), C(1), C(-1) );
}

[soufianekhiat]

Added support for VJP (propagate_adjoing) don't help for explosion but add it for correctness.
sidenote:
For JVP propagation you can check my production/working branch sk/hlsl6_jvp_sioutas_trace (include hlsl6, JVP, sioutas auto scheduler and HalideTraceViz working on windows/linux/macos (based on ffmpeg) and with branch working for VJP and JVP).