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/Func.cpp b/src/Func.cpp index f4e4298fae07..dacdaaa43a27 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) { diff --git a/src/Func.h b/src/Func.h index 0bfb591871c7..308eb868cdc0 100644 --- a/src/Func.h +++ b/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 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..d58a3302aa4c 100644 --- a/src/IR.h +++ b/src/IR.h @@ -620,6 +620,12 @@ 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 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 diff --git a/src/IROperator.cpp b/src/IROperator.cpp index c109f1d2ccff..97fcc31385cb 100644 --- a/src/IROperator.cpp +++ b/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."; diff --git a/src/IROperator.h b/src/IROperator.h index 797f12870f5d..b8526d6b8a6e 100644 --- a/src/IROperator.h +++ b/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::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..c8aab67c9dd7 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,153 @@ 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); + } +} + +// A branch can only be real control flow if its condition is not lane-varying, +// and (on GPU) if the kernel is not vectorized. We check this against the +// schedule rather than silently falling back to a 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 must +// have a lane-varying condition (otherwise it would have been hoisted out of +// it). A real branch can not be taken per SIMD lane, and we deliberately do not +// fall back to a select, so this is a hard error. +class CheckBranchNotVectorized : public IRVisitor { + using IRVisitor::visit; + + const Function &func; + int in_vectorized = 0; + + void visit(const For *op) override { + const bool vec = (op->for_type == ForType::Vectorized); + in_vectorized += vec ? 1 : 0; + IRVisitor::visit(op); + in_vectorized -= vec ? 1 : 0; + } + + void visit(const IfThenElse *op) override { + if (op->else_case.defined() && is_branch_stmt(Stmt(op), func.name())) { + user_assert(in_vectorized == 0) + << "branch() in Func \"" << func.name() << "\" has a condition that " + << "depends on a vectorized dimension, so it can not become a real " + << "control-flow branch (a branch can not be taken per SIMD lane). " + << "Use select() for a lane-varying condition, or do not vectorize " + << "that dimension.\n"; + } + IRVisitor::visit(op); + } + +public: + explicit CheckBranchNotVectorized(const Function &f) + : func(f) { + } +}; + Stmt build_provide_loop_nest(const map &env, const string &prefix, const Function &func, @@ -494,8 +642,27 @@ 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()); + // 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 end up + // buried inside a value only if a branch-defined Func was inlined into a + // consumer, which we can not turn into 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 +684,16 @@ 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 and can not be a real branch. + CheckBranchNotVectorized check(func); + stmt.accept(&check); + } stmt = inject_placeholder_prefetch(stmt, env, prefix, def.schedule().prefetches()); // Make any specialized copies @@ -2606,6 +2783,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/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..915f73a27e6e --- /dev/null +++ b/test/correctness/branch.cpp @@ -0,0 +1,834 @@ +#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 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; + } + } + } + } + +#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 I: a condition that depends on a vectorized dimension is rejected. + { + Func f("vec_branch"); + f(x) = branch(x < 5, x * 2, x + 1); + f.vectorize(x, 8); + bool threw = false; + try { + f.compile_to_module({}, "vec_branch"); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch() on a vectorized dimension 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; +}