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

src/Bounds.cpp

@@ -1179,7 +1179,7 @@ class Bounds : public IRVisitor {
                 return abs(op->args[0] - op->args[1]);
             } else if (op->is_intrinsic(Call::return_second)) {
                 return op->args[1];
-            } else if (op->is_intrinsic(Call::if_then_else)) {
+            } else if (op->is_intrinsic(Call::if_then_else) || op->is_intrinsic(Call::branch)) {
                 // Probably more conservative than necessary
                 return op->args[1];
             } else if (op->is_intrinsic(Call::rounding_shift_right)) {
@@ -1230,7 +1230,7 @@ class Bounds : public IRVisitor {
                 }
             }
         } else if (op->args.size() == 3) {
-            if (op->is_intrinsic(Call::if_then_else)) {
+            if (op->is_intrinsic(Call::if_then_else) || op->is_intrinsic(Call::branch)) {
                 // Probably more conservative than necessary
                 return Select::make(op->args[0], op->args[1], op->args[2]);
             } else if (can_widen_all(op->args)) {

src/CSE.cpp

@@ -169,6 +169,22 @@ class ComputeUseCounts : public IRGraphVisitor {
         // Visit the children if we haven't been here before.
         IRGraphVisitor::include(e);
     }
+
+    void visit(const Call *op) override {
+        // The value arguments of if_then_else are evaluated lazily: only the
+        // taken side actually runs. Don't count uses inside them, so that a
+        // subexpression that occurs only within an arm is never lifted into a
+        // let outside the branch. Doing so would evaluate it unconditionally,
+        // which defeats the point of the branch and could execute an
+        // operation the branch was there to guard against. The condition is
+        // always evaluated, so recurse into it normally.
+        if ((op->is_intrinsic(Call::if_then_else) || op->is_intrinsic(Call::branch)) &&
+            op->args.size() >= 2) {
+            include(op->args[0]);
+        } else {
+            IRGraphVisitor::visit(op);
+        }
+    }
 };
 
 /** Rebuild an expression using a map of replacements. Works on graphs without exploding. */

src/Func.cpp

@@ -24,6 +24,7 @@
 #include "IRMutator.h"
 #include "IROperator.h"
 #include "IRPrinter.h"
+#include "IRVisitor.h"
 #include "ImageParam.h"
 #include "LLVM_Output.h"
 #include "Lower.h"
@@ -3204,6 +3205,16 @@ Stage FuncRef::operator=(const Expr &e) {
     return (*this) = Tuple(e);
 }
 
+Stage FuncRef::operator=(const Branch &b) {
+    // The branch is kept as an intrinsic in the definition's value. It is turned
+    // into real control flow (a point-wise IfThenElse, hoisted to the loop level
+    // where its condition is invariant) in ScheduleFunctions, so that producers
+    // can be compute_at'd inside the branch and have their computation gated by
+    // it. We deliberately do NOT force-inline the arms: that would take away the
+    // user's ability to schedule them.
+    return (*this) = b.expr;
+}
+
 Stage FuncRef::operator=(const Tuple &e) {
     if (!func.has_pure_definition()) {
         for (size_t i = 0; i < args.size(); ++i) {

src/Func.h

@@ -58,6 +58,7 @@ struct VarOrRVar {
 };
 
 class ImageParam;
+struct Branch;
 
 namespace Internal {
 struct AssociativeOp;
@@ -533,6 +534,11 @@ class FuncRef {
      * for a Func with multiple outputs. */
     Stage operator=(const Tuple &);
 
+    /** Define this Func as an explicit branch (see \ref branch). The value arms
+     * are force-inlined so the branch gates real computation, not just a load;
+     * it is an error if a value arm calls a Func that can not be inlined. */
+    Stage operator=(const Branch &);
+
     /** Define a stage that adds the given expression to this Func. If the
      * expression refers to some RDom, this performs a sum reduction of the
      * expression over the domain. If the function does not already have a

src/IR.cpp

@@ -618,6 +618,7 @@ constexpr const char *intrinsic_op_names[] = {
     "bitwise_or",
     "bitwise_xor",
     "bool_to_mask",
+    "branch",
     "bundle",
     "call_cached_indirect_function",
     "cast_mask",

src/IR.h

@@ -620,6 +620,12 @@ struct Call : public ExprNode<Call> {
         // Converts a boolean to a mask. Scalar bools become -1 (all bits set) when true,
         // 0 when false. Vector bools are converted to proper vector masks.
         bool_to_mask,
+        // A user-facing strict select that lowers to a real control-flow
+        // branch and evaluates only the taken side (see branch() in
+        // IROperator.h). It is lowered to if_then_else late in lowering, or
+        // becomes an error if it can not be a real branch (a lane-varying
+        // condition under vectorization, or vectorization inside a GPU kernel).
+        branch,
         // Bundle multiple exprs together temporarily for analysis (e.g. CSE)
         bundle,
         // Takes a sequence of (condition, function) pairs, and calls the first

src/IROperator.cpp

@@ -1537,6 +1537,41 @@ Expr select(Expr condition, Expr true_value, Expr false_value) {
     return Select::make(std::move(condition), std::move(true_value), std::move(false_value));
 }
 
+Branch branch(Expr condition, Expr true_value, Expr false_value) {
+    user_assert(condition.defined()) << "branch() of undefined condition.\n";
+    user_assert(condition.type().is_bool())
+        << "The first argument to branch must be a boolean:\n"
+        << "  " << condition << " has type " << condition.type() << "\n";
+    user_assert(!condition.type().is_vector())
+        << "branch() requires a scalar condition, but its condition\n"
+        << "  " << condition << "\nhas type " << condition.type() << ".\n"
+        << "branch() generates a real control-flow branch that evaluates only "
+        << "one side, which cannot be done per lane. Use select() for a "
+        << "lane-varying condition.\n";
+
+    // Coerce int literals to the type of the other argument, like select().
+    if (as_const_int(true_value)) {
+        true_value = cast(false_value.type(), std::move(true_value));
+    }
+    if (as_const_int(false_value)) {
+        false_value = cast(true_value.type(), std::move(false_value));
+    }
+
+    user_assert(true_value.type() == false_value.type())
+        << "The second and third arguments to branch do not have a matching type:\n"
+        << "  " << true_value << " has type " << true_value.type() << "\n"
+        << "  " << false_value << " has type " << false_value.type() << "\n";
+
+    // Capture the type before moving the args (argument evaluation order is
+    // unspecified, so reading true_value.type() inline with std::move is UB).
+    Type result_type = true_value.type();
+    Expr e = Call::make(result_type,
+                        Call::branch,
+                        {std::move(condition), std::move(true_value), std::move(false_value)},
+                        Call::PureIntrinsic);
+    return Branch{std::move(e)};
+}
+
 Tuple select(const Tuple &condition, const Tuple &true_value, const Tuple &false_value) {
     user_assert(condition.size() == true_value.size() && true_value.size() == false_value.size())
         << "select() on Tuples requires all Tuples to have identical sizes.";

src/IROperator.h

@@ -851,6 +851,48 @@ inline Expr select(const Expr &c0, const FuncRef &v0, const Expr &c1, const Func
 }
 // @}
 
+/** The result of branch(). It is deliberately NOT an Expr: a branch may only
+ * appear as the entire right-hand side of a Func definition
+ * (f(x) = branch(cond, a, b)), never as a sub-expression. That single
+ * restriction is what makes the "evaluate only one side" guarantee tractable
+ * (bounds inference and CSE never see a lazy sub-expression), and it lets the
+ * value arms be force-inlined so the branch gates real computation, not just a
+ * load. */
+struct Branch {
+    /** Internal representation (an if_then_else-style intrinsic). Users never
+     * touch this directly; assign the Branch to a Func instead. */
+    Expr expr;
+};
+
+/** An explicitly branched variant of select. branch(cond, a, b) means the same
+ * thing as select(cond, a, b), but is guaranteed to lower to a real
+ * control-flow branch that evaluates only the taken side, instead of the
+ * branchless select that always evaluates both sides.
+ *
+ * Unlike select(), the result is a \ref Branch, not an Expr, so it can only be
+ * used as the whole right-hand side of a Func definition:
+ *     f(x) = branch(cond, a, b);
+ * The value arms are force-inlined so that only the taken side is actually
+ * computed (it is an error if a value arm calls a Func that can not be inlined,
+ * e.g. one with an update stage - inline it or use select() there instead).
+ *
+ * The condition must be a scalar bool. It works on CPU and on GPU, where it
+ * becomes a real per-thread or per-wave branch. Because a real branch can not
+ * be expressed per lane, it is a hard error (not a silent fallback to select)
+ * to use branch() with a lane-varying condition under vectorization, or to
+ * combine branch() with vectorization inside a GPU kernel. */
+Branch branch(Expr condition, Expr true_value, Expr false_value);
+
+/** A multi-way variant of branch, mirroring the multi-way select. Each
+ * condition becomes its own real branch, forming an if / else-if / else
+ * chain. */
+template<typename... Args,
+         std::enable_if_t<Internal::all_are_convertible<Expr, Args...>::value> * = nullptr>
+inline Branch branch(Expr c0, Expr v0, Expr c1, Expr v1, Args &&...args) {
+    return branch(std::move(c0), std::move(v0),
+                  branch(std::move(c1), std::move(v1), std::forward<Args>(args)...).expr);
+}
+
 /** Oftentimes we want to pack a list of expressions with the same type
  * into a channel dimension, e.g.,
  * img(x, y, c) = select(c == 0, 100, // Red

src/Lower.cpp

@@ -354,6 +354,7 @@ void lower_impl(const vector<Function> &output_funcs,
     s = simplify(s);
     log("Lowering after vectorizing:", s);
 
+
     if (t.has_gpu_feature() ||
         t.has_feature(Target::Vulkan)) {
         debug(1) << "Injecting per-block gpu synchronization...\n";