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/Derivative.cpp

@@ -1183,6 +1183,23 @@ void ReverseAccumulationVisitor::visit(const Call *op) {
             for (const auto &arg : op->args) {
                 accumulate(arg, make_zero(op->type));
             }
+        } else if (op->is_intrinsic(Call::branch)) {
+            // branch(cond, a, b) has the same derivative as select(cond, a, b):
+            // route the incoming adjoint to whichever side is taken. The routing
+            // is a select rather than a branch on purpose. Reverse mode multiplies
+            // each node's local derivative by the incoming adjoint, so routing the
+            // adjoint through a branch does not gate the arm derivatives - both
+            // still get computed (dcoA * branch(cond, adj, 0) keeps dcoA on the
+            // zero side). It would only duplicate the adjoint IR (both arms in each
+            // side, since deriv*0 can not be simplified for floats) with no gain.
+            // True gating of an arm's adjoint needs the arm to be its own Func; the
+            // forward branch itself is still kept as control flow (its recompute
+            // for the backward pass is lifted in ScheduleFunctions).
+            internal_assert(op->args.size() == 3);
+            accumulate(op->args[1],
+                       select(op->args[0], adjoint, make_zero(op->type)));
+            accumulate(op->args[2],
+                       select(op->args[0], make_zero(op->type), adjoint));
         } else {
             user_warning << "Dropping gradients at call to " << op->name << "\n";
             for (const auto &arg : op->args) {

src/DerivativeUtils.cpp

@@ -171,6 +171,15 @@ void ExpressionSorter::visit(const Call *op) {
         return;
     }
 
+    if (op->is_intrinsic(Call::branch)) {
+        // Like Select, ignore the condition (arg 0); its derivative is zero.
+        // Only the value arms carry gradients.
+        internal_assert(op->args.size() == 3);
+        include(op->args[1]);
+        include(op->args[2]);
+        return;
+    }
+
     for (const auto &arg : op->args) {
         include(arg);
     }

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) {
@@ -3276,6 +3287,23 @@ Func define_base_case(const Internal::Function &func, const vector<Expr> &a, con
     return f;
 }
 
+// Distribute `self <op> _` into the arms of a (possibly multi-way) branch, so
+// that f op= branch(cond, a, b) becomes f = branch(cond, self op a, self op b).
+// The result is a whole-value branch, i.e. real control flow that only computes
+// the taken arm.
+template<typename BinaryOp>
+Expr distribute_into_branch(const Expr &self, const Expr &e) {
+    if (const Call *c = Call::as_intrinsic(e, {Call::branch})) {
+        internal_assert(c->args.size() == 3);
+        return Call::make(c->type, Call::branch,
+                          {c->args[0],
+                           distribute_into_branch<BinaryOp>(self, c->args[1]),
+                           distribute_into_branch<BinaryOp>(self, c->args[2])},
+                          Call::PureIntrinsic);
+    }
+    return BinaryOp()(self, e);
+}
+
 }  // namespace
 
 template<typename BinaryOp>
@@ -3307,6 +3335,21 @@ Stage FuncRef::func_ref_update(const Expr &e, int init_val) {
     return self_ref = BinaryOp()(Expr(self_ref), e);
 }
 
+template<typename BinaryOp>
+Stage FuncRef::func_ref_update(const Branch &b, int init_val) {
+    // f op= branch(cond, a, b) distributes to f = branch(cond, f op a, f op b), so
+    // the whole update value stays a branch (real control flow) and only the taken
+    // arm's contribution is actually computed. The distribution is valid for any op
+    // because branch has the same value as select: op(f, select(c, a, b)) ==
+    // select(c, op(f, a), op(f, b)). The only per-op difference is init_val (the
+    // identity: 0 for +/-, 1 for *).
+    const vector<Expr> rhs = {b.expr};
+    const vector<Expr> expanded_args = args_with_implicit_vars(rhs);
+    FuncRef self_ref = define_base_case(func, expanded_args, rhs, init_val)(expanded_args);
+    Expr distributed = distribute_into_branch<BinaryOp>(Expr(self_ref), b.expr);
+    return self_ref = Branch{distributed};
+}
+
 Stage FuncRef::operator+=(const Expr &e) {
     return func_ref_update<std::plus<Expr>>(e, 0);
 }
@@ -3319,6 +3362,10 @@ Stage FuncRef::operator+=(const Tuple &e) {
     }
 }
 
+Stage FuncRef::operator+=(const Branch &b) {
+    return func_ref_update<std::plus<Expr>>(b, 0);
+}
+
 Stage FuncRef::operator+=(const FuncRef &e) {
     if (e.size() == 1) {
         return (*this) += Expr(e);
@@ -3331,6 +3378,12 @@ Stage FuncRef::operator*=(const Expr &e) {
     return func_ref_update<std::multiplies<Expr>>(e, 1);
 }
 
+Stage FuncRef::operator*=(const Branch &b) {
+    // Same distribution as += and -=, but the identity is 1: if the Func has no
+    // pure definition it starts at 1 and each taken arm multiplies into it.
+    return func_ref_update<std::multiplies<Expr>>(b, 1);
+}
+
 Stage FuncRef::operator*=(const Tuple &e) {
     if (e.size() == 1) {
         return (*this) *= e[0];
@@ -3351,6 +3404,10 @@ Stage FuncRef::operator-=(const Expr &e) {
     return func_ref_update<std::minus<Expr>>(e, 0);
 }
 
+Stage FuncRef::operator-=(const Branch &b) {
+    return func_ref_update<std::minus<Expr>>(b, 0);
+}
+
 Stage FuncRef::operator-=(const Tuple &e) {
     if (e.size() == 1) {
         return (*this) -= e[0];
@@ -3371,6 +3428,10 @@ Stage FuncRef::operator/=(const Expr &e) {
     return func_ref_update<std::divides<Expr>>(e, 1);
 }
 
+Stage FuncRef::operator/=(const Branch &b) {
+    return func_ref_update<std::divides<Expr>>(b, 1);
+}
+
 Stage FuncRef::operator/=(const Tuple &e) {
     if (e.size() == 1) {
         return (*this) /= e[0];

src/Func.h

@@ -58,6 +58,7 @@ struct VarOrRVar {
 };
 
 class ImageParam;
+struct Branch;
 
 namespace Internal {
 struct AssociativeOp;
@@ -518,6 +519,15 @@ class FuncRef {
     template<typename BinaryOp>
     Stage func_ref_update(const Expr &e, int init_val);
 
+    /** Helper for function update by a branch (see \ref branch). Distributes the
+     * accumulation into the arms - f op= branch(cond, a, b) becomes
+     * f = branch(cond, f op a, f op b) - so the whole update value stays a branch
+     * (real control flow) and only the taken arm's contribution is computed. If
+     * the function does not already have a pure definition, init_val is used as
+     * the RHS of the initial definition. */
+    template<typename BinaryOp>
+    Stage func_ref_update(const Branch &b, int init_val);
+
 public:
     FuncRef(const Internal::Function &, const std::vector<Expr> &,
             int placeholder_pos = -1, int count = 0);
@@ -533,6 +543,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
@@ -544,6 +559,12 @@ class FuncRef {
     Stage operator+=(const FuncRef &);
     // @}
 
+    /** Accumulate a branch (see \ref branch) into this Func. This distributes
+     * the accumulation into the arms - f += branch(cond, a, b) becomes
+     * f = branch(cond, f + a, f + b) - so the whole update value is a branch
+     * (real control flow), and only the taken arm's contribution is computed. */
+    Stage operator+=(const Branch &);
+
     /** Define a stage that adds the negative of the given expression to this
      * Func. If the expression refers to some RDom, this performs a sum reduction
      * of the negative of the expression over the domain. If the function does
@@ -555,6 +576,12 @@ class FuncRef {
     Stage operator-=(const FuncRef &);
     // @}
 
+    /** Subtract a branch (see \ref branch) from this Func. Distributes the
+     * accumulation into the arms - f -= branch(cond, a, b) becomes
+     * f = branch(cond, f - a, f - b) - so the whole update value is a branch
+     * (real control flow), and only the taken arm's contribution is computed. */
+    Stage operator-=(const Branch &);
+
     /** Define a stage that multiplies this Func by the given expression. If the
      * expression refers to some RDom, this performs a product reduction of the
      * expression over the domain. If the function does not already have a pure
@@ -566,6 +593,14 @@ class FuncRef {
     Stage operator*=(const FuncRef &);
     // @}
 
+    /** Multiply this Func by a branch (see \ref branch). Distributes the
+     * accumulation into the arms - f *= branch(cond, a, b) becomes
+     * f = branch(cond, f * a, f * b) - so the whole update value is a branch
+     * (real control flow), and only the taken arm's contribution is computed.
+     * If the Func has no pure definition it is initialised to 1 (the
+     * multiplicative identity), matching \ref operator*=(const Expr &). */
+    Stage operator*=(const Branch &);
+
     /** Define a stage that divides this Func by the given expression.
      * If the expression refers to some RDom, this performs a product
      * reduction of the inverse of the expression over the domain. If the
@@ -577,6 +612,14 @@ class FuncRef {
     Stage operator/=(const FuncRef &);
     // @}
 
+    /** Divide this Func by a branch (see \ref branch). Distributes the
+     * accumulation into the arms - f /= branch(cond, a, b) becomes
+     * f = branch(cond, f / a, f / b) - so the whole update value is a branch
+     * (real control flow), and only the taken arm's contribution is computed.
+     * If the Func has no pure definition it is initialised to 1, matching
+     * \ref operator/=(const Expr &). */
+    Stage operator/=(const Branch &);
+
     /* Override the usual assignment operator, so that
      * f(x, y) = g(x, y) defines f.
      */

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.";