diff --git a/src/Bounds.cpp b/src/Bounds.cpp index 9087516b54b8..3c0d9ba8feae 100644 --- a/src/Bounds.cpp +++ b/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)) { diff --git a/src/CSE.cpp b/src/CSE.cpp index 95f61f0214e9..30f255bb88c3 100644 --- a/src/CSE.cpp +++ b/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. */ diff --git a/src/Derivative.cpp b/src/Derivative.cpp index a7e9ade253fe..2f8bed0326da 100644 --- a/src/Derivative.cpp +++ b/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) { diff --git a/src/DerivativeUtils.cpp b/src/DerivativeUtils.cpp index 54e1c37e0b84..23878bfa4ba6 100644 --- a/src/DerivativeUtils.cpp +++ b/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); } diff --git a/src/Func.cpp b/src/Func.cpp index f4e4298fae07..dfd105a60b5e 100644 --- a/src/Func.cpp +++ b/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 &a, con return f; } +// Distribute `self _` 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 +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(self, c->args[1]), + distribute_into_branch(self, c->args[2])}, + Call::PureIntrinsic); + } + return BinaryOp()(self, e); +} + } // namespace template @@ -3307,6 +3335,21 @@ Stage FuncRef::func_ref_update(const Expr &e, int init_val) { return self_ref = BinaryOp()(Expr(self_ref), e); } +template +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 rhs = {b.expr}; + const vector 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(Expr(self_ref), b.expr); + return self_ref = Branch{distributed}; +} + Stage FuncRef::operator+=(const Expr &e) { return func_ref_update>(e, 0); } @@ -3319,6 +3362,10 @@ Stage FuncRef::operator+=(const Tuple &e) { } } +Stage FuncRef::operator+=(const Branch &b) { + return func_ref_update>(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>(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>(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>(e, 0); } +Stage FuncRef::operator-=(const Branch &b) { + return func_ref_update>(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>(e, 1); } +Stage FuncRef::operator/=(const Branch &b) { + return func_ref_update>(b, 1); +} + Stage FuncRef::operator/=(const Tuple &e) { if (e.size() == 1) { return (*this) /= e[0]; diff --git a/src/Func.h b/src/Func.h index 0bfb591871c7..faa18f165de7 100644 --- a/src/Func.h +++ b/src/Func.h @@ -58,6 +58,7 @@ struct VarOrRVar { }; class ImageParam; +struct Branch; namespace Internal { struct AssociativeOp; @@ -518,6 +519,15 @@ class FuncRef { template 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 + Stage func_ref_update(const Branch &b, int init_val); + public: FuncRef(const Internal::Function &, const std::vector &, 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. */ diff --git a/src/IR.cpp b/src/IR.cpp index ae0e38a36251..d556e1b68656 100644 --- a/src/IR.cpp +++ b/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", diff --git a/src/IR.h b/src/IR.h index 000d8a994851..4f5ac821939c 100644 --- a/src/IR.h +++ b/src/IR.h @@ -620,6 +620,15 @@ struct Call : public ExprNode { // 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 may only be the whole value of a definition, and is + // turned into a Stmt-level IfThenElse (one store per arm) in + // ScheduleFunctions, then hoisted to the loop level at which its + // condition is invariant. If its condition varies across the lanes of a + // vectorized loop it stays an IfThenElse and becomes predicated + // loads/stores in VectorizeLoops. + branch, // Bundle multiple exprs together temporarily for analysis (e.g. CSE) bundle, // Takes a sequence of (condition, function) pairs, and calls the first diff --git a/src/IROperator.cpp b/src/IROperator.cpp index c109f1d2ccff..044f01a27eb3 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -1537,6 +1537,42 @@ 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, which can not be " + << "taken per lane. Use select() for an already-vector condition. (A " + << "scalar condition that merely varies across a vectorized dimension " + << "is fine: it becomes a predicated load/store.)\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."; diff --git a/src/IROperator.h b/src/IROperator.h index 797f12870f5d..d79a35ccc0ed 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -851,6 +851,63 @@ 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 the definition + * can be lowered to real control flow with one store per arm. */ +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); + * A branch-defined Func can therefore not be inlined into a consumer; schedule + * it with compute_root() or compute_at() instead. + * + * The branch is hoisted to the outermost loop level at which its condition is + * loop-invariant, so a producer compute_at'd at a loop level inside the branch + * is only computed when that side is taken. That is how a branch gates real + * computation rather than just a load. + * + * 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. + * + * If the condition varies across the lanes of a vectorized dimension, it can + * not be a scalar jump, so it instead becomes predication: the loads and stores + * of each arm are masked, and only the active lanes access memory (if an arm + * can not be predicated, the statement is scalarized instead). The untaken side + * still never loads, stores, or faults, which makes branch() a natural way to + * express a boundary condition: + * f(x) = branch(x >= 0 && x < w, in(unsafe_promise_clamped(x, 0, w-1)), 0); + * compiles to a predicated vector load on targets that have one (e.g. AVX-512). + * The unsafe_promise_clamped is what tells bounds inference that the index + * stays in range, so `in` is not required outside its real extent. + * + * Combining branch() with vectorization inside a GPU kernel is a hard error + * (not a silent fallback to select): use select() for the vectorized + * dimension. */ +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::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)...).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 diff --git a/src/Lower.cpp b/src/Lower.cpp index cd7ccc9a03f4..21b2c48a586f 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -354,6 +354,7 @@ void lower_impl(const vector &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"; diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index a0b4a4c18876..220db11911cb 100644 --- a/src/ScheduleFunctions.cpp +++ b/src/ScheduleFunctions.cpp @@ -11,6 +11,7 @@ #include "IRMutator.h" #include "IROperator.h" #include "IRPrinter.h" +#include "IRVisitor.h" #include "Inline.h" #include "Prefetch.h" #include "Qualify.h" @@ -467,6 +468,195 @@ Stmt build_loop_nest( } // Build a loop nest about a provide node using a schedule +// True if the expression contains a branch() intrinsic anywhere. +class ContainsBranch : public IRGraphVisitor { + using IRGraphVisitor::visit; + void visit(const Call *op) override { + if (op->is_intrinsic(Call::branch)) { + found = true; + } + IRGraphVisitor::visit(op); + } + +public: + bool found = false; +}; + +bool expr_contains_branch(const Expr &e) { + ContainsBranch v; + e.accept(&v); + return v.found; +} + +// A branch() value is turned into real point-wise control flow here: an +// IfThenElse around one Provide per arm (recursing for the multi-way form). +// This is the same idea as specialize(), but point-wise rather than func-wise. +Stmt build_branch_provide(const string &name, const Expr &value, const vector &site) { + if (const Call *c = Call::as_intrinsic(value, {Call::branch})) { + internal_assert(c->args.size() == 3); + Stmt then_case = build_branch_provide(name, c->args[1], site); + Stmt else_case = build_branch_provide(name, c->args[2], site); + return IfThenElse::make(c->args[0], then_case, else_case); + } + return Provide::make(name, {value}, site, const_true()); +} + +// True if s is the control-flow tree we just built for a branch: an IfThenElse +// (with a real else) whose leaves are all Provides to this func. Loops are +// allowed inside, so this still matches after we have hoisted it out of a loop. +bool is_branch_stmt(const Stmt &s, const string &name) { + if (const Provide *p = s.as()) { + return p->name == name; + } + if (const IfThenElse *i = s.as()) { + return i->else_case.defined() && + is_branch_stmt(i->then_case, name) && + is_branch_stmt(i->else_case, name); + } + if (const For *f = s.as()) { + return is_branch_stmt(f->body, name); + } + return false; +} + +// Hoist a branch out of every loop whose variable its condition does not use, +// so it ends up at the outermost loop level where it is invariant: +// for v { if (c) A else B } -> if (c) { for v A } else { for v B } +// Producers compute_at'd at a level inside the branch then only run when that +// side is taken, which is what gates their computation (rather than inlining). +class HoistBranchOutOfLoops : public IRMutator { + using IRMutator::visit; + + const string &func_name; + + Stmt visit(const For *op) override { + Stmt body = mutate(op->body); + const IfThenElse *i = body.as(); + if (i != nullptr && + is_branch_stmt(body, func_name) && + !expr_uses_var(i->condition, op->name)) { + Stmt then_loop = For::make(op->name, op->min, op->max, op->for_type, + op->partition_policy, op->device_api, i->then_case); + Stmt else_loop = For::make(op->name, op->min, op->max, op->for_type, + op->partition_policy, op->device_api, i->else_case); + return IfThenElse::make(i->condition, then_loop, else_loop); + } + if (body.same_as(op->body)) { + return op; + } + return For::make(op->name, op->min, op->max, op->for_type, + op->partition_policy, op->device_api, body); + } + +public: + explicit HoistBranchOutOfLoops(const string &n) + : func_name(n) { + } +}; + +// Collect the conditions of a (possibly multi-way) branch value. +void collect_branch_conditions(const Expr &value, vector &conds) { + if (const Call *c = Call::as_intrinsic(value, {Call::branch})) { + conds.push_back(c->args[0]); + collect_branch_conditions(c->args[2], conds); + } +} + +// On CPU a lane-varying condition is fine: it becomes predication (see below). +// Inside a GPU kernel we do not have that fallback, so combining branch() with +// vectorization there is a hard error rather than a silent fallback to select. +void check_branch_schedule(const Function &func, const Definition &def) { + bool any_vectorized = false, any_gpu = false; + for (const Dim &d : def.schedule().dims()) { + any_vectorized |= (d.for_type == ForType::Vectorized); + any_gpu |= (d.for_type == ForType::GPUBlock || + d.for_type == ForType::GPUThread || + d.for_type == ForType::GPULane); + } + user_assert(!(any_gpu && any_vectorized)) + << "branch() in Func \"" << func.name() << "\" is used in a GPU kernel that " + << "is also vectorized. On GPU only a per-thread or per-wave " + << "(non-vectorized) branch is allowed. Use select() for the vectorized " + << "dimension.\n"; +} + +// After hoisting, a branch that is still nested inside a vectorized loop has a +// lane-varying condition (otherwise it would have been hoisted out of that +// loop). We leave it as an IfThenElse and let VectorizeLoops handle it: it +// pushes the vector condition down onto the loads and stores of each arm as a +// predicate, so only the active lanes actually access memory. If some arm can +// not be predicated, VectorizeLoops scalarizes the statement instead, which +// restores true per-lane control flow. Both are faithful to branch(): the +// untaken side never loads, stores, or faults. This is what makes branch() +// useful for boundary conditions, e.g. +// f(x) = branch(x >= 0 && x < w, in(unsafe_promise_clamped(x, 0, w-1)), 0) +// becomes a predicated (masked) vector load rather than a clamp plus a select. + +// Find the condition of some branch() buried in an expression, or an undefined +// Expr if there is none. +class FindBranchCondition : public IRGraphVisitor { + using IRGraphVisitor::visit; + void visit(const Call *op) override { + if (!found.defined() && op->is_intrinsic(Call::branch)) { + found = op->args[0]; + } + IRGraphVisitor::visit(op); + } + +public: + Expr found; +}; + +// Replace every branch() whose condition equals `cond` by one of its arms +// (arg index 1 = true side, 2 = false side), collapsing that one condition. +class PickBranchArm : public IRMutator { + using IRMutator::visit; + Expr cond; + int arm; + Expr visit(const Call *op) override { + if (op->is_intrinsic(Call::branch) && equal(op->args[0], cond)) { + return mutate(op->args[arm]); + } + return IRMutator::visit(op); + } + +public: + PickBranchArm(Expr c, int a) + : cond(std::move(c)), arm(a) { + } +}; + +// Lift every branch() in a point-wise value up to the top, producing a whole- +// value (possibly nested / multi-way) branch. For a point-wise surrounding +// expression E, E[branch(c, a, b)] == branch(c, E[a], E[b]), because branch has +// the same value as select and pure point-wise ops distribute over select. This +// undoes the burying that reverse-mode autodiff does when it multiplies a +// branch-valued forward Func into a larger adjoint expression: the branch is +// recovered as real control flow instead of being an illegal sub-expression. +// Branches that share a condition are resolved together, so k distinct +// conditions give at most 2^k arms (in practice one). +Expr lift_branches(const Expr &value) { + if (const Call *c = Call::as_intrinsic(value, {Call::branch})) { + // Already a top-level branch; just lift any branches inside the arms. + return Call::make(c->type, Call::branch, + {c->args[0], + lift_branches(c->args[1]), + lift_branches(c->args[2])}, + Call::PureIntrinsic); + } + FindBranchCondition f; + value.accept(&f); + if (!f.found.defined()) { + return value; + } + Expr cond = f.found; + Expr t = PickBranchArm(cond, 1)(value); + Expr e = PickBranchArm(cond, 2)(value); + return Call::make(value.type(), Call::branch, + {cond, lift_branches(t), lift_branches(e)}, + Call::PureIntrinsic); +} + Stmt build_provide_loop_nest(const map &env, const string &prefix, const Function &func, @@ -494,8 +684,39 @@ Stmt build_provide_loop_nest(const map &env, debug(3) << "Site " << i << " = " << s << "\n"; } - // Make the (multi-dimensional multi-valued) store node. - Stmt body = Provide::make(func.name(), values, site, const_true()); + // If a branch() got buried inside a larger point-wise value - typically + // because reverse-mode autodiff multiplied a branch-valued forward Func into + // an adjoint expression - lift it back to the top so it can become real + // control flow. E[branch(c, a, b)] == branch(c, E[a], E[b]) for point-wise E. + if (values.size() == 1 && expr_contains_branch(values[0]) && + !Call::as_intrinsic(values[0], {Call::branch})) { + // Flatten Expr-level Lets first: lifting pulls the branch condition to + // the top, and it must not reference a Let variable whose binding lives + // inside the surrounding expression. CSE re-compresses later in lowering. + values[0] = lift_branches(substitute_in_all_lets(values[0])); + } + + // Make the (multi-dimensional multi-valued) store node. If the value is a + // branch(), build real point-wise control flow instead of a single Provide. + vector branch_conds; + Stmt body; + if (values.size() == 1 && Call::as_intrinsic(values[0], {Call::branch})) { + collect_branch_conditions(values[0], branch_conds); + check_branch_schedule(func, def); + body = build_branch_provide(func.name(), values[0], site); + } else { + // A branch may only be the whole value of a definition. It can survive + // buried inside a multi-valued (Tuple) definition, which we can not turn + // into point-wise control flow. + for (const Expr &v : values) { + user_assert(!expr_contains_branch(v)) + << "A Func defined with branch() was inlined into \"" << func.name() + << "\". A branch must be the whole value of a definition, so a " + << "branch-defined Func can not be inlined. Schedule it with " + << "compute_root() or compute_at().\n"; + } + body = Provide::make(func.name(), values, site, const_true()); + } if (def.schedule().atomic()) { // Add atomic node. bool any_unordered_parallel = false; for (const auto &d : def.schedule().dims()) { @@ -517,6 +738,15 @@ Stmt build_provide_loop_nest(const map &env, // Default schedule/values if there is no specialization Stmt stmt = build_loop_nest(body, prefix, start_fuse, func, def); + if (!branch_conds.empty()) { + // Hoist the branch to the outermost loop level where its condition is + // invariant. Producers compute_at'd at a level inside the branch then + // only run when that side is taken, which gates their computation. + stmt = HoistBranchOutOfLoops(func.name())(stmt); + // Anything still stuck inside a vectorized loop has a lane-varying + // condition. It stays an IfThenElse and becomes predicated loads and + // stores (or a scalarized statement) in VectorizeLoops. + } stmt = inject_placeholder_prefetch(stmt, env, prefix, def.schedule().prefetches()); // Make any specialized copies @@ -2606,6 +2836,16 @@ Stmt schedule_functions(const vector &outputs, } if (group_should_be_inlined(funcs)) { + // A branch() becomes real control flow, so a Func defined with one + // needs its own loop nest: inlining it would bury the branch inside + // a consumer's expression, where it can not be turned into a branch. + for (const Expr &v : funcs[0].definition().values()) { + user_assert(!expr_contains_branch(v)) + << "Func \"" << funcs[0].name() << "\" is defined with branch(), " + << "which becomes real control flow, so it can not be inlined " + << "into its consumer. Schedule it with compute_root() or " + << "compute_at().\n"; + } debug(1) << "Inlining " << funcs[0].name() << "\n"; s = inline_function(s, funcs[0]); } else { diff --git a/src/VectorizeLoops.cpp b/src/VectorizeLoops.cpp index f97099287003..94ba852fa957 100644 --- a/src/VectorizeLoops.cpp +++ b/src/VectorizeLoops.cpp @@ -339,8 +339,20 @@ class PredicateLoadStore : public IRMutator { } Expr visit(const Call *op) override { - // We should not vectorize calls with side-effects - valid = valid && op->is_pure(); + // We should not vectorize calls with side-effects. + // + // promise_clamped and unsafe_promise_clamped are rebuilt as impure + // intrinsics by the simplifier (see Simplify_Call.cpp), so that they do + // not get lifted out of the if they sit in. They have no side-effects + // though, and predicating them is safe: it turns the surrounding if + // into a mask on the load/store rather than moving the promise. Without + // this, a boundary condition written as + // f(x) = branch(x >= 0 && x < w, in(unsafe_promise_clamped(x, ...)), 0) + // would scalarize instead of becoming a predicated load. Other passes + // whitelist these the same way (see Deinterleave.cpp). + valid = valid && (op->is_pure() || + op->is_intrinsic(Call::unsafe_promise_clamped) || + op->is_intrinsic(Call::promise_clamped)); return IRMutator::visit(op); } diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 66e57183ad63..10e75535bd1e 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -33,6 +33,7 @@ tests( bounds_of_split.cpp bounds_query.cpp bounds_query_respects_specialize_fail.cpp + branch.cpp buffer_t.cpp c_function.cpp callable.cpp diff --git a/test/correctness/branch.cpp b/test/correctness/branch.cpp new file mode 100644 index 000000000000..d406a3aad84a --- /dev/null +++ b/test/correctness/branch.cpp @@ -0,0 +1,1134 @@ +#include "Halide.h" +#include +#include +#include +#include + +using namespace Halide; +using namespace Halide::Internal; + +// Test branch(cond, a, b): a strict variant of select. branch() returns a +// Branch (not an Expr) and may only be the whole right-hand side of a Func +// definition. It is turned into real point-wise control flow in +// ScheduleFunctions (an IfThenElse with one store per arm, like specialize but +// point-wise), hoisted to the outermost loop level where its condition is +// invariant, so producers can be compute_at'd inside it and have their +// computation gated by the branch. + +namespace { + +// Is there a Stmt-level IfThenElse (with a real else) anywhere in the body? +class HasRealBranch : public IRVisitor { + using IRVisitor::visit; + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + found = true; + } + IRVisitor::visit(op); + } + +public: + bool found = false; +}; + +// Is a ProducerConsumer "produce " nested inside an IfThenElse? That is +// what tells us the producer's computation is gated by the branch. +class ProduceInsideBranch : public IRVisitor { + using IRVisitor::visit; + + int depth = 0; + + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + depth++; + IRVisitor::visit(op); + depth--; + } else { + IRVisitor::visit(op); + } + } + + void visit(const ProducerConsumer *op) override { + if (op->is_producer && op->name == target && depth > 0) { + found = true; + } + IRVisitor::visit(op); + } + +public: + std::string target; + bool found = false; + explicit ProduceInsideBranch(std::string t) + : target(std::move(t)) { + } +}; + +template +bool scan_module(const Module &m, V &v) { + for (const auto &fn : m.functions()) { + fn.body.accept(&v); + } + return v.found; +} + +// Count Stmt-level IfThenElse nodes that have a real else (i.e. branches). +class CountRealBranches : public IRVisitor { + using IRVisitor::visit; + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + count++; + } + IRVisitor::visit(op); + } + +public: + int count = 0; +}; + +int count_real_branches(const Module &m) { + CountRealBranches v; + for (const auto &fn : m.functions()) { + fn.body.accept(&v); + } + return v.count; +} + +// Count "produce " nodes inside a Stmt. +class CountProduce : public IRVisitor { + using IRVisitor::visit; + void visit(const ProducerConsumer *op) override { + if (op->is_producer && op->name == target) { + count++; + } + IRVisitor::visit(op); + } + +public: + std::string target; + int count = 0; + explicit CountProduce(std::string t) + : target(std::move(t)) { + } +}; + +int count_produce(const Stmt &s, const std::string &name) { + CountProduce v(name); + s.accept(&v); + return v.count; +} + +// For the outermost real branch, how many times is "produce " found in +// its then-side vs its else-side. Used to check a producer is only computed in +// the arm that actually uses it. +class BranchArmProduce : public IRVisitor { + using IRVisitor::visit; + void visit(const IfThenElse *op) override { + if (!found && op->else_case.defined()) { + found = true; + then_count = count_produce(op->then_case, target); + else_count = count_produce(op->else_case, target); + return; + } + IRVisitor::visit(op); + } + +public: + std::string target; + bool found = false; + int then_count = 0, else_count = 0; + explicit BranchArmProduce(std::string t) + : target(std::move(t)) { + } +}; + +BranchArmProduce arm_produce(const Module &m, const std::string &name) { + BranchArmProduce v(name); + for (const auto &fn : m.functions()) { + fn.body.accept(&v); + } + return v; +} + +// True if some Load has a non-trivial predicate, i.e. it became a predicated +// (masked) load rather than an unconditional one. +class HasPredicatedLoad : public IRVisitor { + using IRVisitor::visit; + void visit(const Load *op) override { + if (!is_const_one(op->predicate)) { + found = true; + } + IRVisitor::visit(op); + } + +public: + bool found = false; +}; + +// True if a real branch appears nested inside any For loop. +class BranchInsideAnyLoop : public IRVisitor { + using IRVisitor::visit; + + int depth = 0; + + void visit(const For *op) override { + depth++; + IRVisitor::visit(op); + depth--; + } + + void visit(const IfThenElse *op) override { + if (op->else_case.defined() && depth > 0) { + found = true; + } + IRVisitor::visit(op); + } + +public: + bool found = false; +}; + +// True if some branch (IfThenElse with a real else) has a For loop inside it, +// i.e. the branch was hoisted above that loop. +class BranchAboveLoop : public IRVisitor { + using IRVisitor::visit; + + bool in_branch = false; + + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + bool old = in_branch; + in_branch = true; + IRVisitor::visit(op); + in_branch = old; + } else { + IRVisitor::visit(op); + } + } + + void visit(const For *op) override { + if (in_branch) { + found = true; + } + IRVisitor::visit(op); + } + +public: + bool found = false; +}; + +} // namespace + +int main(int argc, char **argv) { + Var x("x"), y("y"); + + // Part A: branch() returns a Branch wrapping the branch intrinsic. + { + Expr cond = x > 5; + Expr t = x * 2; + Expr f = x + 1; + + Branch b = branch(cond, t, f); + const Call *c = Call::as_intrinsic(b.expr, {Call::branch}); + if (!c || c->args.size() != 3 || + !equal(c->args[0], cond) || !equal(c->args[1], t) || !equal(c->args[2], f)) { + printf("branch() did not produce a well-formed branch intrinsic\n"); + return 1; + } + } + + // Part B: branch() is numerically equivalent to select(). + { + Expr cond = (x % 3) == 0; + Expr t = x * x + 7; + Expr f = x * 5 - 2; + + Func sel("sel"), br("br"); + sel(x) = select(cond, t, f); + br(x) = branch(cond, t, f); + + Buffer rs = sel.realize({256}); + Buffer rb = br.realize({256}); + for (int i = 0; i < rs.width(); i++) { + if (rs(i) != rb(i)) { + printf("Mismatch at %d: select=%d branch=%d\n", i, rs(i), rb(i)); + return 1; + } + } + } + + // Part C: the multi-way branch produces correct values. + { + Func f("multi"); + f(x) = branch(x < 10, 1, x < 20, 2, 3); + Buffer r = f.realize({30}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 10 ? 1 : (i < 20 ? 2 : 3); + if (r(i) != correct) { + printf("multi-way mismatch at %d: got %d want %d\n", i, r(i), correct); + return 1; + } + } + } + + // Part D: the branch becomes REAL control flow (a Stmt-level IfThenElse + // with a store per arm), not an if_then_else expression in a single store. + { + Func f("real_cf"); + f(x) = branch(x > 3, x * 2, x + 1); + Module m = f.compile_to_module({}, "real_cf"); + + HasRealBranch hb; + if (!scan_module(m, hb)) { + printf("branch did not become a Stmt-level IfThenElse\n"); + return 1; + } + } + + // Part E: when the condition depends on the innermost var, the branch sits + // inside that loop, so there is no loop level inside the branch to + // compute_at onto. The producer is therefore NOT gated here (that is + // inherent, not a bug) - we only check the result is still correct. + { + Func p("gated"), f("uses_gated"); + p(x) = x * x + 7; + f(x) = branch(x < 128, p(x), x + 1); + p.compute_at(f, x); + + Buffer r = f.realize({256}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 128 ? (i * i + 7) : (i + 1); + if (r(i) != correct) { + printf("producer mismatch at %d: got %d want %d\n", i, r(i), correct); + return 1; + } + } + } + + // Part F: the condition is hoisted out of loops it does not depend on, so a + // producer can be compute_at'd at an inner level inside the hoisted branch. + { + Func p("hoist_p"), f("hoist_f"); + p(x, y) = x + y; + f(x, y) = branch(y < 8, p(x, y), x - y); // condition uses only y + p.compute_at(f, x); + + Module m = f.compile_to_module({}, "hoist_f"); + ProduceInsideBranch pib("hoist_p"); + if (!scan_module(m, pib)) { + printf("branch was not hoisted / producer not gated\n"); + return 1; + } + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? (i + j) : (i - j); + if (r(i, j) != correct) { + printf("hoist mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part G: branch() works inside a GPU kernel (a real per-thread branch). + { + Target t = get_jit_target_from_environment(); + if (t.has_gpu_feature()) { + Var xo("xo"), xi("xi"); + Func gsel("gsel"), gbr("gbr"); + gsel(x) = select(x % 2 == 0, x * x + 7, x * 5 - 2); + gbr(x) = branch(x % 2 == 0, x * x + 7, x * 5 - 2); + gsel.gpu_tile(x, xo, xi, 16); + gbr.gpu_tile(x, xo, xi, 16); + Buffer rs = gsel.realize({256}, t); + Buffer rb = gbr.realize({256}, t); + rs.copy_to_host(); + rb.copy_to_host(); + for (int i = 0; i < 256; i++) { + if (rs(i) != rb(i)) { + printf("GPU mismatch at %d: select=%d branch=%d\n", i, rs(i), rb(i)); + return 1; + } + } + } else { + printf("[gpu part skipped: jit target has no gpu feature]\n"); + } + } + + // Part L: a multi-way branch becomes a real if / else-if / else chain, not + // just a correct value. + { + Func f("multi_cf"); + f(x) = branch(x < 10, 1, x < 20, 2, 3); + Module m = f.compile_to_module({}, "multi_cf"); + if (count_real_branches(m) < 2) { + printf("multi-way branch did not become a chain of real branches\n"); + return 1; + } + } + + // Part M: the branch is hoisted out of loops its condition does not use, + // and is NOT hoisted when the condition uses the innermost loop var. + { + Func fo("hoist_outer"); + fo(x, y) = branch(y < 8, x + y, x - y); // condition uses only y + Module mo = fo.compile_to_module({}, "hoist_outer"); + BranchAboveLoop above; + if (!scan_module(mo, above)) { + printf("branch was not hoisted above the inner loop\n"); + return 1; + } + + Func fi("hoist_inner"); + fi(x, y) = branch(x < 8, x + y, x - y); // condition uses innermost var + Module mi = fi.compile_to_module({}, "hoist_inner"); + BranchAboveLoop above2; + if (scan_module(mi, above2)) { + printf("branch was hoisted out of a loop its condition uses\n"); + return 1; + } + } + + // Part N: a compute_root producer is NOT gated by the branch (its produce + // loop sits outside). This is the inherent bounds limitation, not a bug. + { + Func p("root_p"), f("root_f"); + p(x) = x * x; + f(x) = branch(x < 128, p(x), x + 1); + p.compute_root(); + + Module m = f.compile_to_module({}, "root_f"); + ProduceInsideBranch pib("root_p"); + if (scan_module(m, pib)) { + printf("a compute_root producer should not be inside the branch\n"); + return 1; + } + + Buffer r = f.realize({256}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 128 ? (i * i) : (i + 1); + if (r(i) != correct) { + printf("compute_root mismatch at %d\n", i); + return 1; + } + } + } + + // Part O: a branch works as an update definition over an RDom, updating + // only part of the domain and leaving the rest from the pure definition. + { + Func f("upd_branch"); + f(x) = -1; + RDom r(0, 6); + f(r) = branch(r < 3, 100, 200); + Buffer res = f.realize({10}); + for (int i = 0; i < res.width(); i++) { + int correct = (i < 3) ? 100 : (i < 6 ? 200 : -1); + if (res(i) != correct) { + printf("update-stage branch mismatch at %d: got %d want %d\n", + i, res(i), correct); + return 1; + } + } + } + + // Part P: branch composes with specialize (both make control flow). + { + Param sp("sp"); + Func f("spec_branch"); + f(x) = branch(x < 5, x * 2, x + 1); + f.specialize(sp); + + for (bool v : {true, false}) { + sp.set(v); + Buffer r = f.realize({10}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 5 ? (i * 2) : (i + 1); + if (r(i) != correct) { + printf("specialize+branch mismatch at %d (sp=%d)\n", i, (int)v); + return 1; + } + } + } + } + + // Part Q: branch on a non-integer type. + { + Func f("float_branch"); + f(x) = branch(x < 5, cast(x) * 2.5f, cast(x) + 1.5f); + Buffer r = f.realize({10}); + for (int i = 0; i < r.width(); i++) { + float correct = i < 5 ? (i * 2.5f) : (i + 1.5f); + if (std::abs(r(i) - correct) > 1e-5f) { + printf("float branch mismatch at %d: got %f want %f\n", + i, r(i), correct); + return 1; + } + } + } + + // Part R: a producer used by BOTH arms, compute_at'd inside the hoisted + // branch, is computed in whichever side runs. + { + Func p("shared_p"), f("shared_f"); + p(x, y) = x + y; + f(x, y) = branch(y < 8, p(x, y) * 2, p(x, y) + 1); + p.compute_at(f, x); + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? ((i + j) * 2) : ((i + j) + 1); + if (r(i, j) != correct) { + printf("shared-producer mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part S: branch inside a parallel loop. + { + Func f("par_branch"); + f(x, y) = branch(y < 8, x + y, x - y); + f.parallel(y); + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? (i + j) : (i - j); + if (r(i, j) != correct) { + printf("parallel branch mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part T: SELECTIVE gating. Each arm uses a different producer; with both + // compute_at'd inside the hoisted branch, each producer should only be + // computed in the arm that actually uses it. + { + Func g("sel_g"), h("sel_h"), f("selective"); + g(x, y) = x + y; + h(x, y) = x * y; + f(x, y) = branch(y < 8, g(x, y), h(x, y)); + g.compute_at(f, x); + h.compute_at(f, x); + + Module m = f.compile_to_module({}, "selective"); + BranchArmProduce pg = arm_produce(m, "sel_g"); + BranchArmProduce ph = arm_produce(m, "sel_h"); + if (!pg.found) { + printf("selective: no real branch found\n"); + return 1; + } + if (pg.then_count < 1 || pg.else_count != 0) { + printf("selective: sel_g should be produced only in the then arm " + "(then=%d else=%d)\n", + pg.then_count, pg.else_count); + return 1; + } + if (ph.else_count < 1 || ph.then_count != 0) { + printf("selective: sel_h should be produced only in the else arm " + "(then=%d else=%d)\n", + ph.then_count, ph.else_count); + return 1; + } + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? (i + j) : (i * j); + if (r(i, j) != correct) { + printf("selective mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part U: a fully loop-invariant condition (a Param) hoists above every + // loop, so the whole loop nest sits inside the branch. + { + Param q("q"); + Func f("invariant"); + f(x, y) = branch(q > 0, x + y, x - y); + + Module m = f.compile_to_module({q}, "invariant"); + BranchInsideAnyLoop inside; + if (scan_module(m, inside)) { + printf("a loop-invariant branch was not hoisted above all loops\n"); + return 1; + } + + q.set(1); + Buffer r = f.realize({8, 8}); + for (int j = 0; j < 8; j++) { + for (int i = 0; i < 8; i++) { + if (r(i, j) != i + j) { + printf("invariant branch mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part V: a multi-way branch whose conditions live at DIFFERENT loop levels. + // The outer condition (on y) can hoist out of the x loop, but the inner one + // (on x) must not be dragged out with it. + { + Func f("mixed_levels"); + f(x, y) = branch(y < 8, 1, x < 4, 2, 3); + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? 1 : (i < 4 ? 2 : 3); + if (r(i, j) != correct) { + printf("mixed-level mismatch at (%d,%d): got %d want %d\n", + i, j, r(i, j), correct); + return 1; + } + } + } + } + + // Part W: a branch-defined Func used as a compute_root intermediate (it can + // not be inlined, so it must be scheduled). + { + Func b("mid_branch"), f("mid_consumer"); + b(x) = branch(x < 5, x * 2, x + 1); + b.compute_root(); + f(x) = b(x) + 100; + + Buffer r = f.realize({10}); + for (int i = 0; i < r.width(); i++) { + int correct = (i < 5 ? i * 2 : i + 1) + 100; + if (r(i) != correct) { + printf("intermediate branch mismatch at %d\n", i); + return 1; + } + } + } + + // Part X: chained branch Funcs (a branch consuming another branch). + { + Func b1("chain1"), b2("chain2"), f("chain_out"); + b1(x) = branch(x < 5, x * 2, x + 1); + b1.compute_root(); + b2(x) = branch(x % 2 == 0, b1(x), 0); + b2.compute_root(); + f(x) = b2(x); + + Buffer r = f.realize({10}); + for (int i = 0; i < r.width(); i++) { + int inner = (i < 5) ? i * 2 : i + 1; + int correct = (i % 2 == 0) ? inner : 0; + if (r(i) != correct) { + printf("chained branch mismatch at %d: got %d want %d\n", + i, r(i), correct); + return 1; + } + } + } + + // Part Y: branch survives tiling (the condition ends up in terms of the + // split vars). + { + Func f("tiled"); + Var xo("xo"), yo("yo"), xi("xi"), yi("yi"); + f(x, y) = branch(y < 8, x + y, x - y); + f.tile(x, y, xo, yo, xi, yi, 4, 4); + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? (i + j) : (i - j); + if (r(i, j) != correct) { + printf("tiled branch mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part Z: branch survives reorder (the condition var becomes the inner loop). + { + Func f("reordered"); + f(x, y) = branch(y < 8, x + y, x - y); + f.reorder(y, x); // y is now the inner loop -> can not hoist out of it + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? (i + j) : (i - j); + if (r(i, j) != correct) { + printf("reordered branch mismatch at (%d,%d)\n", i, j); + return 1; + } + } + } + } + + // Part AA: a stencil producer (needs a window) compute_at'd inside a hoisted + // branch still gets correct bounds. + { + Func p("stencil_p"), f("stencil_f"); + p(x, y) = x + y; + f(x, y) = branch(y < 8, p(x - 1, y) + p(x + 1, y), 0); + p.compute_at(f, x); + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = (j < 8) ? ((i - 1 + j) + (i + 1 + j)) : 0; + if (r(i, j) != correct) { + printf("stencil mismatch at (%d,%d): got %d want %d\n", + i, j, r(i, j), correct); + return 1; + } + } + } + } + + // Part AB: branch survives unrolling. + { + Func f("unrolled"); + Var xo("xo"), xi("xi"); + f(x) = branch(x < 5, x * 2, x + 1); + f.split(x, xo, xi, 4).unroll(xi); + + Buffer r = f.realize({16}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 5 ? i * 2 : i + 1; + if (r(i) != correct) { + printf("unrolled branch mismatch at %d\n", i); + return 1; + } + } + } + + // Part AC: branch in an atomic update. The atomic node wraps the branch, so + // the branch is not hoisted there, but it must still be correct. + { + Func f("atomic_branch"); + f(x) = 0; + RDom r(0, 10); + f(r) = branch(r < 5, 100, 200); + f.update(0).atomic(true); + + Buffer res = f.realize({10}); + for (int i = 0; i < res.width(); i++) { + int correct = i < 5 ? 100 : 200; + if (res(i) != correct) { + printf("atomic branch mismatch at %d: got %d want %d\n", + i, res(i), correct); + return 1; + } + } + } + + // Part AD: branch composes with fused loop nests (compute_with). + { + Func fa("cw_a"), fb("cw_b"), out("cw_out"); + fa(x, y) = branch(y < 8, x + y, x - y); + fb(x, y) = x * y; + out(x, y) = fa(x, y) + fb(x, y); + fa.compute_root(); + fb.compute_root(); + fb.compute_with(fa, y); + + Buffer r = out.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = ((j < 8) ? (i + j) : (i - j)) + i * j; + if (r(i, j) != correct) { + printf("compute_with branch mismatch at (%d,%d): got %d want %d\n", + i, j, r(i, j), correct); + return 1; + } + } + } + } + + // Part AF: reverse-mode (VJP) autodiff, the default Halide autodiff. The + // gradient through a branch must match the gradient through the equivalent + // select. + { + Param p("pv"); + p.set(3.0f); + Var xx("xx"); + Expr cond = xx < 5; + Expr a = p * cast(xx); + Expr b = p * p; + Func fs("vjp_sel"), fb("vjp_br"); + fs(xx) = select(cond, a, b); + fb(xx) = branch(cond, a, b); + fs.compute_root(); + fb.compute_root(); // branch must be a real Func, not inlined into loss + RDom r(0, 10); + Func loss_s("loss_s"), loss_b("loss_b"); + loss_s() = sum(fs(r)); + loss_b() = sum(fb(r)); + Func gs = propagate_adjoints(loss_s)(p); + Func gb = propagate_adjoints(loss_b)(p); + Buffer bs = gs.realize(); + Buffer bb = gb.realize(); + if (std::abs(bs() - bb()) > 1e-3f) { + printf("VJP gradient mismatch: select=%f branch=%f\n", bs(), bb()); + return 1; + } + } + + // Part AH: f(x) += branch(cond, a, b) distributes the accumulation into the + // arms - f = branch(cond, f + a, f + b) - so the whole update value is a + // branch (real control flow), not a select. A producer used by one arm and + // compute_at'd inside the update is gated to that arm only. + { + Func p("acc_p"), f("acc_f"); + p(x, y) = x + y; // stand-in for an expensive producer + f(x, y) = 100; + f(x, y) += branch(y < 8, p(x, y), x - y); + p.compute_at(f, x); // lands inside the branched update + + Module m = f.compile_to_module({}, "acc_f"); + if (count_real_branches(m) < 1) { + printf("f += branch did not lower to a real branch\n"); + return 1; + } + BranchArmProduce pg = arm_produce(m, "acc_p"); + if (!pg.found) { + printf("f += branch: no branch found around producer\n"); + return 1; + } + if (pg.then_count < 1 || pg.else_count != 0) { + printf("f += branch: producer not gated to the taken arm (then=%d else=%d)\n", + pg.then_count, pg.else_count); + return 1; + } + + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = 100 + ((j < 8) ? (i + j) : (i - j)); + if (r(i, j) != correct) { + printf("f += branch mismatch at (%d,%d): got %d want %d\n", + i, j, r(i, j), correct); + return 1; + } + } + } + } + + // Part AH-sub: f(x) -= branch(cond, a, b) distributes the same way - + // f = branch(cond, f - a, f - b) - real control flow with the producer gated + // to the taken arm. + { + Func p("sub_p"), f("sub_f"); + p(x, y) = x + y; + f(x, y) = 100; + f(x, y) -= branch(y < 8, p(x, y), x - y); + p.compute_at(f, x); + + Module m = f.compile_to_module({}, "sub_f"); + BranchArmProduce pg = arm_produce(m, "sub_p"); + if (count_real_branches(m) < 1 || !pg.found || + pg.then_count < 1 || pg.else_count != 0) { + printf("f -= branch: not a gated real branch (then=%d else=%d)\n", + pg.then_count, pg.else_count); + return 1; + } + Buffer r = f.realize({16, 16}); + for (int j = 0; j < 16; j++) { + for (int i = 0; i < 16; i++) { + int correct = 100 - ((j < 8) ? (i + j) : (i - j)); + if (r(i, j) != correct) { + printf("f -= branch mismatch at (%d,%d): got %d want %d\n", + i, j, r(i, j), correct); + return 1; + } + } + } + } + + // Part AH-mul: f(x) *= branch(cond, a, b) distributes to + // f = branch(cond, f * a, f * b). The only difference from += / -= is the + // identity: with no pure definition the Func initialises to 1, not 0. + { + // With an explicit pure definition. + Func f("mul_f"); + f(x) = 2; + f(x) *= branch(x < 5, x + 1, 3); + Module m = f.compile_to_module({}, "mul_f"); + if (count_real_branches(m) < 1) { + printf("f *= branch did not lower to a real branch\n"); + return 1; + } + Buffer r = f.realize({10}); + for (int i = 0; i < 10; i++) { + int correct = 2 * ((i < 5) ? (i + 1) : 3); + if (r(i) != correct) { + printf("f *= branch mismatch at %d: got %d want %d\n", i, r(i), correct); + return 1; + } + } + + // With no pure definition: initialises to 1 (multiplicative identity). + Func g("mul_g"); + g(x) *= branch(x < 5, x + 1, 3); + Buffer rg = g.realize({10}); + for (int i = 0; i < 10; i++) { + int correct = (i < 5) ? (i + 1) : 3; // 1 * taken arm + if (rg(i) != correct) { + printf("g *= branch (no pure def) mismatch at %d: got %d want %d\n", + i, rg(i), correct); + return 1; + } + } + } + + // Part AH-div: f(x) /= branch(cond, a, b) becomes f = branch(cond, f/a, f/b), + // init 1 with no pure definition (float, to avoid integer division). + { + Func f("div_f"); + f(x) = 120.0f; + f(x) /= branch(x < 5, cast(x + 1), 4.0f); + if (count_real_branches(f.compile_to_module({}, "div_f")) < 1) { + printf("f /= branch did not lower to a real branch\n"); + return 1; + } + Buffer r = f.realize({10}); + for (int i = 0; i < 10; i++) { + float correct = 120.0f / ((i < 5) ? (i + 1) : 4); + if (std::abs(r(i) - correct) > 1e-3f) { + printf("f /= branch mismatch at %d: got %f want %f\n", i, r(i), correct); + return 1; + } + } + } + + // Part AH-ad: reverse-mode (VJP) through a compound branch update. + // f += branch(cond,a,b) is a branch-valued update f = branch(cond, f+a, f+b), + // semantically identical to f += select(cond,a,b), so the gradient must match + // the select version. Covers the self-reference-in-both-arms case. (This + // branch is VJP-only; the JVP counterpart lives on the forward-diff branch.) + { + auto reduce_grad = [&](bool use_branch) -> float { + Param p("adp"); + p.set(3.0f); + RDom r(0, 10); + Func f(use_branch ? "adbv" : "adsv"); + f() = 0.0f; + if (use_branch) { + f() += branch(r < 5, p * cast(r), p * p); + } else { + f() += select(r < 5, p * cast(r), p * p); + } + Buffer d = propagate_adjoints(f)(p).realize(); + return d(); + }; + // f = sum_{r<5} p r + sum_{r>=5} p^2 = 10p + 5p^2; df/dp = 10 + 10p = 40 at p=3. + float vjp_s = reduce_grad(false), vjp_b = reduce_grad(true); + if (std::abs(vjp_b - vjp_s) > 1e-3f || std::abs(vjp_b - 40.0f) > 1e-3f) { + printf("VJP through += branch: select=%f branch=%f (want 40)\n", vjp_s, vjp_b); + return 1; + } + + // Product reduction (*=), where the derivative is the product rule. + // The branch operator only has to be faithful to the select version: + // Halide's own reverse-mode AD of a self-referencing product scan is + // unreliable (it does not match the analytic gradient) for BOTH select + // and branch, so we only assert branch == select. + auto prod_grad = [&](bool use_branch) -> float { + Param p("mdp"); + p.set(3.0f); + RDom r(0, 3); + Func f(use_branch ? "mdbv" : "mdsv"); + f() = 1.0f; + if (use_branch) { + f() *= branch(r < 2, p, 2.0f); + } else { + f() *= select(r < 2, p, 2.0f); + } + Buffer d = propagate_adjoints(f)(p).realize(); + return d(); + }; + float pv_s = prod_grad(false), pv_b = prod_grad(true); + if (std::abs(pv_b - pv_s) > 1e-3f) { + printf("VJP through *= branch not faithful: select=%f branch=%f\n", pv_s, pv_b); + return 1; + } + } + + // Part AH2: reverse-mode (VJP) through a branch with a NONLINEAR loss. The + // backward pass needs the forward value of the branch-valued Func, which used + // to bury the branch inside an adjoint expression and fail lowering with + // "a branch-defined Func can not be inlined". It now lifts back to real + // control flow, and the gradient matches both select and the analytic value. + { + auto grad = [&](bool use_branch) -> float { + Param p("pn"); + p.set(3.0f); + Var xx("xx"); + Expr cond = xx < 5; + Expr a = p * cast(xx); // da/dp = xx + Expr b = p * p; // db/dp = 2p + Func fwd(use_branch ? "vjpn_b" : "vjpn_s"); + if (use_branch) { + fwd(xx) = branch(cond, a, b); + } else { + fwd(xx) = select(cond, a, b); + } + fwd.compute_root(); + RDom r(0, 10); + Func loss(use_branch ? "lossn_b" : "lossn_s"); + loss() = sum(fwd(r) * fwd(r)); // nonlinear in fwd + loss.compute_root(); + Buffer g = propagate_adjoints(loss)(p).realize(); + return g(); + }; + float gs = grad(false), gb = grad(true); + // sum_{xx<5} 2 p xx^2 + sum_{xx>=5} 4 p^3 = 180 + 540 = 720 at p=3. + if (std::abs(gs - 720.0f) > 1e-2f || std::abs(gb - gs) > 1e-2f) { + printf("nonlinear VJP mismatch: select=%f branch=%f analytic=720\n", gs, gb); + return 1; + } + } + + // Part AI: a condition that varies across the lanes of a vectorized + // dimension can not be a scalar jump. Rather than erroring, it becomes + // predication: the loads and stores of each arm are masked. Only the values + // have to be right here; the masking itself is checked in Part AJ. + { + Func f("vec_branch"); + f(x) = branch(x < 5, x * 2, x + 1); + f.vectorize(x, 8); + + Buffer r = f.realize({64}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 5 ? (i * 2) : (i + 1); + if (r(i) != correct) { + printf("vectorized branch mismatch at %d: got %d want %d\n", + i, r(i), correct); + return 1; + } + } + } + + // Part AJ: the boundary-condition case. A branch on a vectorized dimension + // whose arm loads from an input becomes a PREDICATED (masked) vector load, + // so the out-of-bounds lanes never touch memory. unsafe_promise_clamped is + // what tells bounds inference the index stays in range, so `in` is not + // required outside its real extent. On AVX-512 this is a single masked load + // instead of a clamp plus a select. + { + const int w = 100; + Buffer in(w); + for (int i = 0; i < w; i++) { + in(i) = i * 3 + 1; + } + + Func f("pred_bc"); + f(x) = branch(x >= 0 && x < w, + in(unsafe_promise_clamped(x, 0, w - 1)), + 0); + f.vectorize(x, 16); + + Module m = f.compile_to_module({}, "pred_bc"); + HasPredicatedLoad hpl; + if (!scan_module(m, hpl)) { + printf("vectorized branch did not produce a predicated load\n"); + return 1; + } + + // Realize past the end of `in`: the guarded lanes must not load. + Buffer r = f.realize({128}); + for (int i = 0; i < r.width(); i++) { + int correct = (i < w) ? (i * 3 + 1) : 0; + if (r(i) != correct) { + printf("predicated boundary mismatch at %d: got %d want %d\n", + i, r(i), correct); + return 1; + } + } + } + +#ifdef HALIDE_WITH_EXCEPTIONS + // Part H: a lane-varying (vector) condition is rejected at construction. + { + bool threw = false; + try { + (void)branch(Broadcast::make(x > 5, 4), x * 2, x + 1); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch() with a vector condition should have thrown\n"); + return 1; + } + } + + // Part J: branch() combined with vectorization inside a GPU kernel errors. + { + Param p("p"); + Func h("gpu_vec_branch"); + Var xo("xo"), xi("xi"), xv("xv"); + h(x) = branch(p > 0, x * 2, x + 1); + h.split(x, xo, xi, 64) + .split(xi, xi, xv, 4) + .gpu_blocks(xo) + .gpu_threads(xi) + .vectorize(xv); + bool threw = false; + try { + h.compile_to_module({p}, "gpu_vec_branch", + get_host_target().with_feature(Target::CUDA)); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch() with vectorization in a GPU kernel should have thrown\n"); + return 1; + } + } + + // Part AE: a branch produces a single value, so it can not define a stage of + // a tuple-valued Func. That must be a clean error, not a crash. + { + bool threw = false; + try { + Func f("tuple_branch"); + f(x) = Tuple(0, 1); // two values + f(x) = branch(x < 5, 100, 200); // branch gives one -> mismatch + f.realize({10}); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch on a tuple-valued Func should have thrown\n"); + return 1; + } + } + + // Part K: a branch-defined Func can not be inlined (the branch must be the + // whole value of a definition), so using one inline is an error. + { + Func b("inlined_branch"), f("consumer"); + b(x) = branch(x < 5, x * 2, x + 1); + f(x) = b(x) + 1; // b is inlined by default -> branch gets buried + bool threw = false; + try { + f.compile_to_module({}, "consumer"); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("inlining a branch-defined Func should have thrown\n"); + return 1; + } + } +#endif + + printf("Success!\n"); + return 0; +}