From 0121bcad3a3fa428cf7e49eb0dca1d44b676231e Mon Sep 17 00:00:00 2001 From: Soufiane KHIAT Date: Wed, 8 Jul 2026 11:25:51 +0200 Subject: [PATCH 1/4] Allow .branch() for select() --- src/CSE.cpp | 15 +++ src/Expr.h | 15 +++ src/IROperator.cpp | 28 +++++ test/correctness/CMakeLists.txt | 1 + test/correctness/branch.cpp | 210 ++++++++++++++++++++++++++++++++ 5 files changed, 269 insertions(+) create mode 100644 test/correctness/branch.cpp diff --git a/src/CSE.cpp b/src/CSE.cpp index 95f61f0214e9..25270fa083f1 100644 --- a/src/CSE.cpp +++ b/src/CSE.cpp @@ -169,6 +169,21 @@ 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->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/Expr.h b/src/Expr.h index c36a4a1a5a2b..0f7e95a8080e 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -340,6 +340,21 @@ struct Expr : public Internal::IRHandle { Type type() const { return get()->type; } + + /** If this Expr is the result of select(cond, a, b), return an + * equivalent expression that lowers to a real control-flow branch, + * evaluating only the taken side, instead of the branchless select + * that always evaluates both sides. This is useful when one side is + * much more expensive than the other. It is an error to call this on + * an Expr that is not directly a select(). + * + * Selects that are directly the true or false value of the select are + * converted too, so the multi-way form select(c0, v0, c1, v1, ...) + * becomes a full if/else-if/else chain of branches. Note that a + * per-lane branch cannot be expressed under vectorization or on GPUs, + * so in those contexts the branch degrades back to a select (both + * sides evaluated). */ + Expr branch() const; }; /** This lets you use an Expr as a key in a map of the form diff --git a/src/IROperator.cpp b/src/IROperator.cpp index c109f1d2ccff..fc99840615e5 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -1537,6 +1537,34 @@ Expr select(Expr condition, Expr true_value, Expr false_value) { return Select::make(std::move(condition), std::move(true_value), std::move(false_value)); } +namespace { +// Recursively rewrite a select into the if_then_else intrinsic, which lowers +// to a real control-flow branch that evaluates only one side. Selects that +// are directly the true/false value of a select are rewritten too, so the +// nested-select chain produced by the multi-way select(c0, v0, c1, v1, ...) +// becomes a full chain of branches. Non-select values are left untouched. +Expr rewrite_select_to_branch(const Expr &e) { + const Select *s = e.as() != nullptr) + << "Expr::branch() may only be called on the result of select(...). " + << "This expression is not a select:\n " + << *this << "\n"; + return rewrite_select_to_branch(*this); +} + 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/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..864ff9d35ad0 --- /dev/null +++ b/test/correctness/branch.cpp @@ -0,0 +1,210 @@ +#include "Halide.h" +#include +#include +#include + +using namespace Halide; +using namespace Halide::Internal; + +// Test select(cond, a, b).branch(), which turns the branchless Select +// (both sides evaluated) into the if_then_else intrinsic that lowers to a +// real control-flow branch evaluating only the taken side. + +namespace { + +// Count the number of Calls with the given name in an Expr (counts textual +// occurrences, so a subexpression reached through two parents counts twice). +class CountNamedCall : public IRVisitor { + using IRVisitor::visit; + void visit(const Call *op) override { + if (op->name == name) { + count++; + } + IRVisitor::visit(op); + } + +public: + std::string name; + int count = 0; + explicit CountNamedCall(std::string n) + : name(std::move(n)) { + } +}; + +int count_named_calls(const Expr &e, const std::string &name) { + CountNamedCall c(name); + e.accept(&c); + return c.count; +} + +// Walk a lowered body, counting a target named Call both in total and within +// the true arm of any scalar if_then_else. Used to prove the CSE barrier +// keeps arm-only subexpressions inside the branch. +class BranchBarrierCheck : public IRVisitor { + using IRVisitor::visit; + void visit(const Call *op) override { + if (op->name == target) { + total++; + } + if (op->is_intrinsic(Call::if_then_else) && op->args.size() == 3) { + found_ite = true; + in_true_arm += count_named_calls(op->args[1], target); + } + IRVisitor::visit(op); + } + +public: + std::string target; + bool found_ite = false; + int total = 0; + int in_true_arm = 0; + explicit BranchBarrierCheck(std::string t) + : target(std::move(t)) { + } +}; + +} // namespace + +int main(int argc, char **argv) { + Var x("x"); + + // Part A: the frontend IR shape. + { + Expr cond = x > 5; + Expr t = x * 2; + Expr f = x + 1; + + Expr s = select(cond, t, f); + const Select *sel = s.as()) { + printf("multi-way branch() should recurse into the tail\n"); + return 1; + } + const Call *tail = Call::as_intrinsic(c->args[2], {Call::if_then_else}); + if (!tail) { + printf("multi-way branch() tail should be an if_then_else\n"); + return 1; + } + + Func f("multi"); + f(x) = e; + 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 E: the CSE barrier keeps arm-only subexpressions inside the branch + // rather than hoisting them into a let that runs unconditionally. An + // "expensive" pure extern call is used twice in the true arm; ordinarily + // CSE would lift it into a let above the branch, but branch()'s barrier + // must keep it inside the arm. + { + Expr e = Call::make(Int(32), "expensive_fn", {x}, Call::PureExtern); + Expr arm = min(e, 5) + max(e, 7); // uses e twice, not foldable away + + Func f("barrier"); + f(x) = select(x > 3, arm, 0).branch(); + + Module m = f.compile_to_module({}, "barrier"); + + BranchBarrierCheck check("expensive_fn"); + for (const auto &fn : m.functions()) { + fn.body.accept(&check); + } + + if (!check.found_ite) { + printf("barrier test: no if_then_else survived lowering\n"); + return 1; + } + if (check.in_true_arm < 1) { + printf("barrier failed: expensive call was hoisted out of the arm\n"); + return 1; + } + // Every occurrence of the expensive call must be inside the arm; if + // any were hoisted into an enclosing let, total would exceed in_arm. + if (check.total != check.in_true_arm) { + printf("barrier failed: %d expensive calls total but only %d in the arm\n", + check.total, check.in_true_arm); + return 1; + } + } + + // Part D: calling branch() on something that isn't a select is an error. +#ifdef HALIDE_WITH_EXCEPTIONS + { + bool threw = false; + try { + Expr not_a_select = x + 1; + (void)not_a_select.branch(); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch() on a non-select should have thrown\n"); + return 1; + } + } +#endif + + printf("Success!\n"); + return 0; +} From eb7e67b3b2b7dd3b021c65a880093721cb4a2122 Mon Sep 17 00:00:00 2001 From: Soufiane KHIAT Date: Thu, 9 Jul 2026 23:07:59 +0200 Subject: [PATCH 2/4] Add branch() intrinsic for an explicit select that evaluates only one side --- src/Bounds.cpp | 4 +- src/CSE.cpp | 3 +- src/Expr.h | 15 ----- src/IR.cpp | 1 + src/IR.h | 6 ++ src/IROperator.cpp | 54 +++++++++-------- src/IROperator.h | 22 +++++++ src/Lower.cpp | 53 +++++++++++++++++ test/correctness/branch.cpp | 114 +++++++++++++++++++++--------------- 9 files changed, 182 insertions(+), 90 deletions(-) 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 25270fa083f1..30f255bb88c3 100644 --- a/src/CSE.cpp +++ b/src/CSE.cpp @@ -178,7 +178,8 @@ class ComputeUseCounts : public IRGraphVisitor { // 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->args.size() >= 2) { + 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); diff --git a/src/Expr.h b/src/Expr.h index 0f7e95a8080e..c36a4a1a5a2b 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -340,21 +340,6 @@ struct Expr : public Internal::IRHandle { Type type() const { return get()->type; } - - /** If this Expr is the result of select(cond, a, b), return an - * equivalent expression that lowers to a real control-flow branch, - * evaluating only the taken side, instead of the branchless select - * that always evaluates both sides. This is useful when one side is - * much more expensive than the other. It is an error to call this on - * an Expr that is not directly a select(). - * - * Selects that are directly the true or false value of the select are - * converted too, so the multi-way form select(c0, v0, c1, v1, ...) - * becomes a full if/else-if/else chain of branches. Note that a - * per-lane branch cannot be expressed under vectorization or on GPUs, - * so in those contexts the branch degrades back to a select (both - * sides evaluated). */ - Expr branch() const; }; /** This lets you use an Expr as a key in a map of the form 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..51ce1283ce3a 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 cannot be a real branch (a lane-varying + // condition, or use 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 fc99840615e5..f61048b7ddd3 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -1537,32 +1537,38 @@ Expr select(Expr condition, Expr true_value, Expr false_value) { return Select::make(std::move(condition), std::move(true_value), std::move(false_value)); } -namespace { -// Recursively rewrite a select into the if_then_else intrinsic, which lowers -// to a real control-flow branch that evaluates only one side. Selects that -// are directly the true/false value of a select are rewritten too, so the -// nested-select chain produced by the multi-way select(c0, v0, c1, v1, ...) -// becomes a full chain of branches. Non-select values are left untouched. -Expr rewrite_select_to_branch(const Expr &e) { - const Select *s = e.as() != nullptr) - << "Expr::branch() may only be called on the result of select(...). " - << "This expression is not a select:\n " - << *this << "\n"; - return rewrite_select_to_branch(*this); + 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(); + return Call::make(result_type, + Call::branch, + {std::move(condition), std::move(true_value), std::move(false_value)}, + Call::PureIntrinsic); } Tuple select(const Tuple &condition, const Tuple &true_value, const Tuple &false_value) { diff --git a/src/IROperator.h b/src/IROperator.h index 797f12870f5d..fc4d8f0fb985 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -851,6 +851,28 @@ inline Expr select(const Expr &c0, const FuncRef &v0, const Expr &c1, const Func } // @} +/** 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. This is useful when one + * side is much more expensive than the other, and, more importantly, it lets + * bounds inference assume that only one side runs. + * + * The condition must be a scalar bool. 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 inside a + * GPU kernel. */ +Expr 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 Expr 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)...)); +} + /** 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..21084265dba2 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -91,6 +91,55 @@ using std::vector; namespace { +// Lower branch() intrinsics into if_then_else, which codegen turns into a real +// control-flow branch that evaluates only the taken side. This runs after +// vectorization so we can tell whether the condition became lane-varying. We +// deliberately do NOT silently fall back to a both-sides select when a real +// branch is impossible (a lane-varying condition, or use inside a GPU kernel): +// the whole point of branch() is the guarantee that only one side runs, so we +// hard-error instead. +class LowerStrictBranches : public IRMutator { + using IRMutator::visit; + + int in_gpu_loop = 0; + + Stmt visit(const For *op) override { + const bool gpu = op->for_type == ForType::GPUBlock || + op->for_type == ForType::GPUThread || + op->for_type == ForType::GPULane; + in_gpu_loop += gpu ? 1 : 0; + Stmt s = IRMutator::visit(op); + in_gpu_loop -= gpu ? 1 : 0; + return s; + } + + Expr visit(const Call *op) override { + if (op->is_intrinsic(Call::branch)) { + internal_assert(op->args.size() == 3); + Expr cond = mutate(op->args[0]); + Expr true_value = mutate(op->args[1]); + Expr false_value = mutate(op->args[2]); + // A uniform (broadcast) condition is still a single real branch + // that picks between two whole vectors. + if (const Broadcast *b = cond.as()) { + cond = b->value; + } + user_assert(!cond.type().is_vector()) + << "branch() could not be lowered to a real control-flow branch " + << "because its condition became lane-varying (it depends on a " + << "vectorized dimension). Use select() for a per-lane condition, " + << "or schedule so the condition is not vectorized.\n"; + user_assert(in_gpu_loop == 0) + << "branch() is not supported inside a GPU kernel, where a real " + << "control-flow branch can not guarantee that only one side is " + << "evaluated. Use select() instead.\n"; + return Call::make(op->type, Call::if_then_else, + {cond, true_value, false_value}, Call::PureIntrinsic); + } + return IRMutator::visit(op); + } +}; + class LoweringLogger { Stmt last_written; std::chrono::time_point last_time; @@ -354,6 +403,10 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after vectorizing:", s); + debug(1) << "Lowering strict branches...\n"; + s = LowerStrictBranches()(s); + log("Lowering after lowering strict branches:", s); + if (t.has_gpu_feature() || t.has_feature(Target::Vulkan)) { debug(1) << "Injecting per-block gpu synchronization...\n"; diff --git a/test/correctness/branch.cpp b/test/correctness/branch.cpp index 864ff9d35ad0..229ee6c3fa95 100644 --- a/test/correctness/branch.cpp +++ b/test/correctness/branch.cpp @@ -6,14 +6,13 @@ using namespace Halide; using namespace Halide::Internal; -// Test select(cond, a, b).branch(), which turns the branchless Select -// (both sides evaluated) into the if_then_else intrinsic that lowers to a -// real control-flow branch evaluating only the taken side. +// Test branch(cond, a, b): a strict variant of select that lowers to a real +// control-flow branch evaluating only the taken side, instead of the branchless +// select that always evaluates both sides. namespace { -// Count the number of Calls with the given name in an Expr (counts textual -// occurrences, so a subexpression reached through two parents counts twice). +// Count Calls with the given name in an Expr (textual occurrences). class CountNamedCall : public IRVisitor { using IRVisitor::visit; void visit(const Call *op) override { @@ -37,9 +36,9 @@ int count_named_calls(const Expr &e, const std::string &name) { return c.count; } -// Walk a lowered body, counting a target named Call both in total and within -// the true arm of any scalar if_then_else. Used to prove the CSE barrier -// keeps arm-only subexpressions inside the branch. +// Walk a lowered body, counting a target named Call in total and within the +// true arm of any scalar if_then_else (branch lowers to if_then_else). Used to +// prove the CSE barrier keeps arm-only subexpressions inside the branch. class BranchBarrierCheck : public IRVisitor { using IRVisitor::visit; void visit(const Call *op) override { @@ -68,34 +67,27 @@ class BranchBarrierCheck : public IRVisitor { int main(int argc, char **argv) { Var x("x"); - // Part A: the frontend IR shape. + // Part A: the frontend produces the branch intrinsic and keeps operands. { Expr cond = x > 5; Expr t = x * 2; Expr f = x + 1; - Expr s = select(cond, t, f); - const Select *sel = s.as()) { - printf("multi-way branch() should recurse into the tail\n"); - return 1; - } - const Call *tail = Call::as_intrinsic(c->args[2], {Call::if_then_else}); + const Call *tail = Call::as_intrinsic(c->args[2], {Call::branch}); if (!tail) { - printf("multi-way branch() tail should be an if_then_else\n"); + printf("multi-way branch() tail should itself be a branch\n"); return 1; } @@ -152,17 +138,14 @@ int main(int argc, char **argv) { } } - // Part E: the CSE barrier keeps arm-only subexpressions inside the branch - // rather than hoisting them into a let that runs unconditionally. An - // "expensive" pure extern call is used twice in the true arm; ordinarily - // CSE would lift it into a let above the branch, but branch()'s barrier - // must keep it inside the arm. + // Part D: the CSE barrier keeps arm-only subexpressions inside the branch + // rather than hoisting them into a let that runs unconditionally. { Expr e = Call::make(Int(32), "expensive_fn", {x}, Call::PureExtern); Expr arm = min(e, 5) + max(e, 7); // uses e twice, not foldable away Func f("barrier"); - f(x) = select(x > 3, arm, 0).branch(); + f(x) = branch(x > 3, arm, 0); Module m = f.compile_to_module({}, "barrier"); @@ -172,15 +155,13 @@ int main(int argc, char **argv) { } if (!check.found_ite) { - printf("barrier test: no if_then_else survived lowering\n"); + printf("barrier test: branch did not lower to an if_then_else\n"); return 1; } if (check.in_true_arm < 1) { printf("barrier failed: expensive call was hoisted out of the arm\n"); return 1; } - // Every occurrence of the expensive call must be inside the arm; if - // any were hoisted into an enclosing let, total would exceed in_arm. if (check.total != check.in_true_arm) { printf("barrier failed: %d expensive calls total but only %d in the arm\n", check.total, check.in_true_arm); @@ -188,18 +169,55 @@ int main(int argc, char **argv) { } } - // Part D: calling branch() on something that isn't a select is an error. #ifdef HALIDE_WITH_EXCEPTIONS + // Part E: branch() with a lane-varying (vector) condition is a hard error + // at construction (scalar condition only). + { + bool threw = false; + try { + Expr vec_cond = Broadcast::make(x > 5, 4); + (void)branch(vec_cond, x * 2, x + 1); + } catch (const CompileError &) { + threw = true; + } + if (!threw) { + printf("branch() with a vector condition should have thrown\n"); + return 1; + } + } + + // Part F: a branch whose condition depends on a vectorized dimension is a + // hard error at lowering (no silent fallback to select). + { + 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("vectorized branch() should have thrown at lowering\n"); + return 1; + } + } + + // Part G: a branch inside a GPU kernel is a hard error at lowering. { + Func g("gpu_branch"); + g(x) = branch(x < 5, x * 2, x + 1); + Var xo("xo"), xi("xi"); + g.gpu_tile(x, xo, xi, 16); bool threw = false; try { - Expr not_a_select = x + 1; - (void)not_a_select.branch(); + g.compile_to_module({}, "gpu_branch", get_host_target().with_feature(Target::CUDA)); } catch (const CompileError &) { threw = true; } if (!threw) { - printf("branch() on a non-select should have thrown\n"); + printf("branch() inside a GPU kernel should have thrown at lowering\n"); return 1; } } From 9bd2011b6ea8173d0ab444cee6f0266ec70742e2 Mon Sep 17 00:00:00 2001 From: Soufiane KHIAT Date: Mon, 13 Jul 2026 21:26:41 +0200 Subject: [PATCH 3/4] Move feature to FuncRef and force inline if possible otherwise error --- src/Func.cpp | 53 ++++++++++++ src/Func.h | 6 ++ src/IR.h | 4 +- src/IROperator.cpp | 11 +-- src/IROperator.h | 44 +++++++--- src/Lower.cpp | 24 +++--- test/correctness/branch.cpp | 159 +++++++++++++++++++++++++++++++----- 7 files changed, 252 insertions(+), 49 deletions(-) diff --git a/src/Func.cpp b/src/Func.cpp index f4e4298fae07..22a3708a7639 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,58 @@ Stage FuncRef::operator=(const Expr &e) { return (*this) = Tuple(e); } +namespace { + +// Collect the Halide Funcs called in the value arms of a branch (the true and +// false values, recursing through the nested branches of a multi-way branch). +// Conditions are intentionally skipped: they are always evaluated, so they do +// not need to be inlined for the branch to gate computation. +class CollectValueArmFuncs : public IRGraphVisitor { + using IRGraphVisitor::visit; + using IRGraphVisitor::include; + + void visit(const Call *op) override { + if (op->is_intrinsic(Call::branch) && op->args.size() == 3) { + // Recurse into the value arms (args 1 and 2) only; skip the + // condition (arg 0). + include(op->args[1]); + include(op->args[2]); + return; + } + if (op->call_type == Call::Halide && op->func.defined()) { + funcs.emplace_back(op->func); + } + IRGraphVisitor::visit(op); + } + +public: + std::vector funcs; +}; + +// Force-inline the value arms of a branch so that only the taken side is +// actually computed (not just loaded). Hard-error if a value arm calls a Func +// that can not be inlined. +void force_inline_branch_arms(const Expr &branch_expr) { + CollectValueArmFuncs collector; + branch_expr.accept(&collector); + for (Function &fn : collector.funcs) { + user_assert(!fn.has_update_definition()) + << "branch() value arm calls Func \"" << fn.name() << "\", which has an " + << "update definition and can not be inlined. branch() force-inlines its " + << "value arms so that only the taken side is computed; a Func with an " + << "update stage can not be inlined. Rewrite \"" << fn.name() + << "\" without an update stage, or use select() there instead.\n"; + Func(fn).compute_inline(); + } +} + +} // namespace + +Stage FuncRef::operator=(const Branch &b) { + force_inline_branch_arms(b.expr); + 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.h b/src/IR.h index 51ce1283ce3a..d58a3302aa4c 100644 --- a/src/IR.h +++ b/src/IR.h @@ -623,8 +623,8 @@ struct Call : public ExprNode { // 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 cannot be a real branch (a lane-varying - // condition, or use inside a GPU kernel). + // 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, diff --git a/src/IROperator.cpp b/src/IROperator.cpp index f61048b7ddd3..97fcc31385cb 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -1537,7 +1537,7 @@ Expr select(Expr condition, Expr true_value, Expr false_value) { return Select::make(std::move(condition), std::move(true_value), std::move(false_value)); } -Expr branch(Expr condition, Expr true_value, Expr 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" @@ -1565,10 +1565,11 @@ Expr branch(Expr condition, Expr true_value, Expr false_value) { // 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(); - return Call::make(result_type, - Call::branch, - {std::move(condition), std::move(true_value), std::move(false_value)}, - Call::PureIntrinsic); + 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) { diff --git a/src/IROperator.h b/src/IROperator.h index fc4d8f0fb985..b8526d6b8a6e 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -851,26 +851,46 @@ inline Expr select(const Expr &c0, const FuncRef &v0, const Expr &c1, const Func } // @} -/** 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 +/** 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. This is useful when one - * side is much more expensive than the other, and, more importantly, it lets - * bounds inference assume that only one side runs. + * 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. 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 inside a - * GPU kernel. */ -Expr branch(Expr condition, Expr true_value, Expr false_value); + * 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 Expr 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)...)); +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 diff --git a/src/Lower.cpp b/src/Lower.cpp index 21084265dba2..77228b458252 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -92,12 +92,16 @@ using std::vector; namespace { // Lower branch() intrinsics into if_then_else, which codegen turns into a real -// control-flow branch that evaluates only the taken side. This runs after -// vectorization so we can tell whether the condition became lane-varying. We -// deliberately do NOT silently fall back to a both-sides select when a real -// branch is impossible (a lane-varying condition, or use inside a GPU kernel): -// the whole point of branch() is the guarantee that only one side runs, so we -// hard-error instead. +// control-flow branch that evaluates only the taken side (on CPU and GPU +// alike). This runs after vectorization so we can tell whether the branch got +// vectorized. A real branch is impossible in two cases, and there we +// deliberately do NOT silently fall back to a both-sides select (the whole +// point of branch() is the guarantee that only one side runs), so we +// hard-error instead: +// - a lane-varying (vectorized) condition, since SIMD lanes can not branch +// independently (CPU or GPU); and +// - any vectorization inside a GPU kernel: on GPU only a per-thread or +// per-wave (non-vectorized) branch makes sense. class LowerStrictBranches : public IRMutator { using IRMutator::visit; @@ -129,10 +133,10 @@ class LowerStrictBranches : public IRMutator { << "because its condition became lane-varying (it depends on a " << "vectorized dimension). Use select() for a per-lane condition, " << "or schedule so the condition is not vectorized.\n"; - user_assert(in_gpu_loop == 0) - << "branch() is not supported inside a GPU kernel, where a real " - << "control-flow branch can not guarantee that only one side is " - << "evaluated. Use select() instead.\n"; + user_assert(!(in_gpu_loop > 0 && op->type.is_vector())) + << "branch() can not be combined with vectorization inside a GPU " + << "kernel. On GPU only a per-thread or per-wave (non-vectorized) " + << "branch is allowed. Use select() for the vectorized dimension.\n"; return Call::make(op->type, Call::if_then_else, {cond, true_value, false_value}, Call::PureIntrinsic); } diff --git a/test/correctness/branch.cpp b/test/correctness/branch.cpp index 229ee6c3fa95..483073941b2d 100644 --- a/test/correctness/branch.cpp +++ b/test/correctness/branch.cpp @@ -7,8 +7,10 @@ using namespace Halide; using namespace Halide::Internal; // Test branch(cond, a, b): a strict variant of select that lowers to a real -// control-flow branch evaluating only the taken side, instead of the branchless -// select that always evaluates both sides. +// control-flow branch evaluating only the taken side. Unlike select(), branch() +// returns a Branch (not an Expr) and may only be used as the whole right-hand +// side of a Func definition. Its value arms are force-inlined so the branch +// gates real computation, not just a load. namespace { @@ -62,19 +64,47 @@ class BranchBarrierCheck : public IRVisitor { } }; +// True if the lowered body allocates a buffer whose name starts with `prefix` +// (i.e. the corresponding Func was materialized rather than inlined). +class HasAllocationPrefixed : public IRVisitor { + using IRVisitor::visit; + void visit(const Allocate *op) override { + if (op->name.rfind(prefix, 0) == 0) { + found = true; + } + IRVisitor::visit(op); + } + +public: + std::string prefix; + bool found = false; + explicit HasAllocationPrefixed(std::string p) + : prefix(std::move(p)) { + } +}; + +bool module_allocates(const Module &m, const std::string &prefix) { + HasAllocationPrefixed v(prefix); + for (const auto &fn : m.functions()) { + fn.body.accept(&v); + } + return v.found; +} + } // namespace int main(int argc, char **argv) { Var x("x"); - // Part A: the frontend produces the branch intrinsic and keeps operands. + // Part A: branch() returns a Branch wrapping the branch intrinsic, with the + // operands preserved. { Expr cond = x > 5; Expr t = x * 2; Expr f = x + 1; - Expr b = branch(cond, t, f); - const Call *c = Call::as_intrinsic(b, {Call::branch}); + Branch b = branch(cond, t, f); + const Call *c = Call::as_intrinsic(b.expr, {Call::branch}); if (!c) { printf("branch() did not produce a branch intrinsic\n"); return 1; @@ -86,10 +116,6 @@ int main(int argc, char **argv) { printf("branch() did not preserve its operands\n"); return 1; } - if (b.type() != t.type()) { - printf("branch() has the wrong result type\n"); - return 1; - } } // Part B: branch() is numerically equivalent to select() end-to-end. @@ -114,8 +140,8 @@ int main(int argc, char **argv) { // Part C: the multi-way branch becomes a chain of branch intrinsics. { - Expr e = branch(x < 10, 1, x < 20, 2, 3); - const Call *c = Call::as_intrinsic(e, {Call::branch}); + Branch e = branch(x < 10, 1, x < 20, 2, 3); + const Call *c = Call::as_intrinsic(e.expr, {Call::branch}); if (!c) { printf("multi-way branch() did not produce a branch intrinsic\n"); return 1; @@ -169,8 +195,74 @@ int main(int argc, char **argv) { } } + // Part E: the value arms are force-inlined, so a branch over Func calls + // does NOT materialize the producers - only the taken side is computed. + { + Func arm_g("arm_g"), arm_h("arm_h"), f("force_inline"); + arm_g(x) = x * x + 7; + arm_h(x) = x * 5 - 2; + f(x) = branch(x < 128, arm_g(x), arm_h(x)); + + Module m = f.compile_to_module({}, "force_inline"); + if (module_allocates(m, "arm_g") || module_allocates(m, "arm_h")) { + printf("force-inline failed: an arm Func was materialized, not inlined\n"); + return 1; + } + + Buffer r = f.realize({256}); + for (int i = 0; i < r.width(); i++) { + int correct = i < 128 ? (i * i + 7) : (i * 5 - 2); + if (r(i) != correct) { + printf("force-inline mismatch at %d: got %d want %d\n", i, r(i), correct); + return 1; + } + } + } + + // Part F: force-inline overrides an explicit compute_root on an arm Func. + { + Func arm_g("fi_g"), arm_h("fi_h"), f("force_override"); + arm_g(x) = x * 3; + arm_h(x) = x + 9; + arm_g.compute_root(); // branch() must override this + f(x) = branch(x < 64, arm_g(x), arm_h(x)); + + Module m = f.compile_to_module({}, "force_override"); + if (module_allocates(m, "fi_g")) { + printf("force-inline did not override compute_root on an arm Func\n"); + return 1; + } + } + + // Part G: branch() is allowed inside a GPU kernel, where it becomes a real + // per-thread branch. Runs on the device when one is available (e.g. under a + // d3d12compute jit target), comparing the result against select(). + { + 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"); + } + } + #ifdef HALIDE_WITH_EXCEPTIONS - // Part E: branch() with a lane-varying (vector) condition is a hard error + // Part H: branch() with a lane-varying (vector) condition is a hard error // at construction (scalar condition only). { bool threw = false; @@ -186,7 +278,7 @@ int main(int argc, char **argv) { } } - // Part F: a branch whose condition depends on a vectorized dimension is a + // Part I: a branch whose condition depends on a vectorized dimension is a // hard error at lowering (no silent fallback to select). { Func f("vec_branch"); @@ -204,20 +296,47 @@ int main(int argc, char **argv) { } } - // Part G: a branch inside a GPU kernel is a hard error at lowering. + // Part J: branch() combined with vectorization inside a GPU kernel is a + // hard error (only per-thread / per-wave non-vectorized branches allowed). + { + 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 inside a GPU kernel should have thrown\n"); + return 1; + } + } + + // Part K: a value arm that can not be inlined (a Func with an update stage) + // is a hard error at definition time. { - Func g("gpu_branch"); - g(x) = branch(x < 5, x * 2, x + 1); - Var xo("xo"), xi("xi"); - g.gpu_tile(x, xo, xi, 16); bool threw = false; try { - g.compile_to_module({}, "gpu_branch", get_host_target().with_feature(Target::CUDA)); + Func g("upd"); + g(x) = 0; + RDom r(0, 4); + g(x) += r; // update stage -> not inlinable + Func f("uses_update"); + f(x) = branch(x < 5, g(x), 0); } catch (const CompileError &) { threw = true; } if (!threw) { - printf("branch() inside a GPU kernel should have thrown at lowering\n"); + printf("branch() over a Func with an update stage should have thrown\n"); return 1; } } From 399fcf1aa27dd3b971febaa71b83209cf39b80e0 Mon Sep 17 00:00:00 2001 From: Soufiane KHIAT Date: Thu, 16 Jul 2026 01:21:07 +0200 Subject: [PATCH 4/4] Update --- src/Func.cpp | 54 +-- src/Lower.cpp | 56 --- src/ScheduleFunctions.cpp | 191 ++++++++- test/correctness/branch.cpp | 769 +++++++++++++++++++++++++++++------- 4 files changed, 823 insertions(+), 247 deletions(-) diff --git a/src/Func.cpp b/src/Func.cpp index 22a3708a7639..dacdaaa43a27 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -3205,55 +3205,13 @@ Stage FuncRef::operator=(const Expr &e) { return (*this) = Tuple(e); } -namespace { - -// Collect the Halide Funcs called in the value arms of a branch (the true and -// false values, recursing through the nested branches of a multi-way branch). -// Conditions are intentionally skipped: they are always evaluated, so they do -// not need to be inlined for the branch to gate computation. -class CollectValueArmFuncs : public IRGraphVisitor { - using IRGraphVisitor::visit; - using IRGraphVisitor::include; - - void visit(const Call *op) override { - if (op->is_intrinsic(Call::branch) && op->args.size() == 3) { - // Recurse into the value arms (args 1 and 2) only; skip the - // condition (arg 0). - include(op->args[1]); - include(op->args[2]); - return; - } - if (op->call_type == Call::Halide && op->func.defined()) { - funcs.emplace_back(op->func); - } - IRGraphVisitor::visit(op); - } - -public: - std::vector funcs; -}; - -// Force-inline the value arms of a branch so that only the taken side is -// actually computed (not just loaded). Hard-error if a value arm calls a Func -// that can not be inlined. -void force_inline_branch_arms(const Expr &branch_expr) { - CollectValueArmFuncs collector; - branch_expr.accept(&collector); - for (Function &fn : collector.funcs) { - user_assert(!fn.has_update_definition()) - << "branch() value arm calls Func \"" << fn.name() << "\", which has an " - << "update definition and can not be inlined. branch() force-inlines its " - << "value arms so that only the taken side is computed; a Func with an " - << "update stage can not be inlined. Rewrite \"" << fn.name() - << "\" without an update stage, or use select() there instead.\n"; - Func(fn).compute_inline(); - } -} - -} // namespace - Stage FuncRef::operator=(const Branch &b) { - force_inline_branch_arms(b.expr); + // 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; } diff --git a/src/Lower.cpp b/src/Lower.cpp index 77228b458252..21b2c48a586f 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -91,59 +91,6 @@ using std::vector; namespace { -// Lower branch() intrinsics into if_then_else, which codegen turns into a real -// control-flow branch that evaluates only the taken side (on CPU and GPU -// alike). This runs after vectorization so we can tell whether the branch got -// vectorized. A real branch is impossible in two cases, and there we -// deliberately do NOT silently fall back to a both-sides select (the whole -// point of branch() is the guarantee that only one side runs), so we -// hard-error instead: -// - a lane-varying (vectorized) condition, since SIMD lanes can not branch -// independently (CPU or GPU); and -// - any vectorization inside a GPU kernel: on GPU only a per-thread or -// per-wave (non-vectorized) branch makes sense. -class LowerStrictBranches : public IRMutator { - using IRMutator::visit; - - int in_gpu_loop = 0; - - Stmt visit(const For *op) override { - const bool gpu = op->for_type == ForType::GPUBlock || - op->for_type == ForType::GPUThread || - op->for_type == ForType::GPULane; - in_gpu_loop += gpu ? 1 : 0; - Stmt s = IRMutator::visit(op); - in_gpu_loop -= gpu ? 1 : 0; - return s; - } - - Expr visit(const Call *op) override { - if (op->is_intrinsic(Call::branch)) { - internal_assert(op->args.size() == 3); - Expr cond = mutate(op->args[0]); - Expr true_value = mutate(op->args[1]); - Expr false_value = mutate(op->args[2]); - // A uniform (broadcast) condition is still a single real branch - // that picks between two whole vectors. - if (const Broadcast *b = cond.as()) { - cond = b->value; - } - user_assert(!cond.type().is_vector()) - << "branch() could not be lowered to a real control-flow branch " - << "because its condition became lane-varying (it depends on a " - << "vectorized dimension). Use select() for a per-lane condition, " - << "or schedule so the condition is not vectorized.\n"; - user_assert(!(in_gpu_loop > 0 && op->type.is_vector())) - << "branch() can not be combined with vectorization inside a GPU " - << "kernel. On GPU only a per-thread or per-wave (non-vectorized) " - << "branch is allowed. Use select() for the vectorized dimension.\n"; - return Call::make(op->type, Call::if_then_else, - {cond, true_value, false_value}, Call::PureIntrinsic); - } - return IRMutator::visit(op); - } -}; - class LoweringLogger { Stmt last_written; std::chrono::time_point last_time; @@ -407,9 +354,6 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after vectorizing:", s); - debug(1) << "Lowering strict branches...\n"; - s = LowerStrictBranches()(s); - log("Lowering after lowering strict branches:", s); if (t.has_gpu_feature() || t.has_feature(Target::Vulkan)) { 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/branch.cpp b/test/correctness/branch.cpp index 483073941b2d..915f73a27e6e 100644 --- a/test/correctness/branch.cpp +++ b/test/correctness/branch.cpp @@ -1,4 +1,5 @@ #include "Halide.h" +#include #include #include #include @@ -6,98 +7,206 @@ using namespace Halide; using namespace Halide::Internal; -// Test branch(cond, a, b): a strict variant of select that lowers to a real -// control-flow branch evaluating only the taken side. Unlike select(), branch() -// returns a Branch (not an Expr) and may only be used as the whole right-hand -// side of a Func definition. Its value arms are force-inlined so the branch -// gates real computation, not just a load. +// 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 { -// Count Calls with the given name in an Expr (textual occurrences). -class CountNamedCall : public IRVisitor { +// Is there a Stmt-level IfThenElse (with a real else) anywhere in the body? +class HasRealBranch : public IRVisitor { using IRVisitor::visit; - void visit(const Call *op) override { - if (op->name == name) { - count++; + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + found = true; } IRVisitor::visit(op); } public: - std::string name; - int count = 0; - explicit CountNamedCall(std::string n) - : name(std::move(n)) { + 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)) { } }; -int count_named_calls(const Expr &e, const std::string &name) { - CountNamedCall c(name); - e.accept(&c); - return c.count; +template +bool scan_module(const Module &m, V &v) { + for (const auto &fn : m.functions()) { + fn.body.accept(&v); + } + return v.found; } -// Walk a lowered body, counting a target named Call in total and within the -// true arm of any scalar if_then_else (branch lowers to if_then_else). Used to -// prove the CSE barrier keeps arm-only subexpressions inside the branch. -class BranchBarrierCheck : public IRVisitor { +// Count Stmt-level IfThenElse nodes that have a real else (i.e. branches). +class CountRealBranches : public IRVisitor { using IRVisitor::visit; - void visit(const Call *op) override { - if (op->name == target) { - total++; + void visit(const IfThenElse *op) override { + if (op->else_case.defined()) { + count++; } - if (op->is_intrinsic(Call::if_then_else) && op->args.size() == 3) { - found_ite = true; - in_true_arm += count_named_calls(op->args[1], target); + 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; - bool found_ite = false; - int total = 0; - int in_true_arm = 0; - explicit BranchBarrierCheck(std::string t) + int count = 0; + explicit CountProduce(std::string t) : target(std::move(t)) { } }; -// True if the lowered body allocates a buffer whose name starts with `prefix` -// (i.e. the corresponding Func was materialized rather than inlined). -class HasAllocationPrefixed : public IRVisitor { +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 Allocate *op) override { - if (op->name.rfind(prefix, 0) == 0) { + 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 prefix; + std::string target; bool found = false; - explicit HasAllocationPrefixed(std::string p) - : prefix(std::move(p)) { + int then_count = 0, else_count = 0; + explicit BranchArmProduce(std::string t) + : target(std::move(t)) { } }; -bool module_allocates(const Module &m, const std::string &prefix) { - HasAllocationPrefixed v(prefix); +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.found; + 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"); + Var x("x"), y("y"); - // Part A: branch() returns a Branch wrapping the branch intrinsic, with the - // operands preserved. + // Part A: branch() returns a Branch wrapping the branch intrinsic. { Expr cond = x > 5; Expr t = x * 2; @@ -105,20 +214,14 @@ int main(int argc, char **argv) { Branch b = branch(cond, t, f); const Call *c = Call::as_intrinsic(b.expr, {Call::branch}); - if (!c) { - printf("branch() did not produce a branch intrinsic\n"); - return 1; - } - if (c->args.size() != 3 || - !equal(c->args[0], cond) || - !equal(c->args[1], t) || - !equal(c->args[2], f)) { - printf("branch() did not preserve its operands\n"); + 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() end-to-end. + // Part B: branch() is numerically equivalent to select(). { Expr cond = (x % 3) == 0; Expr t = x * x + 7; @@ -138,22 +241,10 @@ int main(int argc, char **argv) { } } - // Part C: the multi-way branch becomes a chain of branch intrinsics. + // Part C: the multi-way branch produces correct values. { - Branch e = branch(x < 10, 1, x < 20, 2, 3); - const Call *c = Call::as_intrinsic(e.expr, {Call::branch}); - if (!c) { - printf("multi-way branch() did not produce a branch intrinsic\n"); - return 1; - } - const Call *tail = Call::as_intrinsic(c->args[2], {Call::branch}); - if (!tail) { - printf("multi-way branch() tail should itself be a branch\n"); - return 1; - } - Func f("multi"); - f(x) = e; + 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); @@ -164,79 +255,68 @@ int main(int argc, char **argv) { } } - // Part D: the CSE barrier keeps arm-only subexpressions inside the branch - // rather than hoisting them into a let that runs unconditionally. + // 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. { - Expr e = Call::make(Int(32), "expensive_fn", {x}, Call::PureExtern); - Expr arm = min(e, 5) + max(e, 7); // uses e twice, not foldable away - - Func f("barrier"); - f(x) = branch(x > 3, arm, 0); - - Module m = f.compile_to_module({}, "barrier"); + Func f("real_cf"); + f(x) = branch(x > 3, x * 2, x + 1); + Module m = f.compile_to_module({}, "real_cf"); - BranchBarrierCheck check("expensive_fn"); - for (const auto &fn : m.functions()) { - fn.body.accept(&check); - } - - if (!check.found_ite) { - printf("barrier test: branch did not lower to an if_then_else\n"); - return 1; - } - if (check.in_true_arm < 1) { - printf("barrier failed: expensive call was hoisted out of the arm\n"); - return 1; - } - if (check.total != check.in_true_arm) { - printf("barrier failed: %d expensive calls total but only %d in the arm\n", - check.total, check.in_true_arm); + HasRealBranch hb; + if (!scan_module(m, hb)) { + printf("branch did not become a Stmt-level IfThenElse\n"); return 1; } } - // Part E: the value arms are force-inlined, so a branch over Func calls - // does NOT materialize the producers - only the taken side is computed. + // 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 arm_g("arm_g"), arm_h("arm_h"), f("force_inline"); - arm_g(x) = x * x + 7; - arm_h(x) = x * 5 - 2; - f(x) = branch(x < 128, arm_g(x), arm_h(x)); - - Module m = f.compile_to_module({}, "force_inline"); - if (module_allocates(m, "arm_g") || module_allocates(m, "arm_h")) { - printf("force-inline failed: an arm Func was materialized, not inlined\n"); - return 1; - } + 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 * 5 - 2); + int correct = i < 128 ? (i * i + 7) : (i + 1); if (r(i) != correct) { - printf("force-inline mismatch at %d: got %d want %d\n", i, r(i), correct); + printf("producer mismatch at %d: got %d want %d\n", i, r(i), correct); return 1; } } } - // Part F: force-inline overrides an explicit compute_root on an arm Func. + // 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 arm_g("fi_g"), arm_h("fi_h"), f("force_override"); - arm_g(x) = x * 3; - arm_h(x) = x + 9; - arm_g.compute_root(); // branch() must override this - f(x) = branch(x < 64, arm_g(x), arm_h(x)); - - Module m = f.compile_to_module({}, "force_override"); - if (module_allocates(m, "fi_g")) { - printf("force-inline did not override compute_root on an arm Func\n"); + 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() is allowed inside a GPU kernel, where it becomes a real - // per-thread branch. Runs on the device when one is available (e.g. under a - // d3d12compute jit target), comparing the result against select(). + // 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()) { @@ -261,14 +341,407 @@ int main(int argc, char **argv) { } } + // 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: branch() with a lane-varying (vector) condition is a hard error - // at construction (scalar condition only). + // Part H: a lane-varying (vector) condition is rejected at construction. { bool threw = false; try { - Expr vec_cond = Broadcast::make(x > 5, 4); - (void)branch(vec_cond, x * 2, x + 1); + (void)branch(Broadcast::make(x > 5, 4), x * 2, x + 1); } catch (const CompileError &) { threw = true; } @@ -278,8 +751,7 @@ int main(int argc, char **argv) { } } - // Part I: a branch whose condition depends on a vectorized dimension is a - // hard error at lowering (no silent fallback to select). + // 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); @@ -291,13 +763,12 @@ int main(int argc, char **argv) { threw = true; } if (!threw) { - printf("vectorized branch() should have thrown at lowering\n"); + printf("branch() on a vectorized dimension should have thrown\n"); return 1; } } - // Part J: branch() combined with vectorization inside a GPU kernel is a - // hard error (only per-thread / per-wave non-vectorized branches allowed). + // Part J: branch() combined with vectorization inside a GPU kernel errors. { Param p("p"); Func h("gpu_vec_branch"); @@ -316,27 +787,43 @@ int main(int argc, char **argv) { threw = true; } if (!threw) { - printf("branch() with vectorization inside a GPU kernel should have thrown\n"); + 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 value arm that can not be inlined (a Func with an update stage) - // is a hard error at definition time. + // 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 { - Func g("upd"); - g(x) = 0; - RDom r(0, 4); - g(x) += r; // update stage -> not inlinable - Func f("uses_update"); - f(x) = branch(x < 5, g(x), 0); + f.compile_to_module({}, "consumer"); } catch (const CompileError &) { threw = true; } if (!threw) { - printf("branch() over a Func with an update stage should have thrown\n"); + printf("inlining a branch-defined Func should have thrown\n"); return 1; } }