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.04192% 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	77.77%	13 Missing and 9 partials ⚠️
src/IROperator.cpp	60.00%	7 Missing and 3 partials ⚠️
src/CSE.cpp	83.33%	0 Missing and 1 partial ⚠️
src/DerivativeUtils.cpp	80.00%	0 Missing and 1 partial ⚠️
src/VectorizeLoops.cpp	66.66%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #9194   +/-   ##
=======================================
  Coverage        ?   70.24%           
=======================================
  Files           ?      255           
  Lines           ?    79072           
  Branches        ?    18912           
=======================================
  Hits            ?    55543           
  Misses          ?    17877           
  Partials        ?     5652           

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

[abadams]

Possible usage of this feature is to get better codegen for boundary conditions on archs that support predicated loads, e.g. it would be cool if this: f(x) = branch(x >= 0 && x < width, input(x), 0) just compiled to a predicated vector load on avx512. You may need a f(x) = branch(... condition ..., input(promise_clamped(x, ...)), 0) to help out bounds inference.

[soufianekhiat]

Possible usage of this feature is to get better codegen for boundary conditions on archs that support predicated loads, e.g. it would be cool if this: f(x) = branch(x >= 0 && x < width, input(x), 0) just compiled to a predicated vector load on avx512. You may need a f(x) = branch(... condition ..., input(promise_clamped(x, ...)), 0) to help out bounds inference.

Any pointer where to look?

I would like to have CAS (Compare-And-Swap) too but don't know how to handle that.

[alexreinking]

Any pointer where to look?

src/BoundaryConditions.cpp

[soufianekhiat]

I did some experiments to see what branch() really generate when you vectorize. It depend a lot of the conditions, so here the cases I tested:

case	predicated load	predicated store	real if	scalarized
condition vary per lane	yes	yes	-	no
condition loop invariant (Param)	no	no	yes	no
condition on outer dim only	no	no	yes	no
condition with likely()	yes	yes	yes	no
impure call in one arm	no	no	yes	yes
inside atomic()	no	no	yes	yes

So predication is not the normal path, it is the fallback when the condition can not be hoisted out of the vector loop. The common case is line 2 and 3, the condition is hoisted and we get a real scalar branch, and this is the point of the feature.

For the boundary condition case of @abadams it work. Compiled for avx512_skylake the loop body is branch free:

vpcmpltud  %zmm1, %zmm3, %k2               ; mask = (x < 100)
vmovdqu32  (%rdx,%rsi,4), %zmm3 {%k2} {z}  ; predicated load
vmovdqu32  %zmm3, (%rax,%rsi,4) {%k2}      ; masked store, taken lanes
vmovdqu32  %zmm2, (%rax,%rsi,4) {%k1}      ; masked store of 0, other lanes

For comparison the select() version clamp then do a gather, 16 scalar loads with vpinsrd, because a clamped index is not a ramp anymore.

But it did not work at first. unsafe_promise_clamped is created as PureIntrinsic in IROperator.cpp, but Simplify_Call.cpp rebuild it with Call::Intrinsic, so is_pure() become false and PredicateLoadStore refuse to predicate and we fall back to a scalar loop. Funny because the promise_clamped is exactly what @abadams suggest to help bounds inference. I did not change the call type (look intentional, to avoid CSE lifting the promise out of its if), I whitelist the two promise intrinsics in PredicateLoadStore, same way Deinterleave.cpp already do it.

Two notes:

• the two last lines of the table scalarize silently. It stay correct (only one side take effect per lane) but it is a silent perf cliff, maybe we want a warning there.
• the PredicateLoadStore change is not specific to branch(), it help any vectorized if containing a promise. Maybe it should be a separate PR.