From fe05aaa054b8816630c067fa3ecb2f9d002517fe Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 17 Aug 2026 20:51:50 +0200 Subject: [PATCH 01/29] PoC: Aligned split directive. --- src/ApplySplit.cpp | 82 +++++++++++++++-------- src/Func.cpp | 38 +++++++++-- src/Func.h | 5 ++ src/Schedule.h | 1 + test/correctness/CMakeLists.txt | 2 + test/correctness/split_aligned.cpp | 95 +++++++++++++++++++++++++++ test/correctness/split_aligned_2d.cpp | 90 +++++++++++++++++++++++++ 7 files changed, 280 insertions(+), 33 deletions(-) create mode 100644 test/correctness/split_aligned.cpp create mode 100644 test/correctness/split_aligned_2d.cpp diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index ddb9bc1098c5..0fa9cdf21e2f 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -26,7 +26,13 @@ vector apply_split(const Split &split, const string &prefix, dim_extent_alignment[split.inner] = split.factor; - Expr base = outer * split.factor + old_min; + Expr base; + if (split.align.defined()) { + base = outer * split.factor; + } else { + base = outer * split.factor + old_min; + } + string base_name = prefix + split.inner + ".base"; Expr base_var = Variable::make(Int(32), base_name); string old_var_name = prefix + split.old_var; @@ -58,14 +64,16 @@ vector apply_split(const Split &split, const string &prefix, // extent divides the factor. Use predication to guard // the calls and/or provides. - // Bounds inference has trouble exploiting an if - // condition. We'll directly tell it that the loop - // variable is bounded above by the original loop max by - // replacing the variable with a promise-clamped version - // of it. We don't also use the original loop min because - // it needlessly complicates the expressions and doesn't - // actually communicate anything new. - Expr guarded = promise_clamped(old_var, old_var, old_max); + Expr guarded; + if (split.align.defined()) { + // Because the un-rebased base block can start before old_min, + // we must clamp both the minimum and maximum boundaries. + guarded = promise_clamped(old_var, old_min, old_max); + } else { + // Legacy: structurally guaranteed to be >= old_min + guarded = promise_clamped(old_var, old_var, old_max); + } + string guarded_var_name = prefix + split.old_var + ".guarded"; Expr guarded_var = Variable::make(Int(32), guarded_var_name); @@ -76,8 +84,6 @@ vector apply_split(const Split &split, const string &prefix, predicate_type = ApplySplitResult::Predicate; break; case TailStrategy::Predicate: - // This is identical to GuardWithIf, but maybe it makes - // sense to keep it anyways? substitution_type = ApplySplitResult::Substitution; predicate_type = ApplySplitResult::Predicate; break; @@ -97,30 +103,44 @@ vector apply_split(const Split &split, const string &prefix, // for the guarded version. result.emplace_back(prefix + split.old_var, guarded_var, substitution_type); result.emplace_back(guarded_var_name, guarded, ApplySplitResult::LetStmt); - result.emplace_back(likely(old_var <= old_max), predicate_type); + + Expr guard_cond = likely(old_var <= old_max); + if (split.align.defined()) { + guard_cond = likely(old_var >= old_min && old_var <= old_max); + } + result.emplace_back(guard_cond, predicate_type); } else if (tail == TailStrategy::ShiftInwards) { // Adjust the base downwards to not compute off the // end of the realization. - // We'll only mark the base as likely (triggering a loop - // partition) if we're at or inside the innermost - // non-trivial loop. - base = likely_if_innermost(base); - base = Min::make(base, old_max + (1 - split.factor)); + base = likely(base); + if (split.align.defined()) { + base = Max::make(base, old_min - split.align); + base = Min::make(base, old_max + (1 - split.factor) - split.align); + } else { + base = Min::make(base, old_max + (1 - split.factor)); + } } else if (tail == TailStrategy::ShiftInwardsAndBlend) { Expr old_base = base; base = likely(base); - base = Min::make(base, old_max + (1 - split.factor)); + if (split.align.defined()) { + base = Max::make(base, old_min - split.align); + base = Min::make(base, old_max + (1 - split.factor) - split.align); + } else { + base = Min::make(base, old_max + (1 - split.factor)); + } // Make a mask which will be a loop invariant if inner gets // vectorized, and apply it if we're in the tail. Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner >= unwanted_elems; + Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; + Expr mask = zero_based_inner >= unwanted_elems; mask = select(base == old_base, likely(const_true()), mask); result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner < split.factor - unwanted_elems; + Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; + Expr mask = zero_based_inner < split.factor - unwanted_elems; mask = select(outer < outer_max, likely(const_true()), mask); result.emplace_back(mask, ApplySplitResult::BlendProvides); } else { @@ -173,12 +193,22 @@ vector> compute_loop_bounds_after_split(const Split &spl Expr old_var_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); switch (split.split_type) { case Split::SplitVar: { - Expr inner_extent = split.factor; - Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor; - let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); - let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1); - let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0); - let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1); + if (split.align.defined()) { + Expr align = split.align; + Expr outer_min = (old_var_min - align) / split.factor; + Expr outer_max = (old_var_max - align) / split.factor; + let_stmts.emplace_back(prefix + split.inner + ".loop_min", align); + let_stmts.emplace_back(prefix + split.inner + ".loop_max", align + split.factor - 1); + let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min); + let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); + } else { + Expr inner_extent = split.factor; + Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor; + let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); + let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1); + let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0); + let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1); + } } break; case Split::FuseVars: { // Define bounds on the fused var using the bounds on the inner and outer diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..6391eb49297a 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1103,9 +1103,9 @@ Func Stage::rfactor(const vector> &preserved) { return intm; } -void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, bool exact, TailStrategy tail) { +void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, const Expr &align_arg, bool exact, TailStrategy tail) { debug(4) << "In schedule for " << name() << ", split " << old << " into " - << outer << " and " << inner << " with factor of " << factor_arg << "\n"; + << outer << " and " << inner << " with factor of " << factor_arg << " and align " << align_arg << "\n"; user_assert(factor_arg.defined()) << "In schedule for " << name() << ", split factor for splitting " @@ -1115,6 +1115,14 @@ void Stage::split(const string &old, const string &outer, const string &inner, c << old << " has type " << factor_arg.type() << ", which is not representable as int32.\n"; Expr factor = cast(factor_arg); + Expr align; + if (align_arg.defined()) { + user_assert(Int(32).can_represent(align_arg.type())) + << "In schedule for " << name() << ", split align for splitting " + << old << " has type " << align_arg.type() + << ", which is not representable as int32.\n"; + align = cast(align_arg); + } vector &dims = definition.schedule().dims(); @@ -1318,11 +1326,15 @@ void Stage::split(const string &old, const string &outer, const string &inner, c } // Add the split to the splits list - Split split = {old_name, outer_name, inner_name, factor, exact, tail, Split::SplitVar}; + Split split = {old_name, outer_name, inner_name, factor, align, exact, tail, Split::SplitVar}; definition.schedule().splits().push_back(split); } -Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail) { +void Stage::split(const std::string &old, const std::string &outer, const std::string &inner, const Expr &factor, bool exact, TailStrategy tail) { + split(old, outer, inner, factor, Expr(), exact, tail); +} + +Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail) { definition.schedule().touched() = true; if (old.is_rvar) { user_assert(outer.is_rvar) << "Can't split RVar " << old.name() << " into Var " << outer.name() << "\n"; @@ -1331,7 +1343,13 @@ Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVa user_assert(!outer.is_rvar) << "Can't split Var " << old.name() << " into RVar " << outer.name() << "\n"; user_assert(!inner.is_rvar) << "Can't split Var " << old.name() << " into RVar " << inner.name() << "\n"; } - split(old.name(), outer.name(), inner.name(), factor, old.is_rvar, tail); + split(old.name(), outer.name(), inner.name(), factor, align, old.is_rvar, tail); + return *this; +} + +Stage &Stage::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail) { + definition.schedule().touched() = true; + split(old.name(), outer.name(), inner.name(), factor, Expr(), old.is_rvar, tail); return *this; } @@ -1413,7 +1431,7 @@ Stage &Stage::fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRV set_dim_type(fused, dims[inner_pos].for_type); // Add the fuse to the splits list - Split split = {fused_name, outer_name, inner_name, Expr(), true, TailStrategy::RoundUp, Split::FuseVars}; + Split split = {fused_name, outer_name, inner_name, Expr(), Expr(), true, TailStrategy::RoundUp, Split::FuseVars}; definition.schedule().splits().push_back(split); return *this; } @@ -1664,7 +1682,7 @@ Stage &Stage::rename(const VarOrRVar &old_var, const VarOrRVar &new_var) { } if (!found) { - Split split = {old_name, new_name, "", 1, old_var.is_rvar, TailStrategy::RoundUp, Split::RenameVar}; + Split split = {old_name, new_name, "", 1, Expr(), old_var.is_rvar, TailStrategy::RoundUp, Split::RenameVar}; definition.schedule().splits().push_back(split); } @@ -2545,6 +2563,12 @@ Func &Func::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar return *this; } +Func &Func::split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail) { + invalidate_cache(); + Stage(func, func.definition(), 0).split(old, outer, inner, factor, align, tail); + return *this; +} + Func &Func::fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRVar &fused) { invalidate_cache(); Stage(func, func.definition(), 0).fuse(inner, outer, fused); diff --git a/src/Func.h b/src/Func.h index 4df562e272ca..d5f7813c50a7 100644 --- a/src/Func.h +++ b/src/Func.h @@ -81,6 +81,8 @@ class Stage { void set_dim_device_api(const VarOrRVar &var, DeviceAPI device_api); void split(const std::string &old, const std::string &outer, const std::string &inner, const Expr &factor, bool exact, TailStrategy tail); + void split(const std::string &old, const std::string &outer, const std::string &inner, + const Expr &factor, const Expr &align, bool exact, TailStrategy tail); void remove(const std::string &var); const std::vector &storage_dims() const { @@ -365,6 +367,7 @@ class Stage { // @{ Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail = TailStrategy::Auto); Stage &fuse(const VarOrRVar &inner, const VarOrRVar &outer, const VarOrRVar &fused); Stage &serial(const VarOrRVar &var); Stage ¶llel(const VarOrRVar &var); @@ -1519,6 +1522,8 @@ class Func { * factor does not provably divide the extent. */ Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail = TailStrategy::Auto); + /** Join two dimensions into a single fused dimension. The fused dimension * covers the product of the extents of the inner and outer dimensions * given. The loop type (e.g. parallel, vectorized) of the resulting fused diff --git a/src/Schedule.h b/src/Schedule.h index ba3d1eea5ca3..034d74960f6d 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -334,6 +334,7 @@ struct ReductionVariable; struct Split { std::string old_var, outer, inner; Expr factor; + Expr align; bool exact; // Is it required that the factor divides the extent // of the old var. True for splits of RVars. Forces // tail strategy to be GuardWithIf. diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d88c13fa177f..452791193957 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -339,6 +339,8 @@ tests( specialize_to_gpu.cpp specialize_trim_condition.cpp spirv_ir.cpp + split_aligned.cpp + split_aligned_2d.cpp split_by_non_factor.cpp split_factor_type.cpp split_fuse_rvar.cpp diff --git a/test/correctness/split_aligned.cpp b/test/correctness/split_aligned.cpp new file mode 100644 index 000000000000..5801185db446 --- /dev/null +++ b/test/correctness/split_aligned.cpp @@ -0,0 +1,95 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, xo{"xo"}, xi{"xi"}; + for (auto ts : {TailStrategy::ShiftInwards, TailStrategy::GuardWithIf}) { + Func f; + Param offset{"offset"}; + offset.set_range(0, 3); + f(x) = mux((x - offset) % 4, {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); + f + .split(x, xo, xi, 4, offset, ts) + .unroll(xi); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &f : module.functions()) { + f.body.accept(&checker); + } + + for (int i = 0; i < 4; i++) { + printf("Testing runtime alignment: %d\n", i); + offset.set(i); + Buffer im = f.realize({32}); + f.realize(im, get_target_from_environment()); + + for (int x = 0; x < 32; x++) { + int selector = (4 + x - offset.get()) % 4; + int expected; + if (selector == 0) { + expected = x; + } else if (selector == 1) { + expected = x * x; + } else if (selector == 2) { + expected = 2 * x; + } else { + expected = -x * (x + 1); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (selector: %d)\n", x, im(x), expected, selector); + return 1; + } + } + } + + if (ts == Halide::TailStrategy::ShiftInwards) { + if (checker.mux_count != 8) { + std::printf("Expected 8 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 1) { + // The head and tail are reduced to a single iteration, so the loop is stripped. + std::printf("Expected one for loop: %d\n", checker.for_count); + return 1; + } + } else if (ts == Halide::TailStrategy::GuardWithIf) { + if (checker.mux_count != 0) { + std::printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 2) { + // The head and tail are reduced to a single iteration, so the loop is stripped. + std::printf("Expected one for loop: %d\n", checker.for_count); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/split_aligned_2d.cpp b/test/correctness/split_aligned_2d.cpp new file mode 100644 index 000000000000..22b7d8e14763 --- /dev/null +++ b/test/correctness/split_aligned_2d.cpp @@ -0,0 +1,90 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Var y{"y"}, yo{"yo"}, yi{"yi"}; + Func f; + Param offset_x{"offset_x"}, offset_y{"offset_y"}; + offset_x.set_range(0, 1); + offset_y.set_range(0, 1); + auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { + return (2 * ((y - offset_y) % 2)) + ((x - offset_x) % 2); + }; + auto a = [](const auto &x, const auto &y) { return x * x; }; + auto b = [](const auto &x, const auto &y) { return x * y; }; + auto c = [](const auto &x, const auto &y) { return y * y; }; + auto d = [](const auto &x, const auto &y) { return x + y; }; + f(x, y) = mux(idx(x, y, offset_x, offset_y), {a(x, y), b(x, y), c(x, y), d(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + + f + .split(x, xo, xi, 2, offset_x, Halide::TailStrategy::GuardWithIf) + .split(y, yo, yi, 2, offset_y, Halide::TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .parallel(yo); + + Module module = f.compile_to_module({offset_x, offset_y}); + MuxCounter checker; + for (const LoweredFunc &f : module.functions()) { + f.body.accept(&checker); + } + + for (int i = 0; i < 4; i++) { + printf("Testing runtime alignment: x=%d y=%d\n", i / 2, i % 2); + offset_x.set(i / 2); + offset_y.set(i % 2); + Buffer im = f.realize({32, 32}); + f.realize(im, get_target_from_environment()); + + for (int y = 0; y < 32; y++) { + for (int x = 0; x < 32; x++) { + int selector = idx(2 + x, 2 + y, offset_x.get(), offset_y.get()); + int expected = std::vector>{a, b, c, d}[selector](x, y); + if (im(x, y) != expected) { + printf("im(%d, %d) = %d instead of %d (selector: %d)\n", x, y, im(x, y), expected, selector); + return 1; + } + } + } + } + + if (checker.mux_count != 12) { + std::printf("Expected 12 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 3) { + std::printf("Expected 3 for loops: %d\n", checker.for_count); + return 1; + } + + printf("Success!\n"); + return 0; +} From 9ebb0323c5a58d46ac787765eaffef30a2d910bb Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 19:57:15 +0200 Subject: [PATCH 02/29] Restore likely_if_innermost for ShiftInwards. --- src/ApplySplit.cpp | 2 +- test/correctness/split_aligned_2d.cpp | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 0fa9cdf21e2f..f4fcecb1e25d 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -114,7 +114,7 @@ vector apply_split(const Split &split, const string &prefix, // Adjust the base downwards to not compute off the // end of the realization. - base = likely(base); + base = likely_if_innermost(base); if (split.align.defined()) { base = Max::make(base, old_min - split.align); base = Min::make(base, old_max + (1 - split.factor) - split.align); diff --git a/test/correctness/split_aligned_2d.cpp b/test/correctness/split_aligned_2d.cpp index 22b7d8e14763..99ac9743d8c9 100644 --- a/test/correctness/split_aligned_2d.cpp +++ b/test/correctness/split_aligned_2d.cpp @@ -76,11 +76,11 @@ int main(int argc, char **argv) { } } - if (checker.mux_count != 12) { - std::printf("Expected 12 muxes: %d\n", checker.mux_count); + if (checker.mux_count != 0) { + std::printf("Expected 0 muxes: %d\n", checker.mux_count); return 1; } - if (checker.for_count != 3) { + if (checker.for_count != 1) { std::printf("Expected 3 for loops: %d\n", checker.for_count); return 1; } From 1315f7ad41eb4d882ceddb560e96e228d3af9799 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 21:48:45 +0200 Subject: [PATCH 03/29] Add Python binding Co-authored-by: Claude Sonnet 5 --- python_bindings/halide/src/halide_/PyScheduleMethods.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python_bindings/halide/src/halide_/PyScheduleMethods.h b/python_bindings/halide/src/halide_/PyScheduleMethods.h index f528af886dff..7e585690e1e3 100644 --- a/python_bindings/halide/src/halide_/PyScheduleMethods.h +++ b/python_bindings/halide/src/halide_/PyScheduleMethods.h @@ -29,6 +29,8 @@ HALIDE_NEVER_INLINE void add_schedule_methods(PythonClass &class_instance) { .def("split", (T & (T::*)(const VarOrRVar &, const VarOrRVar &, const VarOrRVar &, const Expr &, TailStrategy)) & T::split, py::arg("old"), py::arg("outer"), py::arg("inner"), py::arg("factor"), py::arg("tail") = TailStrategy::Auto) + .def("split", (T & (T::*)(const VarOrRVar &, const VarOrRVar &, const VarOrRVar &, const Expr &, const Expr &, TailStrategy)) & T::split, + py::arg("old"), py::arg("outer"), py::arg("inner"), py::arg("factor"), py::arg("align"), py::arg("tail") = TailStrategy::Auto) .def("fuse", &T::fuse, py::arg("inner"), py::arg("outer"), py::arg("fused")) From 97dfaee1ef7b3dd6bdbc67d03cfb139f53c07bd4 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 21:49:04 +0200 Subject: [PATCH 04/29] Add serialization. Co-authored-by: Claude Sonnet 5 --- src/Deserialization.cpp | 2 ++ src/Serialization.cpp | 4 +++- src/halide_ir.fbs | 1 + 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index f7f8566326db..be75ec77d8da 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -1152,6 +1152,7 @@ Split Deserializer::deserialize_split(const Serialize::Split *split) { const auto exact = split->exact(); const auto tail = deserialize_tail_strategy(split->tail()); const auto split_type = deserialize_split_type(split->split_type()); + const auto align = deserialize_expr(split->align_type(), split->align()); auto hl_split = Split(); hl_split.old_var = old_var; hl_split.outer = outer; @@ -1160,6 +1161,7 @@ Split Deserializer::deserialize_split(const Serialize::Split *split) { hl_split.exact = exact; hl_split.tail = tail; hl_split.split_type = split_type; + hl_split.align = align; return hl_split; } diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 2dd7bf4f33aa..36ea9d84984f 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -1259,10 +1259,12 @@ Offset Serializer::serialize_split(FlatBufferBuilder &builder, const auto exact = split.exact; const auto tail_serialized = serialize_tail_strategy(split.tail); const auto split_type_serialized = serialize_split_type(split.split_type); + const auto align_serialized = serialize_expr(builder, split.align); return Serialize::CreateSplit(builder, old_var_serialized, outer_serialized, inner_serialized, factor_serialized.first, factor_serialized.second, - exact, tail_serialized, split_type_serialized); + exact, tail_serialized, split_type_serialized, + align_serialized.first, align_serialized.second); } Offset Serializer::serialize_dim(FlatBufferBuilder &builder, const Dim &dim) { diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 4bba4bb79a8f..3e81a44fb3b1 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -568,6 +568,7 @@ table Split { exact: bool; tail: TailStrategy; split_type: SplitType; + align: Expr; } enum DimType: ubyte { From 5afb9fc34bd8851b0292e5f7d43cdb2968134179 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:48:57 +0200 Subject: [PATCH 05/29] Fix the incorrectly assumed fast-path for this aligned splits. Co-authored-by: Claude Sonnet 5 --- src/ApplySplit.cpp | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index f4fcecb1e25d..98c88f81905c 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -44,8 +44,17 @@ vector apply_split(const Split &split, const string &prefix, internal_assert(tail != TailStrategy::Auto) << "An explicit tail strategy should exist at this point\n"; + // When align is defined, tiles are anchored to align instead of to + // old_min, so knowing that the factor divides the extent is not + // enough to prove no boundary guard is needed: we additionally need + // the tiling anchored at align to line up with the tiling anchored + // at old_min, i.e. old_min and align must be congruent mod factor. + bool alignment_matches_old_min = !split.align.defined() || + is_const_zero(simplify((old_min - split.align) % split.factor)); + if ((iter != dim_extent_alignment.end()) && - is_const_zero(simplify(iter->second % split.factor))) { + is_const_zero(simplify(iter->second % split.factor)) && + alignment_matches_old_min) { // We have proved that the split factor divides the // old extent. No need to adjust the base or add an if // statement. From 0ab58c424a8abd7516b054bb6f47dcb03a41d869 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:51:44 +0200 Subject: [PATCH 06/29] Add three rfactor + aligned split tests. Co-authored-by: Claude Sonnet 5 Co-authored-by: Gemini Pro 3.1 --- src/Simplify_Add.cpp | 1 + src/Simplify_Mod.cpp | 4 + test/correctness/CMakeLists.txt | 3 + test/correctness/rfactor_split_aligned.cpp | 94 ++++++++++ test/correctness/rfactor_split_aligned_2d.cpp | 101 +++++++++++ .../rfactor_split_aligned_phases.cpp | 168 ++++++++++++++++++ 6 files changed, 371 insertions(+) create mode 100644 test/correctness/rfactor_split_aligned.cpp create mode 100644 test/correctness/rfactor_split_aligned_2d.cpp create mode 100644 test/correctness/rfactor_split_aligned_phases.cpp diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index a07ad1b4464b..a2298c2e019c 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -201,6 +201,7 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite(x + ((c0 - x) / c1) * c1, c0 - ((c0 - x) % c1), c1 > 0) || rewrite(x + ((c0 - x) / c1 + y) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || rewrite(x + (y + (c0 - x) / c1) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || + rewrite(((0 - x) / c0) + ((x % c0 + c1) / c0), (c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || false)))) { return mutate(rewrite.result, info); diff --git a/src/Simplify_Mod.cpp b/src/Simplify_Mod.cpp index 7e5232da0975..0bbbddb4ec34 100644 --- a/src/Simplify_Mod.cpp +++ b/src/Simplify_Mod.cpp @@ -59,6 +59,10 @@ Expr Simplify::visit(const Mod *op, ExprInfo *info) { rewrite((x * c0 - y) % c1, (-y) % c1, c0 % c1 == 0) || rewrite((y - x * c0) % c1, y % c1, c0 % c1 == 0) || rewrite((x - y) % 2, (x + y) % 2) || // Addition and subtraction are the same modulo 2, because -1 == 1 + rewrite((((x * c0) + y) - z) % c0, (y - z) % c0) || + rewrite((((x * c0) + y) + z) % c0, (y + z) % c0) || + rewrite((((x * c0) - y) - z) % c0, (-y - z) % c0) || + rewrite((((x * c0) - y) + z) % c0, (z - y) % c0) || rewrite(ramp(x, c0, c2) % broadcast(c1, c2), broadcast(x, c2) % broadcast(c1, c2), (c0 % c1 == 0)) || rewrite(ramp(x, c0, lanes) % broadcast(c1, lanes), ramp(x % c1, c0, lanes), diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 452791193957..3bc62f47c261 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -483,6 +483,9 @@ tests( random.cpp reorder_rvars.cpp rfactor.cpp + rfactor_split_aligned.cpp + rfactor_split_aligned_2d.cpp + rfactor_split_aligned_phases.cpp ring_buffer.cpp stream_compaction.cpp thread_safety.cpp diff --git a/test/correctness/rfactor_split_aligned.cpp b/test/correctness/rfactor_split_aligned.cpp new file mode 100644 index 000000000000..19f51dfa63bc --- /dev/null +++ b/test/correctness/rfactor_split_aligned.cpp @@ -0,0 +1,94 @@ +#include "Halide.h" +#include + +// rfactor() eagerly applies any splits present on the RVar(s) it's given (see +// Stage::rfactor / project_rdom in Func.cpp), so it needs to tolerate splits +// that carry an alignment (Stage::split's 'align' argument) just as well as +// ordinary ones. This test factors the *outer* half of an aligned split of +// the reduction variable out into a parallel-reducible intermediate Func, +// while unrolling the *inner* (aligned) half in the reducing computation. +// Because the inner half is not itself preserved by rfactor(), it keeps the +// exact loop bounds computed by compute_loop_bounds_after_split (rather than +// being re-derived by general bounds inference), so unrolling it still lets +// the compiler resolve the runtime-offset mux() to a compile-time constant +// per lane, exactly as it does without rfactor in split_aligned.cpp. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ro, u); + intm.compute_root(); + intm.update(0).parallel(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes (the aligned+unrolled inner split var should " + "resolve the mux at compile time even after rfactor): %d\n", + checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + int selector = (4 + r - off) % 4; + int term; + if (selector == 0) { + term = r + x; + } else if (selector == 1) { + term = r * r + x; + } else if (selector == 2) { + term = 2 * r + x; + } else { + term = -r * (r + 1) + x; + } + expected += term; + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, off); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor_split_aligned_2d.cpp b/test/correctness/rfactor_split_aligned_2d.cpp new file mode 100644 index 000000000000..9120c26a4578 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_2d.cpp @@ -0,0 +1,101 @@ +#include "Halide.h" +#include + +// A 2D companion to rfactor_split_aligned.cpp. Here rfactor() is applied to +// an RVar (r.x) that is unrelated to the one carrying the aligned split +// (r.y), which is the more common pattern in practice: factor out one +// reduction dimension for parallel/vector reduction while a separate +// dimension is scheduled with an alignment-aware split so a +// runtime-offset-dependent mux() can be resolved statically once its half of +// the split is unrolled. Since r.y's split is entirely unrelated to the +// preserved var, both halves of the split remain ordinary (non-preserved) +// reduction variables of the intermediate Func, retaining their exact +// compile-time loop bounds and so still collapsing the mux to nothing. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, y{"y"}; + Func f{"f"}; + RDom r(0, 20, 0, 16, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x, y) = 0; + f(x, y) += mux((r.y - offset) % 4, + {r.x + r.y + x + y, + r.x * r.y + x - y, + 2 * r.x - r.y + x, + -r.x * (r.y + 1) + y}) * + select(r.x % 2 == 0, 1, -1); + + RVar ryo{"ryo"}, ryi{"ryi"}; + f.update(0) + .split(r.y, ryo, ryi, 4, offset, TailStrategy::GuardWithIf) + .unroll(ryi); + + Var u{"u"}; + Func intm = f.update(0).rfactor(r.x, u); + intm.compute_root(); + intm.update(0).parallel(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({6, 6}); + for (int y = 0; y < 6; y++) { + for (int x = 0; x < 6; x++) { + int expected = 0; + for (int rx = 0; rx < 20; rx++) { + for (int ry = 0; ry < 16; ry++) { + int selector = (4 + ry - off) % 4; + int term; + if (selector == 0) { + term = rx + ry + x + y; + } else if (selector == 1) { + term = rx * ry + x - y; + } else if (selector == 2) { + term = 2 * rx - ry + x; + } else { + term = -rx * (ry + 1) + y; + } + term *= (rx % 2 == 0) ? 1 : -1; + expected += term; + } + } + if (im(x, y) != expected) { + printf("im(%d, %d) = %d instead of %d (offset: %d)\n", x, y, im(x, y), expected, off); + return 1; + } + } + } + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor_split_aligned_phases.cpp b/test/correctness/rfactor_split_aligned_phases.cpp new file mode 100644 index 000000000000..2f1cfd277205 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_phases.cpp @@ -0,0 +1,168 @@ +#include "Halide.h" +#include + +// A variant of rfactor_split_aligned.cpp that preserves the *inner* (aligned, +// unrolled) half of the split via rfactor() instead of the outer half, +// turning it into four separate per-phase partial-sum accumulators that get +// combined at the end. rfactor() must still produce correct results here: +// this is precisely the "does rfactor tolerate splits with an alignment" +// question, exercised in the case where the aligned split is the one being +// preserved (and therefore promoted from an RVar with exact, +// compute_loop_bounds_after_split-derived bounds to an ordinary pure Var of +// the intermediate Func, whose bounds are instead re-derived by general +// bounds inference). That promotion means the compiler can no longer read +// off the new pure var's range directly from the split; it has to prove it +// symbolically from the surrounding min/max clamps instead, which is what +// the mux_count checks below are exercising. +// +// The second case additionally makes the RDom's own extent a runtime Param +// rather than a compile-time constant, so the split's "factor provably +// divides the extent" fast path (see apply_split in ApplySplit.cpp) can't +// fire either, and everything -- the boundary guard, the alignment, and the +// mux resolution -- has to come out of the general GuardWithIf path instead. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int expected_value(int r, int x, int off) { + int selector = (4 + r - off) % 4; + if (selector == 0) { + return r + x; + } else if (selector == 1) { + return r * r + x; + } else if (selector == 2) { + return 2 * r + x; + } else { + return -r * (r + 1) + x; + } +} + +int test_fixed_extent() { + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ri, u); + intm.compute_root(); + intm.update(0).unroll(u); + + Module module = f.compile_to_module({offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + printf("Testing runtime alignment: %d\n", off); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + expected += expected_value(r, x, off); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, off); + return 1; + } + } + } + + return 0; +} + +int test_param_extent() { + Var x{"x"}; + Func f{"f"}; + Param extent{"extent"}; + RDom r(0, extent, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}; + Func intm = f.update(0).rfactor(ri, u); + intm.compute_root(); + intm.update(0).unroll(u); + + Module module = f.compile_to_module({extent, offset}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes (with a Param extent): %d\n", checker.mux_count); + return 1; + } + + // 40 is a multiple of the split factor; 37 is not, so it also exercises + // the tail of the RDom's own range. + for (int ext : {40, 37}) { + for (int off = 0; off < 4; off++) { + printf("Testing runtime extent %d, alignment %d\n", ext, off); + extent.set(ext); + offset.set(off); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < ext; r++) { + expected += expected_value(r, x, off); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (extent: %d, offset: %d)\n", x, im(x), expected, ext, off); + return 1; + } + } + } + } + + return 0; +} + +int main(int argc, char **argv) { + if (test_fixed_extent()) { + return 1; + } + if (test_param_extent()) { + return 1; + } + + printf("Success!\n"); + return 0; +} From f201909f1190c0e801a517c6bbab83eb2fad64d4 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:54:25 +0200 Subject: [PATCH 07/29] Add test for nested aligned splits. Co-authored-by: Claude Sonnet 5 --- src/Schedule.cpp | 6 ++ test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_nested.cpp | 84 +++++++++++++++++++++++ 3 files changed, 91 insertions(+) create mode 100644 test/correctness/split_aligned_nested.cpp diff --git a/src/Schedule.cpp b/src/Schedule.cpp index 948233112b7c..77f27d8e89a6 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -340,6 +340,9 @@ struct StageScheduleContents { if (s.factor.defined()) { s.factor = mutator(s.factor); } + if (s.align.defined()) { + s.align = mutator(s.align); + } } for (PrefetchDirective &p : prefetches) { if (p.offset.defined()) { @@ -702,6 +705,9 @@ void StageSchedule::accept(IRVisitor *visitor) const { if (s.factor.defined()) { s.factor.accept(visitor); } + if (s.align.defined()) { + s.align.accept(visitor); + } } for (const PrefetchDirective &p : prefetches()) { if (p.offset.defined()) { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 3bc62f47c261..d289e86c027f 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -341,6 +341,7 @@ tests( spirv_ir.cpp split_aligned.cpp split_aligned_2d.cpp + split_aligned_nested.cpp split_by_non_factor.cpp split_factor_type.cpp split_fuse_rvar.cpp diff --git a/test/correctness/split_aligned_nested.cpp b/test/correctness/split_aligned_nested.cpp new file mode 100644 index 000000000000..774e653efd6a --- /dev/null +++ b/test/correctness/split_aligned_nested.cpp @@ -0,0 +1,84 @@ +#include "Halide.h" +#include + +// Nests two aligned splits: x is split into (xo, xi) aligned to p1, and then +// the resulting outer var xo is itself split into (xoo, xoi) aligned to a +// second, independent runtime Param p2. This exercises the aligned-split +// machinery (ApplySplit.cpp's apply_split/compute_loop_bounds_after_split) +// on a var whose own loop_min is not a compile-time constant (it comes from +// the first split's outer bound, which is a function of p1), stacked with a +// second, unrelated alignment. The mux selector only depends on p1, so this +// is primarily a correctness test of composing aligned splits -- the +// reconstruction of x from xoo, xoi, and xi has to be correct for every +// combination of the two independently-varying runtime alignments. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + Var x{"x"}, xo{"xo"}, xi{"xi"}, xoo{"xoo"}, xoi{"xoi"}; + Func f{"f"}; + Param p1{"p1"}, p2{"p2"}; + p1.set_range(0, 3); + p2.set_range(0, 2); + + f(x) = mux((x - p1) % 4, {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); + + f.split(x, xo, xi, 4, p1, TailStrategy::GuardWithIf) + .split(xo, xoo, xoi, 3, p2, TailStrategy::GuardWithIf) + .unroll(xi); + + Module module = f.compile_to_module({p1, p2}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int a1 = 0; a1 < 4; a1++) { + for (int a2 = 0; a2 < 3; a2++) { + printf("Testing runtime alignment: p1=%d p2=%d\n", a1, a2); + p1.set(a1); + p2.set(a2); + Buffer im = f.realize({61}); + for (int x = 0; x < 61; x++) { + int selector = (4 + x - a1) % 4; + int expected; + if (selector == 0) { + expected = x; + } else if (selector == 1) { + expected = x * x; + } else if (selector == 2) { + expected = 2 * x; + } else { + expected = -x * (x + 1); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (p1: %d, p2: %d)\n", x, im(x), expected, a1, a2); + return 1; + } + } + } + } + + printf("Success!\n"); + return 0; +} From bacdbc8e7ed5cf19f11d75340e7f3afcb89c715f Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:58:52 +0200 Subject: [PATCH 08/29] Test varying tail strategies for nested aligned splits. Co-authored-by: Claude Sonnet 5 --- test/correctness/split_aligned_nested.cpp | 96 +++++++++++++---------- 1 file changed, 56 insertions(+), 40 deletions(-) diff --git a/test/correctness/split_aligned_nested.cpp b/test/correctness/split_aligned_nested.cpp index 774e653efd6a..59b411c614f6 100644 --- a/test/correctness/split_aligned_nested.cpp +++ b/test/correctness/split_aligned_nested.cpp @@ -11,6 +11,15 @@ // is primarily a correctness test of composing aligned splits -- the // reconstruction of x from xoo, xoi, and xi has to be correct for every // combination of the two independently-varying runtime alignments. +// +// Both splits are tried with both GuardWithIf and ShiftInwards (as in +// split_aligned.cpp): correctness must hold for all four combinations, and +// as in split_aligned.cpp the mux only fully resolves at compile time (0 +// muxes) when the split that carries the selector's alignment (the first +// one, on x) uses GuardWithIf; ShiftInwards leaves 8 muxes unresolved +// because the clamped base is no longer a compile-time-constant offset from +// the unrolled lane on every iteration. The tail strategy of the second, +// unrelated split (on xo) doesn't affect that count either way. using namespace Halide; using namespace Halide::Internal; @@ -30,50 +39,57 @@ class MuxCounter : public IRVisitor { }; int main(int argc, char **argv) { - Var x{"x"}, xo{"xo"}, xi{"xi"}, xoo{"xoo"}, xoi{"xoi"}; - Func f{"f"}; - Param p1{"p1"}, p2{"p2"}; - p1.set_range(0, 3); - p2.set_range(0, 2); + for (auto ts1 : {TailStrategy::GuardWithIf, TailStrategy::ShiftInwards}) { + for (auto ts2 : {TailStrategy::GuardWithIf, TailStrategy::ShiftInwards}) { + printf("Testing tail strategies: ts1=%d ts2=%d\n", (int)ts1, (int)ts2); - f(x) = mux((x - p1) % 4, {x, x * x, 2 * x, -x * (x + 1)}); - f.output_buffer().dim(0).set_min(0); + Var x{"x"}, xo{"xo"}, xi{"xi"}, xoo{"xoo"}, xoi{"xoi"}; + Func f{"f"}; + Param p1{"p1"}, p2{"p2"}; + p1.set_range(0, 3); + p2.set_range(0, 2); - f.split(x, xo, xi, 4, p1, TailStrategy::GuardWithIf) - .split(xo, xoo, xoi, 3, p2, TailStrategy::GuardWithIf) - .unroll(xi); + f(x) = mux((x - p1) % 4, {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); - Module module = f.compile_to_module({p1, p2}); - MuxCounter checker; - for (const LoweredFunc &lf : module.functions()) { - lf.body.accept(&checker); - } - if (checker.mux_count != 0) { - printf("Expected 0 muxes: %d\n", checker.mux_count); - return 1; - } + f.split(x, xo, xi, 4, p1, ts1) + .split(xo, xoo, xoi, 3, p2, ts2) + .unroll(xi); - for (int a1 = 0; a1 < 4; a1++) { - for (int a2 = 0; a2 < 3; a2++) { - printf("Testing runtime alignment: p1=%d p2=%d\n", a1, a2); - p1.set(a1); - p2.set(a2); - Buffer im = f.realize({61}); - for (int x = 0; x < 61; x++) { - int selector = (4 + x - a1) % 4; - int expected; - if (selector == 0) { - expected = x; - } else if (selector == 1) { - expected = x * x; - } else if (selector == 2) { - expected = 2 * x; - } else { - expected = -x * (x + 1); - } - if (im(x) != expected) { - printf("im(%d) = %d instead of %d (p1: %d, p2: %d)\n", x, im(x), expected, a1, a2); - return 1; + Module module = f.compile_to_module({p1, p2}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + int expected_mux_count = (ts1 == TailStrategy::GuardWithIf) ? 0 : 8; + if (checker.mux_count != expected_mux_count) { + printf("Expected %d muxes: %d\n", expected_mux_count, checker.mux_count); + return 1; + } + + for (int a1 = 0; a1 < 4; a1++) { + for (int a2 = 0; a2 < 3; a2++) { + p1.set(a1); + p2.set(a2); + Buffer im = f.realize({61}); + for (int x = 0; x < 61; x++) { + int selector = (4 + x - a1) % 4; + int expected; + if (selector == 0) { + expected = x; + } else if (selector == 1) { + expected = x * x; + } else if (selector == 2) { + expected = 2 * x; + } else { + expected = -x * (x + 1); + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (p1: %d, p2: %d, ts1: %d, ts2: %d)\n", + x, im(x), expected, a1, a2, (int)ts1, (int)ts2); + return 1; + } + } } } } From 1a5f500a518cea3aa7cce7c774962cf740ce4fac Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 23:36:57 +0200 Subject: [PATCH 09/29] Documentation for the aligned split. Co-authored-by: Claude Sonnet 5 --- src/Func.h | 27 +++++++++++++++++++++++++++ src/Schedule.h | 3 ++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/Func.h b/src/Func.h index d5f7813c50a7..2de86684a8e7 100644 --- a/src/Func.h +++ b/src/Func.h @@ -1522,6 +1522,33 @@ class Func { * factor does not provably divide the extent. */ Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); + /** A version of split() that additionally takes a runtime-valued + * phase, 'align', which need not be known at compile time. Instead + * of the inner dimension always iterating over [0, factor-1], it + * iterates over [align, align+factor-1]. This may increase the + * number of iterations over the outer loop by 1 compared to an + * unaligned split. + * + * This is useful when an algorithm selects between cases using an + * expression like ``(x - offset) % factor``, where 'offset' is a + * value only known at runtime (e.g. a Param). Passing that same + * 'offset' as 'align' makes ``(x - offset) % factor`` a + * compile-time constant on each unrolled iteration of the inner + * loop, so that a mux() indexed by it can be resolved statically + * instead of compiling to a runtime select: + \code + Var x, xo, xi; + Param offset; + f(x) = mux((x - offset) % 4, {a(x), b(x), c(x), d(x)}); + f.split(x, xo, xi, 4, offset, TailStrategy::GuardWithIf) + .unroll(xi); + \endcode + * Without 'align', the compiler can't tell at compile time which of + * the four mux() cases applies to a given unrolled value of 'xi', + * because that depends on the runtime value of 'offset'. With it, + * ``(x - offset) % 4`` simplifies to a distinct compile-time + * constant for each unrolled value of 'xi', and each mux() call + * collapses to its selected case. */ Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, const Expr &align, TailStrategy tail = TailStrategy::Auto); /** Join two dimensions into a single fused dimension. The fused dimension diff --git a/src/Schedule.h b/src/Schedule.h index 034d74960f6d..7bfa92981ac1 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -334,7 +334,8 @@ struct ReductionVariable; struct Split { std::string old_var, outer, inner; Expr factor; - Expr align; + Expr align; // If defined, the inner var loops over [align, + // align + factor - 1] instead of [0, factor - 1]. bool exact; // Is it required that the factor divides the extent // of the old var. True for splits of RVars. Forces // tail strategy to be GuardWithIf. From 31bd47a34c6e57d0c91c8e77c1eaf66665391d11 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 23:47:20 +0200 Subject: [PATCH 10/29] Fix ShiftInwardsAndBlend, RoundUpAndBlend. Claude rederived the masks those blend operations in case of aligned splits. Co-authored-by: Claude Sonnet 5 --- src/ApplySplit.cpp | 87 +++++++++++-- test/correctness/CMakeLists.txt | 1 + .../rfactor_split_aligned_nested.cpp | 114 ++++++++++++++++++ 3 files changed, 191 insertions(+), 11 deletions(-) create mode 100644 test/correctness/rfactor_split_aligned_nested.cpp diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 98c88f81905c..6df71d634b9e 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -23,6 +23,7 @@ vector apply_split(const Split &split, const string &prefix, Expr old_max = Variable::make(Int(32), prefix + split.old_var + ".loop_max"); Expr old_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); Expr old_extent = (old_max - old_min) + 1; + Expr outer_min = Variable::make(Int(32), prefix + split.outer + ".loop_min"); dim_extent_alignment[split.inner] = split.factor; @@ -131,26 +132,90 @@ vector apply_split(const Split &split, const string &prefix, base = Min::make(base, old_max + (1 - split.factor)); } } else if (tail == TailStrategy::ShiftInwardsAndBlend) { + // Unclamped base, saved before the Min/Max below adjust it. Used + // to figure out how much (if at all) the boundary tile got + // shifted, so we know which elements of it are redundant with a + // neighboring tile and must be masked out rather than + // recomputed (to avoid double-counting in a reduction). Expr old_base = base; base = likely(base); + Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; + Expr mask; if (split.align.defined()) { - base = Max::make(base, old_min - split.align); - base = Min::make(base, old_max + (1 - split.factor) - split.align); + // Because base is anchored to align instead of old_min, the + // boundary tile can now be shifted at either end (whereas + // without align only the max end is reachable, since base + // is structurally >= old_min already). Elements shifted in + // from the low end overlap the tile above (mask out the + // last shift_low of them); elements shifted in from the + // high end overlap the tile below (mask out the first + // shift_high of them). + Expr low_bound = old_min - split.align; + Expr high_bound = old_max + (1 - split.factor) - split.align; + Expr shift_low = low_bound - old_base; + Expr shift_high = old_base - high_bound; + base = Max::make(base, low_bound); + base = Min::make(base, high_bound); + Expr mask_low = zero_based_inner < split.factor - shift_low; + Expr mask_high = zero_based_inner >= shift_high; + mask = select(old_base < low_bound, mask_low, + select(old_base > high_bound, mask_high, likely(const_true()))); } else { + // Without align, base is structurally >= old_min (outer + // starts at 0), so only the max end can ever be shifted. base = Min::make(base, old_max + (1 - split.factor)); + Expr unwanted_elems = (-old_extent) % split.factor; + mask = zero_based_inner >= unwanted_elems; + mask = select(base == old_base, likely(const_true()), mask); } - // Make a mask which will be a loop invariant if inner gets - // vectorized, and apply it if we're in the tail. - Expr unwanted_elems = (-old_extent) % split.factor; - Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; - Expr mask = zero_based_inner >= unwanted_elems; - mask = select(base == old_base, likely(const_true()), mask); result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { - Expr unwanted_elems = (-old_extent) % split.factor; Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; - Expr mask = zero_based_inner < split.factor - unwanted_elems; - mask = select(outer < outer_max, likely(const_true()), mask); + Expr mask; + if (split.align.defined()) { + // Unlike ShiftInwardsAndBlend, the max end is intentionally + // left unclamped here (RoundUp relies on padding, not on + // shifting, to handle overrun at the max end) -- but the min + // end still needs clamping: align can make the min-end tile + // start before old_min, and unlike ShiftInwards/blend at the + // max end, there's no padding below old_min to absorb an + // underrun into, so it has to be prevented outright. + // + // The mask below compares old_base (the unclamped base) + // against low_bound/high_bound directly, rather than + // comparing outer against outer_min/outer_max: the latter + // needs loop partitioning to split the loop into three + // pieces (prologue/steady-state/epilogue) to stay correct, + // and partition_loops doesn't reliably do that here when + // both boundaries are data-dependent, silently dropping the + // last tile. Comparing old_base against the bounds directly + // is correct regardless of how (or whether) the loop gets + // partitioned, matching the approach already proven correct + // above for ShiftInwardsAndBlend. + Expr old_base = base; + Expr low_bound = old_min - split.align; + Expr high_bound = old_max + (1 - split.factor) - split.align; + Expr shift_low = low_bound - old_base; + Expr shift_high = old_base - high_bound; + base = Max::make(likely(base), low_bound); + // The min end is clamped (shifted forward), so its overlap + // is with the tile *above* -- same geometry as + // ShiftInwardsAndBlend, mask out the trailing shift_low + // elements. The max end is left unclamped, so shift_high + // counts a genuine overrun past old_max with no + // neighboring tile to defer to -- mask out the trailing + // shift_high elements too (the opposite convention from + // ShiftInwardsAndBlend's clamped max end, which instead + // masks out the *leading* elements of a shifted-back tile). + Expr mask_low = zero_based_inner < split.factor - shift_low; + Expr mask_high = zero_based_inner < split.factor - shift_high; + mask = select(old_base < low_bound, mask_low, + select(old_base > high_bound, mask_high, likely(const_true()))); + } else { + Expr unwanted_elems = (-old_extent) % split.factor; + Expr fresh_high = zero_based_inner < split.factor - unwanted_elems; + mask = select(outer < outer_max, likely(const_true()), fresh_high); + } result.emplace_back(mask, ApplySplitResult::BlendProvides); } else { internal_assert(tail == TailStrategy::RoundUp); diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d289e86c027f..f0b3866e2671 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -486,6 +486,7 @@ tests( rfactor.cpp rfactor_split_aligned.cpp rfactor_split_aligned_2d.cpp + rfactor_split_aligned_nested.cpp rfactor_split_aligned_phases.cpp ring_buffer.cpp stream_compaction.cpp diff --git a/test/correctness/rfactor_split_aligned_nested.cpp b/test/correctness/rfactor_split_aligned_nested.cpp new file mode 100644 index 000000000000..58c30b7056d5 --- /dev/null +++ b/test/correctness/rfactor_split_aligned_nested.cpp @@ -0,0 +1,114 @@ +#include "Halide.h" +#include + +// A companion to rfactor_split_aligned.cpp and split_aligned_nested.cpp: +// after r's aligned split (factor 4, aligned to offset) is rfactored on its +// outer half into a preserved pure var u, u is itself split again with a +// second, independent alignment (p2), tried with GuardWithIf, +// RoundUpAndBlend, and ShiftInwardsAndBlend. +// +// GuardWithIf and Predicate are the only tail strategies Stage::split allows +// on an RVar (splitting r itself), because RoundUp/ShiftInwards-family +// strategies would change the meaning of a reduction by recomputing or +// overrunning it -- but u is an ordinary pure Var of the intermediate +// Func's own update definition, so RoundUpAndBlend/ShiftInwardsAndBlend +// (the update-definition-safe counterparts of RoundUp/ShiftInwards) are +// legal there, and are exactly the tail strategies meant for vectorizing +// an update like this one. +// +// This combination exercises boundary handling in ApplySplit.cpp +// (apply_split's ShiftInwardsAndBlend/RoundUpAndBlend branches) that plain, +// unnested aligned splits don't: u's own old_min is not a compile-time +// constant (it comes from r's split, a function of the runtime offset +// Param), so both the low and high boundary tiles of u's split can only be +// distinguished from the interior at runtime. + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + +public: + int mux_count{0}; +}; + +int main(int argc, char **argv) { + for (auto ts : {TailStrategy::GuardWithIf, TailStrategy::RoundUpAndBlend, TailStrategy::ShiftInwardsAndBlend}) { + printf("Testing tail strategy: %d\n", (int)ts); + + Var x{"x"}; + Func f{"f"}; + RDom r(0, 40, "r"); + Param offset{"offset"}; + offset.set_range(0, 3); + + f(x) = 0; + f(x) += mux((r - offset) % 4, {r + x, r * r + x, 2 * r + x, -r * (r + 1) + x}); + + RVar ro{"ro"}, ri{"ri"}; + f.update(0) + .split(r, ro, ri, 4, offset, TailStrategy::GuardWithIf) + .unroll(ri); + + Var u{"u"}, uo{"uo"}, ui{"ui"}; + Param p2{"p2"}; + p2.set_range(0, 1); + + Func intm = f.update(0).rfactor(ro, u); + intm.compute_root(); + intm.update(0) + .split(u, uo, ui, 2, p2, ts) + .vectorize(ui); + + Module module = f.compile_to_module({offset, p2}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + if (checker.mux_count != 0) { + printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + + for (int off = 0; off < 4; off++) { + for (int a2 = 0; a2 < 2; a2++) { + offset.set(off); + p2.set(a2); + Buffer im = f.realize({10}); + for (int x = 0; x < 10; x++) { + int expected = 0; + for (int r = 0; r < 40; r++) { + int selector = (4 + r - off) % 4; + int term; + if (selector == 0) { + term = r + x; + } else if (selector == 1) { + term = r * r + x; + } else if (selector == 2) { + term = 2 * r + x; + } else { + term = -r * (r + 1) + x; + } + expected += term; + } + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d, p2: %d, ts: %d)\n", + x, im(x), expected, off, a2, (int)ts); + return 1; + } + } + } + } + } + + printf("Success!\n"); + return 0; +} From 884b58cfaf49f2b4b986f9bce0b3fb9d755a7e05 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 19 Aug 2026 17:00:59 +0200 Subject: [PATCH 11/29] Add simplifier rules for broadcast() <= ramp() && ramp() <= broadcast(). Fix old copy-paste bug in simplifier rules. Co-Authored-By: Claude Sonnet 5 --- src/Simplify_Exprs.cpp | 19 ++++++++++++++++-- test/correctness/simplify.cpp | 37 +++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index c19fa2e7fed8..b3ce824f0cf1 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -214,8 +214,23 @@ Expr Simplify::visit(const VectorReduce *op, ExprInfo *info) { x + max(y * (arg_lanes - 1), 0) <= z) || rewrite(h_and(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), x < y + min(z * (arg_lanes - 1), 0)) || - rewrite(h_and(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), + rewrite(h_and(broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes), 1), x <= y + min(z * (arg_lanes - 1), 0)) || + + // The "all lanes of a ramp lie within [lo, hi]" check loop + // partitioning builds (a lower-bound comparison ANDed with an + // upper-bound comparison, both against the same stride, e.g. + // (0 <= ramp(b0, s, n)) && (ramp(b1, s, n) <= extent)) + rewrite(h_and((broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)) && + (ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)), + 1), + (x <= y + min(z * (arg_lanes - 1), 0)) && + (w + max(z * (arg_lanes - 1), 0) <= u)) || + rewrite(h_and((ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)) && + (broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)), + 1), + (w + max(z * (arg_lanes - 1), 0) <= u) && + (x <= y + min(z * (arg_lanes - 1), 0))) || false) { return mutate(rewrite.result, info); } @@ -237,7 +252,7 @@ Expr Simplify::visit(const VectorReduce *op, ExprInfo *info) { x + min(y * (arg_lanes - 1), 0) <= z) || rewrite(h_or(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), x < y + max(z * (arg_lanes - 1), 0)) || - rewrite(h_or(broadcast(x, arg_lanes) < ramp(y, z, arg_lanes), 1), + rewrite(h_or(broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes), 1), x <= y + max(z * (arg_lanes - 1), 0)) || false) { return mutate(rewrite.result, info); diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 981a00e6f0ee..538ef1982303 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -1767,6 +1767,43 @@ void check_boolean() { check(ramp(x * 8 + 5, -1, 4) < broadcast(y * 8, 4), broadcast(x < y, 4)); check(ramp(x * 8 - 1, -1, 4) < broadcast(y * 8, 4), broadcast(x < y + 1, 4)); + // A horizontal AND/OR of a single ramp/broadcast comparison collapses to + // a plain scalar comparison on the ramp's endpoints, for both orderings + // of ramp vs broadcast and both '<' and '<='. + check(VectorReduce::make(VectorReduce::And, ramp(x, y, 4) < broadcast(z, 4), 1), + max(y, 0) * 3 + x < z); + check(VectorReduce::make(VectorReduce::And, ramp(x, y, 4) <= broadcast(z, 4), 1), + max(y, 0) * 3 + x <= z); + check(VectorReduce::make(VectorReduce::And, broadcast(x, 4) < ramp(y, z, 4), 1), + x < min(z, 0) * 3 + y); + check(VectorReduce::make(VectorReduce::And, broadcast(x, 4) <= ramp(y, z, 4), 1), + x <= min(z, 0) * 3 + y); + + check(VectorReduce::make(VectorReduce::Or, ramp(x, y, 4) < broadcast(z, 4), 1), + min(y, 0) * 3 + x < z); + check(VectorReduce::make(VectorReduce::Or, ramp(x, y, 4) <= broadcast(z, 4), 1), + min(y, 0) * 3 + x <= z); + check(VectorReduce::make(VectorReduce::Or, broadcast(x, 4) < ramp(y, z, 4), 1), + x < max(z, 0) * 3 + y); + check(VectorReduce::make(VectorReduce::Or, broadcast(x, 4) <= ramp(y, z, 4), 1), + x <= max(z, 0) * 3 + y); + + // The "all lanes of a ramp lie within [lo, hi]" shape loop partitioning + // builds -- a lower-bound comparison ANDed with an upper-bound + // comparison, both against the same stride -- fuses to a plain And of + // two scalar comparisons, regardless of clause order. + { + Expr u = Var("u"); + check(VectorReduce::make(VectorReduce::And, + (broadcast(x, 4) <= ramp(y, z, 4)) && (ramp(w, z, 4) <= broadcast(u, 4)), + 1), + (x <= min(z, 0) * 3 + y) && (max(z, 0) * 3 + w <= u)); + check(VectorReduce::make(VectorReduce::And, + (ramp(w, z, 4) <= broadcast(u, 4)) && (broadcast(x, 4) <= ramp(y, z, 4)), + 1), + (max(z, 0) * 3 + w <= u) && (x <= min(z, 0) * 3 + y)); + } + // Check anded conditions apply to the then case only check(IfThenElse::make(x == 4 && y == 5, not_no_op(z + x + y), From 978ab0895f61c132ce5e15dbb4830ffb3dee8ccb Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 21 Aug 2026 00:06:08 +0200 Subject: [PATCH 12/29] Test simple aligned split in an RVar. Co-Authored-By: Claude Sonnet 5 --- test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_reduction.cpp | 66 ++++++++++++++++++++ 2 files changed, 67 insertions(+) create mode 100644 test/correctness/split_aligned_reduction.cpp diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index f0b3866e2671..a3ea2d7a2290 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -342,6 +342,7 @@ tests( split_aligned.cpp split_aligned_2d.cpp split_aligned_nested.cpp + split_aligned_reduction.cpp split_by_non_factor.cpp split_factor_type.cpp split_fuse_rvar.cpp diff --git a/test/correctness/split_aligned_reduction.cpp b/test/correctness/split_aligned_reduction.cpp new file mode 100644 index 000000000000..65c1fc66a9e8 --- /dev/null +++ b/test/correctness/split_aligned_reduction.cpp @@ -0,0 +1,66 @@ +#include "Halide.h" +#include + +// A simple reduction (no rfactor) with a single aligned split, tried with +// GuardWithIf, RoundUpAndBlend, and ShiftInwardsAndBlend. +// +// The split here is of the pure var x, not of the RDom's r: Stage::split +// only allows GuardWithIf or Predicate when splitting an RVar itself (see +// Func.cpp), since RoundUp/ShiftInwards-family strategies would change the +// meaning of the reduction by recomputing or overrunning it. Splitting a +// pure var of an update definition doesn't have that restriction, and +// RoundUpAndBlend/ShiftInwardsAndBlend are exactly the tail strategies +// meant for vectorizing an update like this one (see their doc comments in +// Schedule.h). +// +// This is the same boundary-handling code in ApplySplit.cpp's +// ShiftInwardsAndBlend/RoundUpAndBlend branches exercised by +// rfactor_split_aligned_nested.cpp, but without rfactor's extra layer of +// indirection (splitting a var that's already itself the result of an +// aligned split) -- here x's own bounds are simple compile-time constants, +// so this isolates the aligned-split-plus-blend mechanics on their own. + +using namespace Halide; + +int main(int argc, char **argv) { + for (auto ts : {TailStrategy::GuardWithIf, TailStrategy::RoundUpAndBlend, TailStrategy::ShiftInwardsAndBlend}) { + printf("Testing tail strategy: %d\n", (int)ts); + + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Func h{"h"}; + RDom r(0, 5, "r"); + Param p{"p"}; + p.set_range(0, 3); + + h(x) = 0; + h(x) += x + r; + h.compute_root(); + + h.update(0) + .split(x, xo, xi, 4, p, ts) + .vectorize(xi); + + // h is read through a further Func rather than realized directly, + // so that RoundUpAndBlend/ShiftInwardsAndBlend get an + // internally-allocated (and thus paddable) buffer to blend into, + // instead of a caller-provided one of a fixed, non-factor-multiple + // size. + Func out{"out"}; + out(x) = h(x); + + for (int a = 0; a < 4; a++) { + p.set(a); + Buffer im = out.realize({37}); + for (int x = 0; x < 37; x++) { + int expected = 5 * x + 10; + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (p: %d, ts: %d)\n", x, im(x), expected, a, (int)ts); + return 1; + } + } + } + } + + printf("Success!\n"); + return 0; +} From a7405cdf9ee77672dd9c564a7678abeb69f70522 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 24 Aug 2026 15:05:40 +0200 Subject: [PATCH 13/29] Reduce aligned-split compute_at mux test to 3x3 and fix its checks Rename split_aligned_2d_6x6.cpp to split_aligned_2d_3x3.cpp and shrink the pattern to 3x3, which reproduces the surviving mux with a much smaller amount of IR to read. Also fix the test itself: realize the 3-D output with a 3-D shape, check all three channels, sweep all nine (offset_x, offset_y) alignments, and include c in the reorder so it stays innermost. With c left outermost it was unrolled around the xo/yo nest, triplicating the loop nest and recomputing R/G/B once per channel. The test currently fails at the mux count (27 = 9 tile positions x 3 channels); the runtime results are correct for every alignment. --- test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_2d_3x3.cpp | 115 ++++++++++++++++++++++ 2 files changed, 116 insertions(+) create mode 100644 test/correctness/split_aligned_2d_3x3.cpp diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index a3ea2d7a2290..d55bf1246168 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -341,6 +341,7 @@ tests( spirv_ir.cpp split_aligned.cpp split_aligned_2d.cpp + split_aligned_2d_3x3.cpp split_aligned_nested.cpp split_aligned_reduction.cpp split_by_non_factor.cpp diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp new file mode 100644 index 000000000000..f72ca58c3c9e --- /dev/null +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -0,0 +1,115 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +class MuxCounter : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + + void visit(const For *op) override { + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count{0}; + int mux_count{0}; +}; + +template +T produce_mux_argument(int i, const T &x, const T &y) { + return x * (i % 3) * 3 + y * (i % 3); +} + +int main(int argc, char **argv) { + Var c{"c"}; + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Var y{"y"}, yo{"yo"}, yi{"yi"}; + Func f("f"), R("R"), G("G"), B("B"); + Param offset_x{"offset_x"}, offset_y{"offset_y"}; + auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { + return (3 * ((y - offset_y) % 3)) + ((x - offset_x) % 3); + }; + std::vector ways; + ways.reserve(9); + for (int i = 0; i < 9; ++i) { + ways.push_back(produce_mux_argument(i, x, y)); + } + R(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + G(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + B(x, y) = mux(idx(x, y, offset_x, offset_y), ways); + f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + + // Split both dimensions so that the inner loops iterate over exactly one + // 3x3 tile of the repeating pattern, anchored at (offset_x, offset_y). + // Unrolling those inner loops should give each mux a constant index, so + // every mux folds away to the single way it selects. + f + .split(x, xo, xi, 3, offset_x, Halide::TailStrategy::GuardWithIf) + .split(y, yo, yi, 3, offset_y, Halide::TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(c, xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .bound(c, 0, 3) + .unroll(c); + + for (Func *channel : {&R, &G, &B}) { + channel->compute_at(f, xo).unroll(x).unroll(y).never_partition_all(); + } + + Module module = f.compile_to_module({offset_x, offset_y}); + MuxCounter checker; + for (const LoweredFunc &lf : module.functions()) { + lf.body.accept(&checker); + } + + const int W = 32, H = 32; + for (int oy = 0; oy < 3; oy++) { + for (int ox = 0; ox < 3; ox++) { + printf("Testing runtime alignment: x=%d y=%d\n", ox, oy); + offset_x.set(ox); + offset_y.set(oy); + Buffer im = f.realize({W, H, 3}); + + for (int cc = 0; cc < 3; cc++) { + for (int y = 0; y < H; y++) { + for (int x = 0; x < W; x++) { + // Bias by 3 so the operands of % stay non-negative, + // where C++'s truncated % agrees with Halide's + // Euclidean %. + int selector = idx(3 + x, 3 + y, ox, oy); + int expected = produce_mux_argument(selector, x, y); + if (im(x, y, cc) != expected) { + printf("im(%d, %d, %d) = %d instead of %d (selector: %d)\n", + x, y, cc, im(x, y, cc), expected, selector); + return 1; + } + } + } + } + } + } + + if (checker.mux_count != 0) { + printf("Expected 0 muxes, got: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 2) { + printf("Expected 2 for loops, got: %d\n", checker.for_count); + return 1; + } + + printf("Success!\n"); + return 0; +} From 546075e7c14b8e804a0884eed2c03ea7ccf1e5fa Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sat, 29 Aug 2026 19:47:03 +0200 Subject: [PATCH 14/29] Rewrite aligned split inner loops to range from 0 to factor. Co-Authored-By: Claude Sonnet 5 --- src/ApplySplit.cpp | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 6df71d634b9e..8dc177f896ab 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -139,7 +139,6 @@ vector apply_split(const Split &split, const string &prefix, // recomputed (to avoid double-counting in a reduction). Expr old_base = base; base = likely(base); - Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; Expr mask; if (split.align.defined()) { // Because base is anchored to align instead of old_min, the @@ -156,8 +155,8 @@ vector apply_split(const Split &split, const string &prefix, Expr shift_high = old_base - high_bound; base = Max::make(base, low_bound); base = Min::make(base, high_bound); - Expr mask_low = zero_based_inner < split.factor - shift_low; - Expr mask_high = zero_based_inner >= shift_high; + Expr mask_low = inner < split.factor - shift_low; + Expr mask_high = inner >= shift_high; mask = select(old_base < low_bound, mask_low, select(old_base > high_bound, mask_high, likely(const_true()))); } else { @@ -165,12 +164,11 @@ vector apply_split(const Split &split, const string &prefix, // starts at 0), so only the max end can ever be shifted. base = Min::make(base, old_max + (1 - split.factor)); Expr unwanted_elems = (-old_extent) % split.factor; - mask = zero_based_inner >= unwanted_elems; + mask = inner >= unwanted_elems; mask = select(base == old_base, likely(const_true()), mask); } result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { - Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; Expr mask; if (split.align.defined()) { // Unlike ShiftInwardsAndBlend, the max end is intentionally @@ -207,13 +205,13 @@ vector apply_split(const Split &split, const string &prefix, // shift_high elements too (the opposite convention from // ShiftInwardsAndBlend's clamped max end, which instead // masks out the *leading* elements of a shifted-back tile). - Expr mask_low = zero_based_inner < split.factor - shift_low; - Expr mask_high = zero_based_inner < split.factor - shift_high; + Expr mask_low = inner < split.factor - shift_low; + Expr mask_high = inner < split.factor - shift_high; mask = select(old_base < low_bound, mask_low, select(old_base > high_bound, mask_high, likely(const_true()))); } else { Expr unwanted_elems = (-old_extent) % split.factor; - Expr fresh_high = zero_based_inner < split.factor - unwanted_elems; + Expr fresh_high = inner < split.factor - unwanted_elems; mask = select(outer < outer_max, likely(const_true()), fresh_high); } result.emplace_back(mask, ApplySplitResult::BlendProvides); @@ -221,6 +219,17 @@ vector apply_split(const Split &split, const string &prefix, internal_assert(tail == TailStrategy::RoundUp); } + // Add align back in last, after all tail-strategy clamping/masking is + // done in terms of the unaligned base: this keeps align as a bare + // top-level addend in the final expressions (so e.g. it can still + // cancel algebraically against a matching subtraction elsewhere) + // rather than being smeared into a Max/Min-clamped expression, while + // letting the inner loop variable itself range over the simple, + // often-constant [0, factor) instead of [align, align + factor). + if (split.align.defined()) { + base = base + split.align; + } + // Define the original variable as the base value computed above plus the inner loop variable. result.emplace_back(old_var_name, base_var + inner, ApplySplitResult::LetStmt); result.emplace_back(base_name, base, ApplySplitResult::LetStmt); @@ -271,8 +280,8 @@ vector> compute_loop_bounds_after_split(const Split &spl Expr align = split.align; Expr outer_min = (old_var_min - align) / split.factor; Expr outer_max = (old_var_max - align) / split.factor; - let_stmts.emplace_back(prefix + split.inner + ".loop_min", align); - let_stmts.emplace_back(prefix + split.inner + ".loop_max", align + split.factor - 1); + let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); + let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 1); let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min); let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); } else { From 6f12d1429b8afe06fb198daceb3519a036e72ea1 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sat, 29 Aug 2026 19:48:23 +0200 Subject: [PATCH 15/29] Override the compute and storage bounds of the test to relieve the simplifier. --- test/correctness/split_aligned_2d_3x3.cpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp index f72ca58c3c9e..39fd1a54f117 100644 --- a/test/correctness/split_aligned_2d_3x3.cpp +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -35,6 +35,8 @@ int main(int argc, char **argv) { Var y{"y"}, yo{"yo"}, yi{"yi"}; Func f("f"), R("R"), G("G"), B("B"); Param offset_x{"offset_x"}, offset_y{"offset_y"}; + offset_x.set_range(0, 2); + offset_y.set_range(0, 2); auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { return (3 * ((y - offset_y) % 3)) + ((x - offset_x) % 3); }; @@ -65,7 +67,16 @@ int main(int argc, char **argv) { .unroll(c); for (Func *channel : {&R, &G, &B}) { - channel->compute_at(f, xo).unroll(x).unroll(y).never_partition_all(); + channel->compute_at(f, xo) + .unroll(x) + .bound_extent(x, 3) + .bound_storage(x, 3) + .align_bounds(x, 3, offset_x) + .unroll(y) + .bound_extent(y, 3) + .bound_storage(y, 3) + .align_bounds(y, 3, offset_y) + .never_partition_all(); } Module module = f.compile_to_module({offset_x, offset_y}); From 483927f65cad1ff28f0417ee3e6c9cd5b0fef302 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 25 Aug 2026 10:08:54 +0200 Subject: [PATCH 16/29] Add a test for an aligned split feeding loop partitioning A loop of eight whose first and last iterations are special and whose interior is periodic with period two. Unrolling the interior by two folds the % away, but only if the unrolled pairs line up with the periodicity, which means the tiles have to start where the interior does. An aligned split says exactly that, and partitioning then peels one iteration at each end rather than two, leaving a steady-state loop of three rather than two. Checks the extent of the remaining loop, that the modulo folded away, and the values. Dropping the alignment from the split fails the extent check, so the test is measuring the thing it claims to. --- test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_partition.cpp | 116 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 test/correctness/split_aligned_partition.cpp diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d55bf1246168..93048b3bed21 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -343,6 +343,7 @@ tests( split_aligned_2d.cpp split_aligned_2d_3x3.cpp split_aligned_nested.cpp + split_aligned_partition.cpp split_aligned_reduction.cpp split_by_non_factor.cpp split_factor_type.cpp diff --git a/test/correctness/split_aligned_partition.cpp b/test/correctness/split_aligned_partition.cpp new file mode 100644 index 000000000000..6e457b70dc2b --- /dev/null +++ b/test/correctness/split_aligned_partition.cpp @@ -0,0 +1,116 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +// Note: this test is built with NDEBUG, so assert() compiles to nothing. +bool check(bool ok, const char *msg) { + if (!ok) { + printf("Failed: %s\n", msg); + } + return ok; +} + +class LoopExtents : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + extents.push_back(simplify(op->extent())); + IRVisitor::visit(op); + } + +public: + std::vector extents; +}; + +class CountMod : public IRVisitor { + using IRVisitor::visit; + + void visit(const Mod *op) override { + count++; + IRVisitor::visit(op); + } + +public: + int count{0}; +}; + +} // namespace + +int main(int argc, char **argv) { + // A loop of eight elements whose first and last iterations are special, + // and whose interior is periodic with period two. Unrolling the interior + // by two turns the % into a constant, but only if the unrolled pairs line + // up with the periodicity -- which means the tiles have to start at x=1, + // where the interior begins, not at x=0. + // + // An aligned split expresses exactly that: split by two, anchored at one. + // Loop partitioning then peels the one iteration at each end that the + // likely() marks as not-steady-state, leaving x in [1, 6]. That's six + // iterations, or three of the unrolled-by-two loop. + // + // Without the alignment the tiles start at x=0 instead, the interior + // doesn't fill a whole number of them, and partitioning has to peel two + // iterations at each end rather than one -- leaving a steady-state loop + // of two rather than three. + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Func f{"f"}; + f(x) = select(x <= 0, 100, + x < 7, likely(x % 2), + 200); + f.bound(x, 0, 8); + f.split(x, xo, xi, 2, 1, TailStrategy::GuardWithIf) + .always_partition(xo) + .unroll(xi); + + Module m = f.compile_to_module({}, "f"); + + LoopExtents loops; + CountMod mods; + for (const LoweredFunc &lf : m.functions()) { + lf.body.accept(&loops); + lf.body.accept(&mods); + } + + printf("Loops:"); + for (const Expr &e : loops.extents) { + std::cout << " " << e; + } + printf("\n"); + + // The two peeled iterations are single elements, so they come out as + // straight-line code rather than loops. What's left is the steady state. + if (!check(loops.extents.size() == 1, "expected exactly one remaining loop")) { + return 1; + } + if (!check(is_const(loops.extents[0], 3), + "expected the steady-state loop to run three times (six " + "elements, unrolled by two)")) { + return 1; + } + + // The whole point of unrolling the interior was to fold away the %. + if (!check(mods.count == 0, "expected the modulo to fold away")) { + return 1; + } + + Buffer out = f.realize({8}); + for (int i = 0; i < 8; i++) { + int expected = 200; + if (i <= 0) { + expected = 100; + } else if (i < 7) { + expected = i % 2; + } + if (out(i) != expected) { + printf("out(%d) = %d instead of %d\n", i, out(i), expected); + return 1; + } + } + + printf("Success!\n"); + return 0; +} From 63fbab364c094dccf8a6722edc469fb35d03e280 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sun, 30 Aug 2026 10:13:01 +0200 Subject: [PATCH 17/29] Explain the schedule of aligned_split_2d. --- test/correctness/split_aligned_2d_3x3.cpp | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp index 39fd1a54f117..7fa73fba8f4b 100644 --- a/test/correctness/split_aligned_2d_3x3.cpp +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -68,15 +68,20 @@ int main(int argc, char **argv) { for (Func *channel : {&R, &G, &B}) { channel->compute_at(f, xo) + .never_partition_all() .unroll(x) + .unroll(y) + .align_bounds(y, 3, offset_y) + .align_bounds(x, 3, offset_x) + // In principle, the above scheduling directives should be enough, but the simplifier + // in Halide fails to determine that the extent is 3. So here, we can help it and + // tell it to use an extent of 3, after which the simplifier manages to prove that + // 3 satisfies the minimum required. The simplifier gap manifests itself twice: once + // for the compute extent, and once for the the storage extent. .bound_extent(x, 3) .bound_storage(x, 3) - .align_bounds(x, 3, offset_x) - .unroll(y) .bound_extent(y, 3) - .bound_storage(y, 3) - .align_bounds(y, 3, offset_y) - .never_partition_all(); + .bound_storage(y, 3); } Module module = f.compile_to_module({offset_x, offset_y}); From 603e14e64cb01f9db7ea786a4e855b3585d45f7a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sun, 30 Aug 2026 10:56:04 +0200 Subject: [PATCH 18/29] Document a scheduling order dependence. --- test/correctness/split_aligned_2d_3x3.cpp | 30 ++++++++++++++--------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp index 7fa73fba8f4b..37e060481141 100644 --- a/test/correctness/split_aligned_2d_3x3.cpp +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -69,19 +69,27 @@ int main(int argc, char **argv) { for (Func *channel : {&R, &G, &B}) { channel->compute_at(f, xo) .never_partition_all() - .unroll(x) - .unroll(y) - .align_bounds(y, 3, offset_y) - .align_bounds(x, 3, offset_x) - // In principle, the above scheduling directives should be enough, but the simplifier - // in Halide fails to determine that the extent is 3. So here, we can help it and - // tell it to use an extent of 3, after which the simplifier manages to prove that - // 3 satisfies the minimum required. The simplifier gap manifests itself twice: once - // for the compute extent, and once for the the storage extent. - .bound_extent(x, 3) + // We want to compute a *whole* tile of the channel: .bound_storage(x, 3) + .bound_extent(x, 3) + .align_bounds(x, 3, offset_x) + .bound_storage(y, 3) .bound_extent(y, 3) - .bound_storage(y, 3); + .align_bounds(y, 3, offset_y) + // In principle, the .align_bounds() should be enough, but the simplifier in Halide + // fails to determine that the producer tile perfectly overlaps with the consumer + // tile. Both have extent 3, but the producer is not getting simplified. So here, + // we can help it and tell it to use an extent of 3, after which the simplifier + // manages to prove that 3 satisfies the minimum required. Proving the expression + // is <= 3 is easier than proving it's == 3. + // The simplifier gap manifests itself twice: once for the compute extent, + // and once for the the storage extent. So we specify both. Unfortunately, + // the scheduling order here is important: .bound_extent() must precede + // .align_bounds() for the bounds inference to do the correct thing. + // + // Finally, we unroll to get rid of the muxes: + .unroll(x) + .unroll(y); } Module module = f.compile_to_module({offset_x, offset_y}); From c8e7c7400a7040b286c109696aa49a95d26bd8ac Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sun, 30 Aug 2026 13:31:00 +0200 Subject: [PATCH 19/29] Tutorial 25 for aligned split. --- src/Func.h | 16 +- tutorial/CMakeLists.txt | 3 +- tutorial/lesson_25_aligned_split.cpp | 605 +++++++++++++++++++++++++++ 3 files changed, 618 insertions(+), 6 deletions(-) create mode 100644 tutorial/lesson_25_aligned_split.cpp diff --git a/src/Func.h b/src/Func.h index 2de86684a8e7..d7a7d64fb7b1 100644 --- a/src/Func.h +++ b/src/Func.h @@ -1523,11 +1523,17 @@ class Func { Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); /** A version of split() that additionally takes a runtime-valued - * phase, 'align', which need not be known at compile time. Instead - * of the inner dimension always iterating over [0, factor-1], it - * iterates over [align, align+factor-1]. This may increase the - * number of iterations over the outer loop by 1 compared to an - * unaligned split. + * phase, 'align', which need not be known at compile time. An + * ordinary split() tiles a Var starting at its own loop_min, so + * tile boundaries fall at old_min, old_min + factor, old_min + 2 * + * factor, and so on. This variant anchors the tiling to 'align' + * instead, so tile boundaries fall at align, align + factor, align + * - factor, and so on. The inner dimension still iterates over + * [0, factor-1], same as an unaligned split; what changes is which + * value of the original Var each (outer, inner) pair reconstructs + * to. This may increase the number of iterations over the outer + * loop by 1 compared to an unaligned split, since align need not + * coincide with old_min. * * This is useful when an algorithm selects between cases using an * expression like ``(x - offset) % factor``, where 'offset' is a diff --git a/tutorial/CMakeLists.txt b/tutorial/CMakeLists.txt index 94c3d9e22be7..cd6091bf13d4 100644 --- a/tutorial/CMakeLists.txt +++ b/tutorial/CMakeLists.txt @@ -341,7 +341,8 @@ if (TARGET Halide::Mullapudi2016) ) endif () -# Lessons 22-24 +# Lessons 22-25 add_tutorial(lesson_22_jit_performance.cpp) add_tutorial(lesson_23_serialization.cpp WITH_IMAGE_IO) add_tutorial(lesson_24_async.cpp GROUPS multithreaded) +add_tutorial(lesson_25_aligned_split.cpp) diff --git a/tutorial/lesson_25_aligned_split.cpp b/tutorial/lesson_25_aligned_split.cpp new file mode 100644 index 000000000000..b5310cb504a7 --- /dev/null +++ b/tutorial/lesson_25_aligned_split.cpp @@ -0,0 +1,605 @@ +// Halide tutorial lesson 25: Aligned splits + +// This lesson demonstrates Func::split's 'align' parameter: a way to anchor +// a split's tile boundaries to a value that's only known at runtime (e.g. a +// Param), instead of to the Var's own min. It shows what that buys you -- +// primarily, letting mux() calls whose selector depends on that same +// runtime value collapse to the single case they select, instead of +// surviving as mux() calls whose case can only be picked at run time -- +// and where the trick currently runs out of steam: crossing a compute_at +// boundary. + +// On linux, you can compile and run it like so: +// g++ lesson_25*.cpp -g -I -L -lHalide -lpthread -ldl -o lesson_25 -std=c++17 +// LD_LIBRARY_PATH= ./lesson_25 + +// On macOS: +// g++ lesson_25*.cpp -g -I -L -lHalide -o lesson_25 -std=c++17 +// DYLD_LIBRARY_PATH= ./lesson_25 + +#include "Halide.h" +#include +#include + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +// A small IRVisitor used throughout this lesson to check what a schedule +// actually did to the compiled code: how many for loops remain, how wide +// they are, and how many mux() and % operations the simplifier managed to +// fold away. Ordinary Halide code never needs to do this -- it's only here +// so this lesson can show what each schedule accomplishes, rather than +// just assert it. +class Counter : public IRVisitor { + using IRVisitor::visit; + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->is_intrinsic(Call::IntrinsicOp::mux)) { + mux_count++; + } + } + void visit(const Mod *op) override { + IRVisitor::visit(op); + mod_count++; + } + void visit(const For *op) override { + extents.push_back(simplify(op->extent())); + IRVisitor::visit(op); + for_count++; + } + +public: + int for_count = 0, mux_count = 0, mod_count = 0; + std::vector extents; +}; + +Counter count(const Module &m) { + Counter c; + for (const LoweredFunc &lf : m.functions()) { + lf.body.accept(&c); + } + return c; +} + +// Finds the "produce { ... }" node for a given Func in a compiled +// Module, so this lesson can print the actual generated code for just that +// Func instead of the whole pipeline. This is the most direct way to see +// what a schedule bought you: read the loop nest and check for yourself +// whether a mux() call turned into a straight-line value or is still there +// as a mux() call with a non-constant selector. Since only that one +// subtree is printed, in a compute_at example you may see it reference a +// variable (e.g. a "let t123 = ...") that's bound just outside the +// printed excerpt, in the surrounding loop nest -- that's expected; the +// thing to focus on is the shape of the code inside, not resolving every +// hoisted temporary. +class FindProducer : public IRVisitor { + using IRVisitor::visit; + + void visit(const ProducerConsumer *p) override { + // When a Func's own name collides with the compiled Module's + // exported function name, Halide disambiguates the internal one by + // appending "$N" (Halide's general name-uniquification suffix), so + // match that too rather than only an exact name. + if (p->is_producer && (p->name == name || starts_with(p->name, name + "$"))) { + producer = p; + } + IRVisitor::visit(p); + } + +public: + FindProducer(std::string name) : name(std::move(name)) {} + + std::string name; + Stmt producer; +}; + +Stmt find_producer(const Module &m, const std::string &name) { + FindProducer f{name}; + for (const LoweredFunc &lf : m.functions()) { + lf.body.accept(&f); + if (f.producer.defined()) { + return f.producer; + } + } + return Stmt{}; +} + +} // namespace + +int main(int argc, char **argv) { + + // Part 1: what does "aligned" actually change? + // + // An ordinary split(x, xo, xi, factor) tiles a Var starting at its own + // loop_min: tile boundaries fall at old_min, old_min + factor, + // old_min + 2*factor, and so on. split(x, xo, xi, factor, align) + // anchors the tiling to 'align' instead, wherever that falls relative + // to old_min. 'align' need not be known at compile time -- that's the + // whole point of the feature -- but to see *why* the anchor point + // matters at all, this first example uses a compile-time-constant + // align of 1, and looks at how it interacts with loop partitioning. + { + // The pipeline below is 8 elements: a boundary value at x == 0, a + // boundary value at x == 7, and a period-2 pattern (x % 2) for the + // six elements in between. likely() marks the interior case as the + // expected one, which is what invites the loop partitioner to try + // to peel the boundary cases into their own prologue/epilogue + // rather than testing for them on every iteration. + // + // First, split by 2 with no alignment at all, and unroll the + // 2-wide inner loop so that x % 2 has a chance to become a + // compile-time constant in each unrolled copy. + Var x("x"), xo("xo"), xi("xi"); + Func plain("plain"); + plain(x) = select(x <= 0, 100, + x < 7, likely(x % 2), + 200); + plain.bound(x, 0, 8); + plain.split(x, xo, xi, 2, TailStrategy::GuardWithIf) + .always_partition(xo) + .unroll(xi); + auto mod_plain = plain.compile_to_module({}, "plain"); + std::cout << "Without alignment, tiled from x == 0:\n" + << find_producer(mod_plain, "plain") << "\n"; + Counter c_plain = count(mod_plain); + + // Now the same thing, but anchored to x == 1 -- where the + // interior's period-2 pattern actually begins -- instead of to + // the Func's own min of 0. (Func's assignment operator is a + // reference, not a clone, so this needs its own Func and Vars, + // not a copy of 'plain'.) + Var x2("x"), xo2("xo"), xi2("xi"); + Func aligned("aligned"); + aligned(x2) = select(x2 <= 0, 100, + x2 < 7, likely(x2 % 2), + 200); + aligned.bound(x2, 0, 8); + aligned.split(x2, xo2, xi2, 2, /* align */ 1, TailStrategy::GuardWithIf) + .always_partition(xo2) + .unroll(xi2); + auto mod_aligned = aligned.compile_to_module({}, "aligned"); + std::cout << "Anchored at x == 1:\n" + << find_producer(mod_aligned, "aligned") << "\n"; + Counter c_aligned = count(mod_aligned); + + // Either way, x % 2 folds away completely: every unrolled + // instance of xi turns (xo*2 + align + xi) % 2 into a compile-time + // constant, regardless of what 'align' is. (Try it: the multiple + // of 2 drops out of the modulo either way.) So this isn't about + // the modulo -- it's about how much of the *boundary* the + // partitioner can peel into single, non-looping iterations. + if (c_plain.mod_count != 0 || c_aligned.mod_count != 0) { + printf("Expected the %% to fold away in both cases\n"); + return 1; + } + + // Without alignment, tiles start at x == 0, so the single-element + // boundary case at x == 0 *is* a whole tile, and the one at x == 7 + // spills into the tile that starts at x == 6. The partitioner has + // to peel a full 2-element tile off of each end to isolate the + // boundary, leaving a steady-state loop that only covers x in + // [2, 5]: two unrolled-by-2 iterations. + // + // Anchored at x == 1, tiles start at 1, 3, 5: exactly where the + // interior's period-2 pattern repeats. Now the single-element + // boundary cases at x == 0 and x == 7 are each a partial tile of + // their own, so the partitioner peels exactly one element off each + // end, leaving a steady-state loop over x in [1, 6]: three + // unrolled-by-2 iterations, one more than the unaligned version. + if (c_plain.for_count != 1 || !is_const(c_plain.extents[0], 2)) { + printf("Expected one steady-state loop of extent 2 without alignment\n"); + return 1; + } + if (c_aligned.for_count != 1 || !is_const(c_aligned.extents[0], 3)) { + printf("Expected one steady-state loop of extent 3 with alignment\n"); + return 1; + } + + // Both schedules compute the same answer either way -- alignment + // only changes how the work is split into loops, never the + // result. + Buffer out = aligned.realize({8}); + for (int i = 0; i < 8; i++) { + int expected = (i <= 0) ? 100 : (i < 7) ? i % 2 : + 200; + if (out(i) != expected) { + printf("out(%d) = %d instead of %d\n", i, out(i), expected); + return 1; + } + } + } + + // Part 2: anchoring a split to a runtime Param + // + // Lesson 13 pointed out that mux(id, {...}) is sugar for a select + // chain, and that the select can be compiled away by bounding and + // unrolling the Var it's indexed by -- but only if 'id' becomes a + // compile-time constant in each unrolled copy. That's easy when 'id' + // is directly the unrolled Var. It's harder when 'id' depends on that + // Var only after subtracting off a value that's not known until + // runtime, as below. + // + // (mux() itself is what you'll see printed below when it doesn't + // fold: the mux() intrinsic survives untouched through every lowering + // pass that runs before compile_to_module returns, which is the level + // this lesson inspects. It isn't rewritten into an actual select() + // chain until final code generation -- see lower_mux() in + // CodeGen_Internal.cpp, called from within CodeGen_LLVM.cpp/ + // CodeGen_C.cpp -- well past what's printed here.) + { + // Each element's treatment depends on (x - offset) % 4: which of + // 4 cases applies to a given x shifts by 'offset', a Param whose + // value isn't known until the pipeline actually runs. + Var x("x"), xo("xo"), xi("xi"); + Func f("f"); + Param offset("offset"); + offset.set_range(0, 3); + f(x) = mux((x - offset) % 4, + {x, x * x, 2 * x, -x * (x + 1)}); + f.output_buffer().dim(0).set_min(0); + + // An ordinary split(x, xo, xi, 4) makes xi range over [0, 3), + // counting up from x's own min -- which has nothing to do with + // 'offset'. (x - offset) % 4 would still depend on 'offset' after + // substituting in xi, so it can never become a compile-time + // constant, and every mux() call has to stay a mux() call with a + // non-constant selector, no matter what you unroll. See for + // yourself: + Func f0("f0"); + f0(x) = mux((x - offset) % 4, + {x, x * x, 2 * x, -x * (x + 1)}); + f0.output_buffer().dim(0).set_min(0); + f0.split(x, xo, xi, 4, TailStrategy::GuardWithIf) + .unroll(xi); + auto mod0 = f0.compile_to_module({offset}, "f0"); + std::cout << "Unaligned split, unrolled: the mux() is still there, selector not constant\n" + << find_producer(mod0, "f0") << "\n"; + + // split(x, xo, xi, 4, offset) anchors tile boundaries to 'offset' + // instead: x becomes xo*4 + offset + xi, with xi still ranging + // over the plain [0, 4). Substitute that into the mux's selector + // and 'offset' cancels algebraically -- (xo*4 + offset + xi - + // offset) % 4 simplifies to xi % 4 -- with no need for xi to be a + // literal yet. Halide's tail strategy here is GuardWithIf: since + // the split factor doesn't necessarily divide the extent, the + // last partial tile is handled by a separate, guarded copy of the + // loop body rather than by recomputing or overrunning storage. + f.split(x, xo, xi, 4, offset, TailStrategy::GuardWithIf) + // xi % 4 is now a compile-time-provable value in [0, 4), but + // it's still a runtime loop variable, not a literal -- so the + // mux's selector is *known to be one of 4 cases*, but not + // *which* one, and it would still be a mux() call with that + // non-constant selector. Unrolling turns each iteration of xi + // into its own copy of the loop body with xi replaced by a + // literal 0, 1, 2, or 3, which is what finally lets each + // mux() collapse to the single case it selects. + .unroll(xi); + + Module m = f.compile_to_module({offset}); + std::cout << "Aligned split, GuardWithIf, unrolled: every mux() is gone\n" + << find_producer(m, "f") << "\n"; + Counter c = count(m); + if (c.mux_count != 0) { + printf("Expected every mux() to fold away, found %d left\n", c.mux_count); + return 1; + } + + // The pipeline still computes the same thing regardless of the + // runtime value of 'offset' -- aligning the split doesn't change + // the algorithm, only how completely the compiler can simplify + // the code that implements it. + for (int o = 0; o < 4; o++) { + offset.set(o); + Buffer im = f.realize({32}); + for (int x = 0; x < 32; x++) { + int selector = (4 + x - o) % 4; + int expected = (selector == 0) ? x : + (selector == 1) ? x * x : + (selector == 2) ? 2 * x : + -x * (x + 1); + if (im(x) != expected) { + printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, o); + return 1; + } + } + } + + // The tail strategy matters here. GuardWithIf isolates the + // partial last tile behind an explicit boundary check, so the + // steady-state loop body only ever sees whole, uniformly-shifted + // tiles, and every mux() in it folds. ShiftInwards instead slides + // the last tile backwards to keep it in bounds, folding it into + // the same loop as everything else -- which keeps the loop count + // down to one, but that shifted tile is no longer offset from + // 'offset' by a compile-time-constant amount, so the muxes serving + // it can't be resolved at compile time and survive as mux() calls + // with a non-constant selector. + Func f2("f2"); + f2(x) = mux((x - offset) % 4, + {x, x * x, 2 * x, -x * (x + 1)}); + f2.output_buffer().dim(0).set_min(0); + f2.split(x, xo, xi, 4, offset, TailStrategy::ShiftInwards) + .unroll(xi); + Module m2 = f2.compile_to_module({offset}); + std::cout << "Aligned split, ShiftInwards, unrolled: some mux() calls remain\n" + << find_producer(m2, "f2") << "\n"; + Counter c2 = count(m2); + if (c2.mux_count == 0) { + printf("Expected ShiftInwards to leave some muxes behind\n"); + return 1; + } + printf("GuardWithIf: %d for loop(s), %d mux() left. ShiftInwards: %d for loop(s), %d mux() left.\n", + c.for_count, c.mux_count, c2.for_count, c2.mux_count); + } + + // Part 3: two dimensions at once + // + // The same idea applies independently in each dimension: split both x + // and y, anchored to their own runtime offsets. + { + Var x("x"), xo("xo"), xi("xi"); + Var y("y"), yo("yo"), yi("yi"); + Func f("f"); + Param offset_x("offset_x"), offset_y("offset_y"); + offset_x.set_range(0, 1); + offset_y.set_range(0, 1); + Expr selector = 2 * ((y - offset_y) % 2) + (x - offset_x) % 2; + f(x, y) = mux(selector, {x * x, x * y, y * y, x + y}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + + f.split(x, xo, xi, 2, offset_x, TailStrategy::GuardWithIf) + .split(y, yo, yi, 2, offset_y, TailStrategy::GuardWithIf) + // Both tail strategies here already guard their own partial + // tile explicitly, so there's nothing for the loop partitioner + // to usefully split further. Turning it off keeps the compiled + // code -- and the for-loop count checked below -- simple and + // predictable. + .never_partition_all() + // xi and yi need to be innermost, and unrolled, for the same + // reason as Part 2: that's what turns them into the literals + // that let the mux's selector become a compile-time constant. + .reorder(xi, yi, xo, yo) + .unroll(xi) + .unroll(yi); + + Module m = f.compile_to_module({offset_x, offset_y}); + std::cout << "2D aligned split, both dims unrolled: no mux() left\n" + << find_producer(m, "f") << "\n"; + Counter c = count(m); + if (c.mux_count != 0) { + printf("Expected every mux() to fold away, found %d left\n", c.mux_count); + return 1; + } + printf("2D case: %d for loop(s), %d mux() left.\n", c.for_count, c.mux_count); + } + + // Part 4: compute_at reintroduces the problem + // + // The mux-folding trick above relies on the unrolled Var being + // directly related to the aligned split's own outer Var by a + // compile-time-constant offset. That relationship doesn't survive + // crossing a compute_at boundary for free: a Func computed at some + // other Func's loop gets its own loop nest, with its own bounds, sized + // by bounds inference -- and by default, those bounds are an interval + // expression in terms of the consumer's loop variables and Params, not + // a literal constant. + { + Var c("c"); + Var x("x"), xo("xo"), xi("xi"); + Var y("y"), yo("yo"), yi("yi"); + Func f("f"), R("R"), G("G"), B("B"); + Param offset_x("offset_x"), offset_y("offset_y"); + offset_x.set_range(0, 2); + offset_y.set_range(0, 2); + Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; + std::vector ways; + for (int i = 0; i < 9; i++) { + ways.push_back(x * (i % 3) * 3 + y * (i % 3)); + } + R(x, y) = mux(selector, ways); + G(x, y) = mux(selector, ways); + B(x, y) = mux(selector, ways); + // f itself picks between the three channels with a second mux, + // indexed by c -- but c's range is a plain compile-time constant + // set by .bound(), so that one folds regardless of anything to do + // with alignment. + f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + f.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) + .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(c, xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .bound(c, 0, 3) + .unroll(c); + + // Compute each channel per tile of f, and try to unroll its x, y + // the same way we unrolled xi, yi above: + for (Func *channel : {&R, &G, &B}) { + channel->compute_at(f, xo) + .never_partition_all() + .unroll(x) + .unroll(y); + } + + // This doesn't even compile. Bounds inference gives R, G, and B's + // own x loop an extent like + // "min(xo*3 + offset_x + 3, f.extent.0) - max(xo*3 + offset_x, 0)" + // -- a runtime expression, not a constant -- because by default it + // only knows the *region required* by this tile of f, not that + // it's exactly 3 wide. unroll() requires a compile-time-constant + // extent, so it fails outright, before the question of whether the + // muxes fold even comes up. + bool got_expected_error = false; + try { + f.compile_to_module({offset_x, offset_y}); + } catch (const Halide::Error &) { + got_expected_error = true; + } + if (!got_expected_error) { + printf("Expected compiling this schedule to fail\n"); + return 1; + } + printf("As expected, compute_at without a fixed extent can't be unrolled.\n"); + + // Fixing just the crash isn't the same as fixing the muxes. Pin the + // computed and allocated size of each tile with bound_extent() and + // bound_storage() -- enough to make unroll() legal again -- but + // stop there, without telling bounds inference anything about + // *where* that tile starts relative to offset_x/offset_y: + Func f2("f2"), R2("R2"), G2("G2"), B2("B2"); + R2(x, y) = mux(selector, ways); + G2(x, y) = mux(selector, ways); + B2(x, y) = mux(selector, ways); + f2(x, y, c) = mux(c, {R2(x, y), G2(x, y), B2(x, y)}); + f2.output_buffer().dim(0).set_min(0); + f2.output_buffer().dim(1).set_min(0); + f2.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) + .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(c, xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .bound(c, 0, 3) + .unroll(c); + for (Func *channel : {&R2, &G2, &B2}) { + channel->compute_at(f2, xo) + .never_partition_all() + .bound_storage(x, 3) + .bound_extent(x, 3) + .bound_storage(y, 3) + .bound_extent(y, 3) + .unroll(x) + .unroll(y); + } + Module m4 = f2.compile_to_module({offset_x, offset_y}); + std::cout << "compute_at with a fixed tile size, but no align_bounds: " + "the muxes are still there\n" + << find_producer(m4, "B2") << "\n"; + Counter c4 = count(m4); + if (c4.mux_count == 0) { + printf("Expected some mux() calls to survive without align_bounds\n"); + return 1; + } + printf("Fixed-size tile, no align_bounds: %d mux() left.\n", c4.mux_count); + } + + // Part 5: circumventing it + // + // bound_extent()/bound_storage() alone got the schedule to compile, but + // the printed IR above still shows mux() calls inside B2: pinning the + // *size* of a tile doesn't tell bounds inference *where* it starts + // relative to offset_x/offset_y, so x and y inside R/G/B's own + // definition aren't provably a compile-time-constant distance from + // offset_x/offset_y the way xi/yi were in Parts 2-3. Func::align_bounds + // is what supplies that: it constrains a Func's computed min to be + // congruent to a given remainder modulo a given modulus -- exactly the + // "phase" information the outer aligned split already has, but that + // doesn't cross the compute_at boundary on its own. + { + Var c("c"); + Var x("x"), xo("xo"), xi("xi"); + Var y("y"), yo("yo"), yi("yi"); + Func f("f"), R("R"), G("G"), B("B"); + Param offset_x("offset_x"), offset_y("offset_y"); + offset_x.set_range(0, 2); + offset_y.set_range(0, 2); + Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; + std::vector ways; + for (int i = 0; i < 9; i++) { + ways.push_back(x * (i % 3) * 3 + y * (i % 3)); + } + R(x, y) = mux(selector, ways); + G(x, y) = mux(selector, ways); + B(x, y) = mux(selector, ways); + f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); + f.output_buffer().dim(0).set_min(0); + f.output_buffer().dim(1).set_min(0); + f.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) + .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(c, xi, yi, xo, yo) + .unroll(xi) + .unroll(yi) + .bound(c, 0, 3) + .unroll(c); + + for (Func *channel : {&R, &G, &B}) { + channel->compute_at(f, xo) + .never_partition_all() + // Fix the computed and allocated size of this tile... + .bound_storage(x, 3) + .bound_extent(x, 3) + // ...*then* tell bounds inference which phase that tile is + // anchored to, matching the outer split above. Order + // matters: bound_extent()/bound_storage() need to run + // first. Do it the other way around -- + // .align_bounds(x, 3, offset_x).bound_extent(x, 3) -- and + // bounds inference derives too small a region for this + // stage (a hard bounds-checking failure at run time, e.g. + // "Bounds given for B in x (from -2 to 0) do not cover + // required region (from -2 to 3)"), rather than merely + // failing to simplify. Getting the order right is a + // correctness requirement here, not just a simplifier + // nicety. + .align_bounds(x, 3, offset_x) + .bound_storage(y, 3) + .bound_extent(y, 3) + .align_bounds(y, 3, offset_y) + // With the tile's region pinned to a known size and a + // known phase relative to offset_x/offset_y, x and y + // inside R/G/B's own definition are now related to + // offset_x/offset_y by a compile-time-constant difference + // once unrolled -- exactly the relationship Part 2 relied + // on -- so the muxes inside R, G, and B fold too. + .unroll(x) + .unroll(y); + } + + Module m = f.compile_to_module({offset_x, offset_y}); + std::cout << "compute_at with bound_extent/bound_storage *and* align_bounds: " + "clean again\n" + << find_producer(m, "B") << "\n"; + Counter cnt = count(m); + if (cnt.mux_count != 0) { + printf("Expected every mux() to fold away, found %d left\n", cnt.mux_count); + return 1; + } + + // And the pipeline still computes the right answer, for every + // runtime alignment of the 3x3 pattern. + const int W = 16, H = 16; + for (int oy = 0; oy < 3; oy++) { + for (int ox = 0; ox < 3; ox++) { + offset_x.set(ox); + offset_y.set(oy); + Buffer im = f.realize({W, H, 3}); + for (int cc = 0; cc < 3; cc++) { + for (int y = 0; y < H; y++) { + for (int x = 0; x < W; x++) { + // Bias by 3 so the operands of % stay + // non-negative, where C++'s truncated % + // agrees with Halide's Euclidean %. + int sel = (3 * ((3 + y - oy) % 3) + (3 + x - ox) % 3); + int expected = x * (sel % 3) * 3 + y * (sel % 3); + if (im(x, y, cc) != expected) { + printf("im(%d, %d, %d) = %d instead of %d\n", + x, y, cc, im(x, y, cc), expected); + return 1; + } + } + } + } + } + } + printf("Full compute_at case: %d for loop(s), %d mux() left.\n", cnt.for_count, cnt.mux_count); + } + + printf("Success!\n"); + return 0; +} From 1bed6555134998c1171732a483ecc26e88f810c9 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 31 Aug 2026 15:16:10 +0200 Subject: [PATCH 20/29] Guard lesson_25's expected-CompileError catch with HALIDE_WITH_EXCEPTIONS The Makefile build (used by CI's Makefile job) defaults WITH_EXCEPTIONS to unset, so libHalide reports errors via abort() instead of throwing. The unguarded try/catch around the intentionally-failing compute_at schedule never got a chance to catch anything in that configuration, so the process aborted instead of exercising the demonstration. Every other place in the tree that expects a Halide::Error already guards this the same way (see test/error/*.cpp). --- tutorial/lesson_25_aligned_split.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tutorial/lesson_25_aligned_split.cpp b/tutorial/lesson_25_aligned_split.cpp index b5310cb504a7..5724cd5b1be0 100644 --- a/tutorial/lesson_25_aligned_split.cpp +++ b/tutorial/lesson_25_aligned_split.cpp @@ -435,6 +435,13 @@ int main(int argc, char **argv) { // it's exactly 3 wide. unroll() requires a compile-time-constant // extent, so it fails outright, before the question of whether the // muxes fold even comes up. + // + // Compiling this schedule throws a Halide::CompileError. Halide can + // be built with error reporting done via abort() instead of C++ + // exceptions (HALIDE_WITH_EXCEPTIONS undefined, e.g. the top-level + // Makefile's default), so only attempt to catch it when exceptions + // are actually the reporting mechanism in this build. +#ifdef HALIDE_WITH_EXCEPTIONS bool got_expected_error = false; try { f.compile_to_module({offset_x, offset_y}); @@ -446,6 +453,9 @@ int main(int argc, char **argv) { return 1; } printf("As expected, compute_at without a fixed extent can't be unrolled.\n"); +#else + printf("Skipping the expected-compile-failure demonstration (built without exceptions).\n"); +#endif // Fixing just the crash isn't the same as fixing the muxes. Pin the // computed and allocated size of each tile with bound_extent() and From d0495db34768c7ff541b5ca7418718b3ad50a07a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 31 Aug 2026 15:29:45 +0200 Subject: [PATCH 21/29] Fix clang-tidy findings in lesson_25_aligned_split - -> (modernize-deprecated-headers) - main() no longer takes unused argc/argv (misc-unused-parameters) - Explicit `protected:` on the visit() overrides in Counter and FindProducer, matching IRVisitor's own visibility instead of the implicit private (misc-override-with-different-visibility) - .extents[0] -> .extents.at(0) (cppcoreguidelines-pro-bounds-avoid-unchecked-container-access) - Nested ternaries replaced with immediately-invoked if/else lambdas (readability-avoid-nested-conditional-operator) - reserve(9) before the two 9-element ways.push_back() loops (performance-inefficient-vector-operation) --- tutorial/lesson_25_aligned_split.cpp | 39 +++++++++++++++++++++------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/tutorial/lesson_25_aligned_split.cpp b/tutorial/lesson_25_aligned_split.cpp index 5724cd5b1be0..c7756c7854ba 100644 --- a/tutorial/lesson_25_aligned_split.cpp +++ b/tutorial/lesson_25_aligned_split.cpp @@ -18,8 +18,8 @@ // DYLD_LIBRARY_PATH= ./lesson_25 #include "Halide.h" +#include #include -#include using namespace Halide; using namespace Halide::Internal; @@ -33,6 +33,7 @@ namespace { // so this lesson can show what each schedule accomplishes, rather than // just assert it. class Counter : public IRVisitor { +protected: using IRVisitor::visit; void visit(const Call *op) override { IRVisitor::visit(op); @@ -75,6 +76,7 @@ Counter count(const Module &m) { // thing to focus on is the shape of the code inside, not resolving every // hoisted temporary. class FindProducer : public IRVisitor { +protected: using IRVisitor::visit; void visit(const ProducerConsumer *p) override { @@ -108,7 +110,7 @@ Stmt find_producer(const Module &m, const std::string &name) { } // namespace -int main(int argc, char **argv) { +int main() { // Part 1: what does "aligned" actually change? // @@ -188,11 +190,11 @@ int main(int argc, char **argv) { // their own, so the partitioner peels exactly one element off each // end, leaving a steady-state loop over x in [1, 6]: three // unrolled-by-2 iterations, one more than the unaligned version. - if (c_plain.for_count != 1 || !is_const(c_plain.extents[0], 2)) { + if (c_plain.for_count != 1 || !is_const(c_plain.extents.at(0), 2)) { printf("Expected one steady-state loop of extent 2 without alignment\n"); return 1; } - if (c_aligned.for_count != 1 || !is_const(c_aligned.extents[0], 3)) { + if (c_aligned.for_count != 1 || !is_const(c_aligned.extents.at(0), 3)) { printf("Expected one steady-state loop of extent 3 with alignment\n"); return 1; } @@ -202,8 +204,15 @@ int main(int argc, char **argv) { // result. Buffer out = aligned.realize({8}); for (int i = 0; i < 8; i++) { - int expected = (i <= 0) ? 100 : (i < 7) ? i % 2 : - 200; + int expected = [&] { + if (i <= 0) { + return 100; + } + if (i < 7) { + return i % 2; + } + return 200; + }(); if (out(i) != expected) { printf("out(%d) = %d instead of %d\n", i, out(i), expected); return 1; @@ -295,10 +304,18 @@ int main(int argc, char **argv) { Buffer im = f.realize({32}); for (int x = 0; x < 32; x++) { int selector = (4 + x - o) % 4; - int expected = (selector == 0) ? x : - (selector == 1) ? x * x : - (selector == 2) ? 2 * x : - -x * (x + 1); + int expected = [&] { + if (selector == 0) { + return x; + } + if (selector == 1) { + return x * x; + } + if (selector == 2) { + return 2 * x; + } + return -x * (x + 1); + }(); if (im(x) != expected) { printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, o); return 1; @@ -396,6 +413,7 @@ int main(int argc, char **argv) { offset_y.set_range(0, 2); Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; std::vector ways; + ways.reserve(9); for (int i = 0; i < 9; i++) { ways.push_back(x * (i % 3) * 3 + y * (i % 3)); } @@ -521,6 +539,7 @@ int main(int argc, char **argv) { offset_y.set_range(0, 2); Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; std::vector ways; + ways.reserve(9); for (int i = 0; i < 9; i++) { ways.push_back(x * (i % 3) * 3 + y * (i % 3)); } From 88a4398533c66b4eb1cc580ffa01d89b595d209b Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 1 Sep 2026 14:17:59 +0200 Subject: [PATCH 22/29] Feedback from Andrew. --- src/ApplySplit.cpp | 9 ++++++++- src/Func.h | 22 +++++++++++----------- src/Simplify_Add.cpp | 2 +- src/Simplify_Mod.cpp | 8 ++++---- 4 files changed, 24 insertions(+), 17 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 8dc177f896ab..1e6c46ee2132 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -74,13 +74,20 @@ vector apply_split(const Split &split, const string &prefix, // extent divides the factor. Use predication to guard // the calls and/or provides. + // Bounds inference has trouble exploiting an if + // condition. We'll directly tell it that the loop + // variable is bounded above by the original loop max by + // replacing the variable with a promise-clamped version + // of it. Expr guarded; if (split.align.defined()) { // Because the un-rebased base block can start before old_min, // we must clamp both the minimum and maximum boundaries. guarded = promise_clamped(old_var, old_min, old_max); } else { - // Legacy: structurally guaranteed to be >= old_min + // We don't also use the original loop min because + // it needlessly complicates the expressions and doesn't + // actually communicate anything new. guarded = promise_clamped(old_var, old_var, old_max); } diff --git a/src/Func.h b/src/Func.h index d7a7d64fb7b1..dc9dcf62a04e 100644 --- a/src/Func.h +++ b/src/Func.h @@ -1523,17 +1523,17 @@ class Func { Func &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto); /** A version of split() that additionally takes a runtime-valued - * phase, 'align', which need not be known at compile time. An - * ordinary split() tiles a Var starting at its own loop_min, so - * tile boundaries fall at old_min, old_min + factor, old_min + 2 * - * factor, and so on. This variant anchors the tiling to 'align' - * instead, so tile boundaries fall at align, align + factor, align - * - factor, and so on. The inner dimension still iterates over - * [0, factor-1], same as an unaligned split; what changes is which - * value of the original Var each (outer, inner) pair reconstructs - * to. This may increase the number of iterations over the outer - * loop by 1 compared to an unaligned split, since align need not - * coincide with old_min. + * 'align' Expr. This version anchors the inner-loop's iterations + * to absolute coordinates, instead of to Halide's inferred loop + * bounds. + * + * As such, in absolute coordinates, the inner-loop boundaries fall + * at ``align``, ``align + factor``, ``align + 2*factor``, and so on. + * The inner dimension still iterates over ``[0, factor-1]``, same as + * an unaligned split. The difference is how the original loop Var + * is reconstructed from the outer and inner loop Vars. + * This may increase the number of iterations over the outer + * loop by 1 compared to an unaligned split. * * This is useful when an algorithm selects between cases using an * expression like ``(x - offset) % factor``, where 'offset' is a diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index a2298c2e019c..56a6a201b98e 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -201,7 +201,7 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite(x + ((c0 - x) / c1) * c1, c0 - ((c0 - x) % c1), c1 > 0) || rewrite(x + ((c0 - x) / c1 + y) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || rewrite(x + (y + (c0 - x) / c1) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || - rewrite(((0 - x) / c0) + ((x % c0 + c1) / c0), (c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || + rewrite(((0 - x) / c0) + ((x % c0 + c1) / c0), fold(c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || false)))) { return mutate(rewrite.result, info); diff --git a/src/Simplify_Mod.cpp b/src/Simplify_Mod.cpp index 0bbbddb4ec34..04b082483119 100644 --- a/src/Simplify_Mod.cpp +++ b/src/Simplify_Mod.cpp @@ -59,10 +59,10 @@ Expr Simplify::visit(const Mod *op, ExprInfo *info) { rewrite((x * c0 - y) % c1, (-y) % c1, c0 % c1 == 0) || rewrite((y - x * c0) % c1, y % c1, c0 % c1 == 0) || rewrite((x - y) % 2, (x + y) % 2) || // Addition and subtraction are the same modulo 2, because -1 == 1 - rewrite((((x * c0) + y) - z) % c0, (y - z) % c0) || - rewrite((((x * c0) + y) + z) % c0, (y + z) % c0) || - rewrite((((x * c0) - y) - z) % c0, (-y - z) % c0) || - rewrite((((x * c0) - y) + z) % c0, (z - y) % c0) || + rewrite((((x * c0) + y) - z) % c1, (y - z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) + y) + z) % c1, (y + z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) - y) - z) % c1, (-y - z) % c1, c0 % c1 == 0) || + rewrite((((x * c0) - y) + z) % c1, (z - y) % c1, c0 % c1 == 0) || rewrite(ramp(x, c0, c2) % broadcast(c1, c2), broadcast(x, c2) % broadcast(c1, c2), (c0 % c1 == 0)) || rewrite(ramp(x, c0, lanes) % broadcast(c1, lanes), ramp(x % c1, c0, lanes), From f1871f40fceb0caa0a6cfe6781dcb0b47a461579 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 1 Sep 2026 14:46:05 +0200 Subject: [PATCH 23/29] Simplify the bounds calculation of a split. --- src/ApplySplit.cpp | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 1e6c46ee2132..14c77ae772d4 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -283,21 +283,18 @@ vector> compute_loop_bounds_after_split(const Split &spl Expr old_var_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); switch (split.split_type) { case Split::SplitVar: { + let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); + let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 1); if (split.align.defined()) { Expr align = split.align; Expr outer_min = (old_var_min - align) / split.factor; Expr outer_max = (old_var_max - align) / split.factor; - let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); - let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 1); let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min); let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); } else { - Expr inner_extent = split.factor; - Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor; - let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); - let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1); + Expr outer_max = (old_var_max - old_var_min) / split.factor; let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0); - let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1); + let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); } } break; case Split::FuseVars: { From c0a1a9a8f4e6132b22b41356965cbc5863be3f8f Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 1 Sep 2026 14:47:21 +0200 Subject: [PATCH 24/29] Peeling variants of the ramp <= broadcast style rules. Co-Authored-By: Claude Sonnet 5 Co-Authored-By: Andrew Adams --- src/Simplify_Exprs.cpp | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/src/Simplify_Exprs.cpp b/src/Simplify_Exprs.cpp index b3ce824f0cf1..07512e9186d0 100644 --- a/src/Simplify_Exprs.cpp +++ b/src/Simplify_Exprs.cpp @@ -218,19 +218,24 @@ Expr Simplify::visit(const VectorReduce *op, ExprInfo *info) { x <= y + min(z * (arg_lanes - 1), 0)) || // The "all lanes of a ramp lie within [lo, hi]" check loop - // partitioning builds (a lower-bound comparison ANDed with an - // upper-bound comparison, both against the same stride, e.g. - // (0 <= ramp(b0, s, n)) && (ramp(b1, s, n) <= extent)) - rewrite(h_and((broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)) && - (ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)), - 1), - (x <= y + min(z * (arg_lanes - 1), 0)) && - (w + max(z * (arg_lanes - 1), 0) <= u)) || - rewrite(h_and((ramp(w, z, arg_lanes) <= broadcast(u, arg_lanes)) && - (broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)), - 1), - (w + max(z * (arg_lanes - 1), 0) <= u) && - (x <= y + min(z * (arg_lanes - 1), 0))) || + // partitioning builds is a chain of these same ramp/broadcast + // comparisons ANDed together (a lower-bound comparison against a + // ramp's minimum lane, an upper-bound comparison against its + // maximum lane, for one or more ramps sharing a stride). Rather + // than special-case every clause count, peel one comparison off + // the tail of the && chain at a time: the four rules above already + // reduce that comparison to a scalar; h_and(w, 1) on what's left + // recurses into this same case, so it either peels the next + // clause or -- once w is down to a single comparison -- hits one + // of the four rules above directly. + rewrite(h_and(w && (ramp(x, y, arg_lanes) < broadcast(z, arg_lanes)), 1), + h_and(w, 1) && (x + max(y * (arg_lanes - 1), 0) < z)) || + rewrite(h_and(w && (ramp(x, y, arg_lanes) <= broadcast(z, arg_lanes)), 1), + h_and(w, 1) && (x + max(y * (arg_lanes - 1), 0) <= z)) || + rewrite(h_and(w && (broadcast(x, arg_lanes) < ramp(y, z, arg_lanes)), 1), + h_and(w, 1) && (x < y + min(z * (arg_lanes - 1), 0))) || + rewrite(h_and(w && (broadcast(x, arg_lanes) <= ramp(y, z, arg_lanes)), 1), + h_and(w, 1) && (x <= y + min(z * (arg_lanes - 1), 0))) || false) { return mutate(rewrite.result, info); } From 7435de5f0728bce7ee3494bc1797763ca3a63085 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 2 Sep 2026 11:21:35 +0200 Subject: [PATCH 25/29] Remove tutorial. --- tutorial/CMakeLists.txt | 3 +- tutorial/lesson_25_aligned_split.cpp | 634 --------------------------- 2 files changed, 1 insertion(+), 636 deletions(-) delete mode 100644 tutorial/lesson_25_aligned_split.cpp diff --git a/tutorial/CMakeLists.txt b/tutorial/CMakeLists.txt index cd6091bf13d4..94c3d9e22be7 100644 --- a/tutorial/CMakeLists.txt +++ b/tutorial/CMakeLists.txt @@ -341,8 +341,7 @@ if (TARGET Halide::Mullapudi2016) ) endif () -# Lessons 22-25 +# Lessons 22-24 add_tutorial(lesson_22_jit_performance.cpp) add_tutorial(lesson_23_serialization.cpp WITH_IMAGE_IO) add_tutorial(lesson_24_async.cpp GROUPS multithreaded) -add_tutorial(lesson_25_aligned_split.cpp) diff --git a/tutorial/lesson_25_aligned_split.cpp b/tutorial/lesson_25_aligned_split.cpp deleted file mode 100644 index c7756c7854ba..000000000000 --- a/tutorial/lesson_25_aligned_split.cpp +++ /dev/null @@ -1,634 +0,0 @@ -// Halide tutorial lesson 25: Aligned splits - -// This lesson demonstrates Func::split's 'align' parameter: a way to anchor -// a split's tile boundaries to a value that's only known at runtime (e.g. a -// Param), instead of to the Var's own min. It shows what that buys you -- -// primarily, letting mux() calls whose selector depends on that same -// runtime value collapse to the single case they select, instead of -// surviving as mux() calls whose case can only be picked at run time -- -// and where the trick currently runs out of steam: crossing a compute_at -// boundary. - -// On linux, you can compile and run it like so: -// g++ lesson_25*.cpp -g -I -L -lHalide -lpthread -ldl -o lesson_25 -std=c++17 -// LD_LIBRARY_PATH= ./lesson_25 - -// On macOS: -// g++ lesson_25*.cpp -g -I -L -lHalide -o lesson_25 -std=c++17 -// DYLD_LIBRARY_PATH= ./lesson_25 - -#include "Halide.h" -#include -#include - -using namespace Halide; -using namespace Halide::Internal; - -namespace { - -// A small IRVisitor used throughout this lesson to check what a schedule -// actually did to the compiled code: how many for loops remain, how wide -// they are, and how many mux() and % operations the simplifier managed to -// fold away. Ordinary Halide code never needs to do this -- it's only here -// so this lesson can show what each schedule accomplishes, rather than -// just assert it. -class Counter : public IRVisitor { -protected: - using IRVisitor::visit; - void visit(const Call *op) override { - IRVisitor::visit(op); - if (op->is_intrinsic(Call::IntrinsicOp::mux)) { - mux_count++; - } - } - void visit(const Mod *op) override { - IRVisitor::visit(op); - mod_count++; - } - void visit(const For *op) override { - extents.push_back(simplify(op->extent())); - IRVisitor::visit(op); - for_count++; - } - -public: - int for_count = 0, mux_count = 0, mod_count = 0; - std::vector extents; -}; - -Counter count(const Module &m) { - Counter c; - for (const LoweredFunc &lf : m.functions()) { - lf.body.accept(&c); - } - return c; -} - -// Finds the "produce { ... }" node for a given Func in a compiled -// Module, so this lesson can print the actual generated code for just that -// Func instead of the whole pipeline. This is the most direct way to see -// what a schedule bought you: read the loop nest and check for yourself -// whether a mux() call turned into a straight-line value or is still there -// as a mux() call with a non-constant selector. Since only that one -// subtree is printed, in a compute_at example you may see it reference a -// variable (e.g. a "let t123 = ...") that's bound just outside the -// printed excerpt, in the surrounding loop nest -- that's expected; the -// thing to focus on is the shape of the code inside, not resolving every -// hoisted temporary. -class FindProducer : public IRVisitor { -protected: - using IRVisitor::visit; - - void visit(const ProducerConsumer *p) override { - // When a Func's own name collides with the compiled Module's - // exported function name, Halide disambiguates the internal one by - // appending "$N" (Halide's general name-uniquification suffix), so - // match that too rather than only an exact name. - if (p->is_producer && (p->name == name || starts_with(p->name, name + "$"))) { - producer = p; - } - IRVisitor::visit(p); - } - -public: - FindProducer(std::string name) : name(std::move(name)) {} - - std::string name; - Stmt producer; -}; - -Stmt find_producer(const Module &m, const std::string &name) { - FindProducer f{name}; - for (const LoweredFunc &lf : m.functions()) { - lf.body.accept(&f); - if (f.producer.defined()) { - return f.producer; - } - } - return Stmt{}; -} - -} // namespace - -int main() { - - // Part 1: what does "aligned" actually change? - // - // An ordinary split(x, xo, xi, factor) tiles a Var starting at its own - // loop_min: tile boundaries fall at old_min, old_min + factor, - // old_min + 2*factor, and so on. split(x, xo, xi, factor, align) - // anchors the tiling to 'align' instead, wherever that falls relative - // to old_min. 'align' need not be known at compile time -- that's the - // whole point of the feature -- but to see *why* the anchor point - // matters at all, this first example uses a compile-time-constant - // align of 1, and looks at how it interacts with loop partitioning. - { - // The pipeline below is 8 elements: a boundary value at x == 0, a - // boundary value at x == 7, and a period-2 pattern (x % 2) for the - // six elements in between. likely() marks the interior case as the - // expected one, which is what invites the loop partitioner to try - // to peel the boundary cases into their own prologue/epilogue - // rather than testing for them on every iteration. - // - // First, split by 2 with no alignment at all, and unroll the - // 2-wide inner loop so that x % 2 has a chance to become a - // compile-time constant in each unrolled copy. - Var x("x"), xo("xo"), xi("xi"); - Func plain("plain"); - plain(x) = select(x <= 0, 100, - x < 7, likely(x % 2), - 200); - plain.bound(x, 0, 8); - plain.split(x, xo, xi, 2, TailStrategy::GuardWithIf) - .always_partition(xo) - .unroll(xi); - auto mod_plain = plain.compile_to_module({}, "plain"); - std::cout << "Without alignment, tiled from x == 0:\n" - << find_producer(mod_plain, "plain") << "\n"; - Counter c_plain = count(mod_plain); - - // Now the same thing, but anchored to x == 1 -- where the - // interior's period-2 pattern actually begins -- instead of to - // the Func's own min of 0. (Func's assignment operator is a - // reference, not a clone, so this needs its own Func and Vars, - // not a copy of 'plain'.) - Var x2("x"), xo2("xo"), xi2("xi"); - Func aligned("aligned"); - aligned(x2) = select(x2 <= 0, 100, - x2 < 7, likely(x2 % 2), - 200); - aligned.bound(x2, 0, 8); - aligned.split(x2, xo2, xi2, 2, /* align */ 1, TailStrategy::GuardWithIf) - .always_partition(xo2) - .unroll(xi2); - auto mod_aligned = aligned.compile_to_module({}, "aligned"); - std::cout << "Anchored at x == 1:\n" - << find_producer(mod_aligned, "aligned") << "\n"; - Counter c_aligned = count(mod_aligned); - - // Either way, x % 2 folds away completely: every unrolled - // instance of xi turns (xo*2 + align + xi) % 2 into a compile-time - // constant, regardless of what 'align' is. (Try it: the multiple - // of 2 drops out of the modulo either way.) So this isn't about - // the modulo -- it's about how much of the *boundary* the - // partitioner can peel into single, non-looping iterations. - if (c_plain.mod_count != 0 || c_aligned.mod_count != 0) { - printf("Expected the %% to fold away in both cases\n"); - return 1; - } - - // Without alignment, tiles start at x == 0, so the single-element - // boundary case at x == 0 *is* a whole tile, and the one at x == 7 - // spills into the tile that starts at x == 6. The partitioner has - // to peel a full 2-element tile off of each end to isolate the - // boundary, leaving a steady-state loop that only covers x in - // [2, 5]: two unrolled-by-2 iterations. - // - // Anchored at x == 1, tiles start at 1, 3, 5: exactly where the - // interior's period-2 pattern repeats. Now the single-element - // boundary cases at x == 0 and x == 7 are each a partial tile of - // their own, so the partitioner peels exactly one element off each - // end, leaving a steady-state loop over x in [1, 6]: three - // unrolled-by-2 iterations, one more than the unaligned version. - if (c_plain.for_count != 1 || !is_const(c_plain.extents.at(0), 2)) { - printf("Expected one steady-state loop of extent 2 without alignment\n"); - return 1; - } - if (c_aligned.for_count != 1 || !is_const(c_aligned.extents.at(0), 3)) { - printf("Expected one steady-state loop of extent 3 with alignment\n"); - return 1; - } - - // Both schedules compute the same answer either way -- alignment - // only changes how the work is split into loops, never the - // result. - Buffer out = aligned.realize({8}); - for (int i = 0; i < 8; i++) { - int expected = [&] { - if (i <= 0) { - return 100; - } - if (i < 7) { - return i % 2; - } - return 200; - }(); - if (out(i) != expected) { - printf("out(%d) = %d instead of %d\n", i, out(i), expected); - return 1; - } - } - } - - // Part 2: anchoring a split to a runtime Param - // - // Lesson 13 pointed out that mux(id, {...}) is sugar for a select - // chain, and that the select can be compiled away by bounding and - // unrolling the Var it's indexed by -- but only if 'id' becomes a - // compile-time constant in each unrolled copy. That's easy when 'id' - // is directly the unrolled Var. It's harder when 'id' depends on that - // Var only after subtracting off a value that's not known until - // runtime, as below. - // - // (mux() itself is what you'll see printed below when it doesn't - // fold: the mux() intrinsic survives untouched through every lowering - // pass that runs before compile_to_module returns, which is the level - // this lesson inspects. It isn't rewritten into an actual select() - // chain until final code generation -- see lower_mux() in - // CodeGen_Internal.cpp, called from within CodeGen_LLVM.cpp/ - // CodeGen_C.cpp -- well past what's printed here.) - { - // Each element's treatment depends on (x - offset) % 4: which of - // 4 cases applies to a given x shifts by 'offset', a Param whose - // value isn't known until the pipeline actually runs. - Var x("x"), xo("xo"), xi("xi"); - Func f("f"); - Param offset("offset"); - offset.set_range(0, 3); - f(x) = mux((x - offset) % 4, - {x, x * x, 2 * x, -x * (x + 1)}); - f.output_buffer().dim(0).set_min(0); - - // An ordinary split(x, xo, xi, 4) makes xi range over [0, 3), - // counting up from x's own min -- which has nothing to do with - // 'offset'. (x - offset) % 4 would still depend on 'offset' after - // substituting in xi, so it can never become a compile-time - // constant, and every mux() call has to stay a mux() call with a - // non-constant selector, no matter what you unroll. See for - // yourself: - Func f0("f0"); - f0(x) = mux((x - offset) % 4, - {x, x * x, 2 * x, -x * (x + 1)}); - f0.output_buffer().dim(0).set_min(0); - f0.split(x, xo, xi, 4, TailStrategy::GuardWithIf) - .unroll(xi); - auto mod0 = f0.compile_to_module({offset}, "f0"); - std::cout << "Unaligned split, unrolled: the mux() is still there, selector not constant\n" - << find_producer(mod0, "f0") << "\n"; - - // split(x, xo, xi, 4, offset) anchors tile boundaries to 'offset' - // instead: x becomes xo*4 + offset + xi, with xi still ranging - // over the plain [0, 4). Substitute that into the mux's selector - // and 'offset' cancels algebraically -- (xo*4 + offset + xi - - // offset) % 4 simplifies to xi % 4 -- with no need for xi to be a - // literal yet. Halide's tail strategy here is GuardWithIf: since - // the split factor doesn't necessarily divide the extent, the - // last partial tile is handled by a separate, guarded copy of the - // loop body rather than by recomputing or overrunning storage. - f.split(x, xo, xi, 4, offset, TailStrategy::GuardWithIf) - // xi % 4 is now a compile-time-provable value in [0, 4), but - // it's still a runtime loop variable, not a literal -- so the - // mux's selector is *known to be one of 4 cases*, but not - // *which* one, and it would still be a mux() call with that - // non-constant selector. Unrolling turns each iteration of xi - // into its own copy of the loop body with xi replaced by a - // literal 0, 1, 2, or 3, which is what finally lets each - // mux() collapse to the single case it selects. - .unroll(xi); - - Module m = f.compile_to_module({offset}); - std::cout << "Aligned split, GuardWithIf, unrolled: every mux() is gone\n" - << find_producer(m, "f") << "\n"; - Counter c = count(m); - if (c.mux_count != 0) { - printf("Expected every mux() to fold away, found %d left\n", c.mux_count); - return 1; - } - - // The pipeline still computes the same thing regardless of the - // runtime value of 'offset' -- aligning the split doesn't change - // the algorithm, only how completely the compiler can simplify - // the code that implements it. - for (int o = 0; o < 4; o++) { - offset.set(o); - Buffer im = f.realize({32}); - for (int x = 0; x < 32; x++) { - int selector = (4 + x - o) % 4; - int expected = [&] { - if (selector == 0) { - return x; - } - if (selector == 1) { - return x * x; - } - if (selector == 2) { - return 2 * x; - } - return -x * (x + 1); - }(); - if (im(x) != expected) { - printf("im(%d) = %d instead of %d (offset: %d)\n", x, im(x), expected, o); - return 1; - } - } - } - - // The tail strategy matters here. GuardWithIf isolates the - // partial last tile behind an explicit boundary check, so the - // steady-state loop body only ever sees whole, uniformly-shifted - // tiles, and every mux() in it folds. ShiftInwards instead slides - // the last tile backwards to keep it in bounds, folding it into - // the same loop as everything else -- which keeps the loop count - // down to one, but that shifted tile is no longer offset from - // 'offset' by a compile-time-constant amount, so the muxes serving - // it can't be resolved at compile time and survive as mux() calls - // with a non-constant selector. - Func f2("f2"); - f2(x) = mux((x - offset) % 4, - {x, x * x, 2 * x, -x * (x + 1)}); - f2.output_buffer().dim(0).set_min(0); - f2.split(x, xo, xi, 4, offset, TailStrategy::ShiftInwards) - .unroll(xi); - Module m2 = f2.compile_to_module({offset}); - std::cout << "Aligned split, ShiftInwards, unrolled: some mux() calls remain\n" - << find_producer(m2, "f2") << "\n"; - Counter c2 = count(m2); - if (c2.mux_count == 0) { - printf("Expected ShiftInwards to leave some muxes behind\n"); - return 1; - } - printf("GuardWithIf: %d for loop(s), %d mux() left. ShiftInwards: %d for loop(s), %d mux() left.\n", - c.for_count, c.mux_count, c2.for_count, c2.mux_count); - } - - // Part 3: two dimensions at once - // - // The same idea applies independently in each dimension: split both x - // and y, anchored to their own runtime offsets. - { - Var x("x"), xo("xo"), xi("xi"); - Var y("y"), yo("yo"), yi("yi"); - Func f("f"); - Param offset_x("offset_x"), offset_y("offset_y"); - offset_x.set_range(0, 1); - offset_y.set_range(0, 1); - Expr selector = 2 * ((y - offset_y) % 2) + (x - offset_x) % 2; - f(x, y) = mux(selector, {x * x, x * y, y * y, x + y}); - f.output_buffer().dim(0).set_min(0); - f.output_buffer().dim(1).set_min(0); - - f.split(x, xo, xi, 2, offset_x, TailStrategy::GuardWithIf) - .split(y, yo, yi, 2, offset_y, TailStrategy::GuardWithIf) - // Both tail strategies here already guard their own partial - // tile explicitly, so there's nothing for the loop partitioner - // to usefully split further. Turning it off keeps the compiled - // code -- and the for-loop count checked below -- simple and - // predictable. - .never_partition_all() - // xi and yi need to be innermost, and unrolled, for the same - // reason as Part 2: that's what turns them into the literals - // that let the mux's selector become a compile-time constant. - .reorder(xi, yi, xo, yo) - .unroll(xi) - .unroll(yi); - - Module m = f.compile_to_module({offset_x, offset_y}); - std::cout << "2D aligned split, both dims unrolled: no mux() left\n" - << find_producer(m, "f") << "\n"; - Counter c = count(m); - if (c.mux_count != 0) { - printf("Expected every mux() to fold away, found %d left\n", c.mux_count); - return 1; - } - printf("2D case: %d for loop(s), %d mux() left.\n", c.for_count, c.mux_count); - } - - // Part 4: compute_at reintroduces the problem - // - // The mux-folding trick above relies on the unrolled Var being - // directly related to the aligned split's own outer Var by a - // compile-time-constant offset. That relationship doesn't survive - // crossing a compute_at boundary for free: a Func computed at some - // other Func's loop gets its own loop nest, with its own bounds, sized - // by bounds inference -- and by default, those bounds are an interval - // expression in terms of the consumer's loop variables and Params, not - // a literal constant. - { - Var c("c"); - Var x("x"), xo("xo"), xi("xi"); - Var y("y"), yo("yo"), yi("yi"); - Func f("f"), R("R"), G("G"), B("B"); - Param offset_x("offset_x"), offset_y("offset_y"); - offset_x.set_range(0, 2); - offset_y.set_range(0, 2); - Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; - std::vector ways; - ways.reserve(9); - for (int i = 0; i < 9; i++) { - ways.push_back(x * (i % 3) * 3 + y * (i % 3)); - } - R(x, y) = mux(selector, ways); - G(x, y) = mux(selector, ways); - B(x, y) = mux(selector, ways); - // f itself picks between the three channels with a second mux, - // indexed by c -- but c's range is a plain compile-time constant - // set by .bound(), so that one folds regardless of anything to do - // with alignment. - f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); - f.output_buffer().dim(0).set_min(0); - f.output_buffer().dim(1).set_min(0); - f.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) - .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) - .never_partition_all() - .reorder(c, xi, yi, xo, yo) - .unroll(xi) - .unroll(yi) - .bound(c, 0, 3) - .unroll(c); - - // Compute each channel per tile of f, and try to unroll its x, y - // the same way we unrolled xi, yi above: - for (Func *channel : {&R, &G, &B}) { - channel->compute_at(f, xo) - .never_partition_all() - .unroll(x) - .unroll(y); - } - - // This doesn't even compile. Bounds inference gives R, G, and B's - // own x loop an extent like - // "min(xo*3 + offset_x + 3, f.extent.0) - max(xo*3 + offset_x, 0)" - // -- a runtime expression, not a constant -- because by default it - // only knows the *region required* by this tile of f, not that - // it's exactly 3 wide. unroll() requires a compile-time-constant - // extent, so it fails outright, before the question of whether the - // muxes fold even comes up. - // - // Compiling this schedule throws a Halide::CompileError. Halide can - // be built with error reporting done via abort() instead of C++ - // exceptions (HALIDE_WITH_EXCEPTIONS undefined, e.g. the top-level - // Makefile's default), so only attempt to catch it when exceptions - // are actually the reporting mechanism in this build. -#ifdef HALIDE_WITH_EXCEPTIONS - bool got_expected_error = false; - try { - f.compile_to_module({offset_x, offset_y}); - } catch (const Halide::Error &) { - got_expected_error = true; - } - if (!got_expected_error) { - printf("Expected compiling this schedule to fail\n"); - return 1; - } - printf("As expected, compute_at without a fixed extent can't be unrolled.\n"); -#else - printf("Skipping the expected-compile-failure demonstration (built without exceptions).\n"); -#endif - - // Fixing just the crash isn't the same as fixing the muxes. Pin the - // computed and allocated size of each tile with bound_extent() and - // bound_storage() -- enough to make unroll() legal again -- but - // stop there, without telling bounds inference anything about - // *where* that tile starts relative to offset_x/offset_y: - Func f2("f2"), R2("R2"), G2("G2"), B2("B2"); - R2(x, y) = mux(selector, ways); - G2(x, y) = mux(selector, ways); - B2(x, y) = mux(selector, ways); - f2(x, y, c) = mux(c, {R2(x, y), G2(x, y), B2(x, y)}); - f2.output_buffer().dim(0).set_min(0); - f2.output_buffer().dim(1).set_min(0); - f2.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) - .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) - .never_partition_all() - .reorder(c, xi, yi, xo, yo) - .unroll(xi) - .unroll(yi) - .bound(c, 0, 3) - .unroll(c); - for (Func *channel : {&R2, &G2, &B2}) { - channel->compute_at(f2, xo) - .never_partition_all() - .bound_storage(x, 3) - .bound_extent(x, 3) - .bound_storage(y, 3) - .bound_extent(y, 3) - .unroll(x) - .unroll(y); - } - Module m4 = f2.compile_to_module({offset_x, offset_y}); - std::cout << "compute_at with a fixed tile size, but no align_bounds: " - "the muxes are still there\n" - << find_producer(m4, "B2") << "\n"; - Counter c4 = count(m4); - if (c4.mux_count == 0) { - printf("Expected some mux() calls to survive without align_bounds\n"); - return 1; - } - printf("Fixed-size tile, no align_bounds: %d mux() left.\n", c4.mux_count); - } - - // Part 5: circumventing it - // - // bound_extent()/bound_storage() alone got the schedule to compile, but - // the printed IR above still shows mux() calls inside B2: pinning the - // *size* of a tile doesn't tell bounds inference *where* it starts - // relative to offset_x/offset_y, so x and y inside R/G/B's own - // definition aren't provably a compile-time-constant distance from - // offset_x/offset_y the way xi/yi were in Parts 2-3. Func::align_bounds - // is what supplies that: it constrains a Func's computed min to be - // congruent to a given remainder modulo a given modulus -- exactly the - // "phase" information the outer aligned split already has, but that - // doesn't cross the compute_at boundary on its own. - { - Var c("c"); - Var x("x"), xo("xo"), xi("xi"); - Var y("y"), yo("yo"), yi("yi"); - Func f("f"), R("R"), G("G"), B("B"); - Param offset_x("offset_x"), offset_y("offset_y"); - offset_x.set_range(0, 2); - offset_y.set_range(0, 2); - Expr selector = 3 * ((y - offset_y) % 3) + (x - offset_x) % 3; - std::vector ways; - ways.reserve(9); - for (int i = 0; i < 9; i++) { - ways.push_back(x * (i % 3) * 3 + y * (i % 3)); - } - R(x, y) = mux(selector, ways); - G(x, y) = mux(selector, ways); - B(x, y) = mux(selector, ways); - f(x, y, c) = mux(c, {R(x, y), G(x, y), B(x, y)}); - f.output_buffer().dim(0).set_min(0); - f.output_buffer().dim(1).set_min(0); - f.split(x, xo, xi, 3, offset_x, TailStrategy::GuardWithIf) - .split(y, yo, yi, 3, offset_y, TailStrategy::GuardWithIf) - .never_partition_all() - .reorder(c, xi, yi, xo, yo) - .unroll(xi) - .unroll(yi) - .bound(c, 0, 3) - .unroll(c); - - for (Func *channel : {&R, &G, &B}) { - channel->compute_at(f, xo) - .never_partition_all() - // Fix the computed and allocated size of this tile... - .bound_storage(x, 3) - .bound_extent(x, 3) - // ...*then* tell bounds inference which phase that tile is - // anchored to, matching the outer split above. Order - // matters: bound_extent()/bound_storage() need to run - // first. Do it the other way around -- - // .align_bounds(x, 3, offset_x).bound_extent(x, 3) -- and - // bounds inference derives too small a region for this - // stage (a hard bounds-checking failure at run time, e.g. - // "Bounds given for B in x (from -2 to 0) do not cover - // required region (from -2 to 3)"), rather than merely - // failing to simplify. Getting the order right is a - // correctness requirement here, not just a simplifier - // nicety. - .align_bounds(x, 3, offset_x) - .bound_storage(y, 3) - .bound_extent(y, 3) - .align_bounds(y, 3, offset_y) - // With the tile's region pinned to a known size and a - // known phase relative to offset_x/offset_y, x and y - // inside R/G/B's own definition are now related to - // offset_x/offset_y by a compile-time-constant difference - // once unrolled -- exactly the relationship Part 2 relied - // on -- so the muxes inside R, G, and B fold too. - .unroll(x) - .unroll(y); - } - - Module m = f.compile_to_module({offset_x, offset_y}); - std::cout << "compute_at with bound_extent/bound_storage *and* align_bounds: " - "clean again\n" - << find_producer(m, "B") << "\n"; - Counter cnt = count(m); - if (cnt.mux_count != 0) { - printf("Expected every mux() to fold away, found %d left\n", cnt.mux_count); - return 1; - } - - // And the pipeline still computes the right answer, for every - // runtime alignment of the 3x3 pattern. - const int W = 16, H = 16; - for (int oy = 0; oy < 3; oy++) { - for (int ox = 0; ox < 3; ox++) { - offset_x.set(ox); - offset_y.set(oy); - Buffer im = f.realize({W, H, 3}); - for (int cc = 0; cc < 3; cc++) { - for (int y = 0; y < H; y++) { - for (int x = 0; x < W; x++) { - // Bias by 3 so the operands of % stay - // non-negative, where C++'s truncated % - // agrees with Halide's Euclidean %. - int sel = (3 * ((3 + y - oy) % 3) + (3 + x - ox) % 3); - int expected = x * (sel % 3) * 3 + y * (sel % 3); - if (im(x, y, cc) != expected) { - printf("im(%d, %d, %d) = %d instead of %d\n", - x, y, cc, im(x, y, cc), expected); - return 1; - } - } - } - } - } - } - printf("Full compute_at case: %d for loop(s), %d mux() left.\n", cnt.for_count, cnt.mux_count); - } - - printf("Success!\n"); - return 0; -} From f13a94d1b8aa767186c19b980a8428629a11c478 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 3 Sep 2026 19:40:03 +0200 Subject: [PATCH 26/29] WIP checkpoint before substitution-based apply_split rewrite Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- src/ApplySplit.cpp | 159 ++++---------------------- src/SimplifyCorrelatedDifferences.cpp | 7 ++ src/Simplify_Add.cpp | 2 + src/Simplify_Div.cpp | 3 + src/Simplify_Mod.cpp | 4 + 5 files changed, 40 insertions(+), 135 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 14c77ae772d4..ab2542f77660 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -23,17 +23,10 @@ vector apply_split(const Split &split, const string &prefix, Expr old_max = Variable::make(Int(32), prefix + split.old_var + ".loop_max"); Expr old_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); Expr old_extent = (old_max - old_min) + 1; - Expr outer_min = Variable::make(Int(32), prefix + split.outer + ".loop_min"); dim_extent_alignment[split.inner] = split.factor; - Expr base; - if (split.align.defined()) { - base = outer * split.factor; - } else { - base = outer * split.factor + old_min; - } - + Expr base = outer * split.factor + old_min; string base_name = prefix + split.inner + ".base"; Expr base_var = Variable::make(Int(32), base_name); string old_var_name = prefix + split.old_var; @@ -45,17 +38,8 @@ vector apply_split(const Split &split, const string &prefix, internal_assert(tail != TailStrategy::Auto) << "An explicit tail strategy should exist at this point\n"; - // When align is defined, tiles are anchored to align instead of to - // old_min, so knowing that the factor divides the extent is not - // enough to prove no boundary guard is needed: we additionally need - // the tiling anchored at align to line up with the tiling anchored - // at old_min, i.e. old_min and align must be congruent mod factor. - bool alignment_matches_old_min = !split.align.defined() || - is_const_zero(simplify((old_min - split.align) % split.factor)); - if ((iter != dim_extent_alignment.end()) && - is_const_zero(simplify(iter->second % split.factor)) && - alignment_matches_old_min) { + is_const_zero(simplify(iter->second % split.factor))) { // We have proved that the split factor divides the // old extent. No need to adjust the base or add an if // statement. @@ -78,19 +62,10 @@ vector apply_split(const Split &split, const string &prefix, // condition. We'll directly tell it that the loop // variable is bounded above by the original loop max by // replacing the variable with a promise-clamped version - // of it. - Expr guarded; - if (split.align.defined()) { - // Because the un-rebased base block can start before old_min, - // we must clamp both the minimum and maximum boundaries. - guarded = promise_clamped(old_var, old_min, old_max); - } else { - // We don't also use the original loop min because - // it needlessly complicates the expressions and doesn't - // actually communicate anything new. - guarded = promise_clamped(old_var, old_var, old_max); - } - + // of it. We don't also use the original loop min because + // it needlessly complicates the expressions and doesn't + // actually communicate anything new. + Expr guarded = promise_clamped(old_var, old_var, old_max); string guarded_var_name = prefix + split.old_var + ".guarded"; Expr guarded_var = Variable::make(Int(32), guarded_var_name); @@ -101,6 +76,8 @@ vector apply_split(const Split &split, const string &prefix, predicate_type = ApplySplitResult::Predicate; break; case TailStrategy::Predicate: + // This is identical to GuardWithIf, but maybe it makes + // sense to keep it anyways? substitution_type = ApplySplitResult::Substitution; predicate_type = ApplySplitResult::Predicate; break; @@ -120,123 +97,36 @@ vector apply_split(const Split &split, const string &prefix, // for the guarded version. result.emplace_back(prefix + split.old_var, guarded_var, substitution_type); result.emplace_back(guarded_var_name, guarded, ApplySplitResult::LetStmt); - - Expr guard_cond = likely(old_var <= old_max); - if (split.align.defined()) { - guard_cond = likely(old_var >= old_min && old_var <= old_max); - } - result.emplace_back(guard_cond, predicate_type); + result.emplace_back(likely(old_var <= old_max), predicate_type); } else if (tail == TailStrategy::ShiftInwards) { // Adjust the base downwards to not compute off the // end of the realization. + // We'll only mark the base as likely (triggering a loop + // partition) if we're at or inside the innermost + // non-trivial loop. base = likely_if_innermost(base); - if (split.align.defined()) { - base = Max::make(base, old_min - split.align); - base = Min::make(base, old_max + (1 - split.factor) - split.align); - } else { - base = Min::make(base, old_max + (1 - split.factor)); - } + base = Min::make(base, old_max + (1 - split.factor)); } else if (tail == TailStrategy::ShiftInwardsAndBlend) { - // Unclamped base, saved before the Min/Max below adjust it. Used - // to figure out how much (if at all) the boundary tile got - // shifted, so we know which elements of it are redundant with a - // neighboring tile and must be masked out rather than - // recomputed (to avoid double-counting in a reduction). Expr old_base = base; base = likely(base); - Expr mask; - if (split.align.defined()) { - // Because base is anchored to align instead of old_min, the - // boundary tile can now be shifted at either end (whereas - // without align only the max end is reachable, since base - // is structurally >= old_min already). Elements shifted in - // from the low end overlap the tile above (mask out the - // last shift_low of them); elements shifted in from the - // high end overlap the tile below (mask out the first - // shift_high of them). - Expr low_bound = old_min - split.align; - Expr high_bound = old_max + (1 - split.factor) - split.align; - Expr shift_low = low_bound - old_base; - Expr shift_high = old_base - high_bound; - base = Max::make(base, low_bound); - base = Min::make(base, high_bound); - Expr mask_low = inner < split.factor - shift_low; - Expr mask_high = inner >= shift_high; - mask = select(old_base < low_bound, mask_low, - select(old_base > high_bound, mask_high, likely(const_true()))); - } else { - // Without align, base is structurally >= old_min (outer - // starts at 0), so only the max end can ever be shifted. - base = Min::make(base, old_max + (1 - split.factor)); - Expr unwanted_elems = (-old_extent) % split.factor; - mask = inner >= unwanted_elems; - mask = select(base == old_base, likely(const_true()), mask); - } + base = Min::make(base, old_max + (1 - split.factor)); + // Make a mask which will be a loop invariant if inner gets + // vectorized, and apply it if we're in the tail. + Expr unwanted_elems = (-old_extent) % split.factor; + Expr mask = inner >= unwanted_elems; + mask = select(base == old_base, likely(const_true()), mask); result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { - Expr mask; - if (split.align.defined()) { - // Unlike ShiftInwardsAndBlend, the max end is intentionally - // left unclamped here (RoundUp relies on padding, not on - // shifting, to handle overrun at the max end) -- but the min - // end still needs clamping: align can make the min-end tile - // start before old_min, and unlike ShiftInwards/blend at the - // max end, there's no padding below old_min to absorb an - // underrun into, so it has to be prevented outright. - // - // The mask below compares old_base (the unclamped base) - // against low_bound/high_bound directly, rather than - // comparing outer against outer_min/outer_max: the latter - // needs loop partitioning to split the loop into three - // pieces (prologue/steady-state/epilogue) to stay correct, - // and partition_loops doesn't reliably do that here when - // both boundaries are data-dependent, silently dropping the - // last tile. Comparing old_base against the bounds directly - // is correct regardless of how (or whether) the loop gets - // partitioned, matching the approach already proven correct - // above for ShiftInwardsAndBlend. - Expr old_base = base; - Expr low_bound = old_min - split.align; - Expr high_bound = old_max + (1 - split.factor) - split.align; - Expr shift_low = low_bound - old_base; - Expr shift_high = old_base - high_bound; - base = Max::make(likely(base), low_bound); - // The min end is clamped (shifted forward), so its overlap - // is with the tile *above* -- same geometry as - // ShiftInwardsAndBlend, mask out the trailing shift_low - // elements. The max end is left unclamped, so shift_high - // counts a genuine overrun past old_max with no - // neighboring tile to defer to -- mask out the trailing - // shift_high elements too (the opposite convention from - // ShiftInwardsAndBlend's clamped max end, which instead - // masks out the *leading* elements of a shifted-back tile). - Expr mask_low = inner < split.factor - shift_low; - Expr mask_high = inner < split.factor - shift_high; - mask = select(old_base < low_bound, mask_low, - select(old_base > high_bound, mask_high, likely(const_true()))); - } else { - Expr unwanted_elems = (-old_extent) % split.factor; - Expr fresh_high = inner < split.factor - unwanted_elems; - mask = select(outer < outer_max, likely(const_true()), fresh_high); - } + Expr unwanted_elems = (-old_extent) % split.factor; + Expr mask = inner < split.factor - unwanted_elems; + mask = select(outer < outer_max, likely(const_true()), mask); result.emplace_back(mask, ApplySplitResult::BlendProvides); } else { internal_assert(tail == TailStrategy::RoundUp); } - // Add align back in last, after all tail-strategy clamping/masking is - // done in terms of the unaligned base: this keeps align as a bare - // top-level addend in the final expressions (so e.g. it can still - // cancel algebraically against a matching subtraction elsewhere) - // rather than being smeared into a Max/Min-clamped expression, while - // letting the inner loop variable itself range over the simple, - // often-constant [0, factor) instead of [align, align + factor). - if (split.align.defined()) { - base = base + split.align; - } - // Define the original variable as the base value computed above plus the inner loop variable. result.emplace_back(old_var_name, base_var + inner, ApplySplitResult::LetStmt); result.emplace_back(base_name, base, ApplySplitResult::LetStmt); @@ -286,9 +176,8 @@ vector> compute_loop_bounds_after_split(const Split &spl let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 1); if (split.align.defined()) { - Expr align = split.align; - Expr outer_min = (old_var_min - align) / split.factor; - Expr outer_max = (old_var_max - align) / split.factor; + Expr outer_min = (old_var_min - split.align) / split.factor; + Expr outer_max = (old_var_max - split.align) / split.factor; let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min); let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); } else { diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index f7d1049ea4f0..0c41b498292d 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -53,6 +53,13 @@ class PartiallyCancelDifferences : public IRMutator { rewrite(min(x, y) - max(x, z), min(min(x, y) - max(x, z), 0)) || rewrite(max(x, y) - min(x, z), max(max(x, y) - min(x, z), 0)) || + // min(x + c0, y) - max(x, z) + // = min(x, y - c0) - max(x, z) + c0 + // = min(min(x, y - c0) - max(x, z), 0) + c0 + // = min(min(x, y - c0) - max(x, z) + c0, c0) + // = min(min(x + c0, y) - max(x, z), c0) + rewrite(min(x + c0, y) - max(x, z), min(min(x, y) - max(x, z), c0)) || + rewrite(min(x + c0, y) - select(z, min(x, y) + c1, x), select(z, (max(min(y - x, c0), 0) - c1), min(y - x, c0)), c0 > 0) || rewrite(min(y, x + c0) - select(z, min(y, x) + c1, x), select(z, (max(min(y - x, c0), 0) - c1), min(y - x, c0)), c0 > 0) || diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index 56a6a201b98e..4a0e5e5cd19c 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -198,6 +198,8 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite((x / w) * w + (z + x % w), select(w == 0, 0, x) + z) || rewrite(x / 2 + x % 2, (x + 1) / 2) || + rewrite((0 - (x % 2)) / 2 * 2 + (x % 2), 0 - (x % 2)) || + rewrite(x + ((c0 - x) / c1) * c1, c0 - ((c0 - x) % c1), c1 > 0) || rewrite(x + ((c0 - x) / c1 + y) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || rewrite(x + (y + (c0 - x) / c1) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..13bef635b8bd 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -95,6 +95,7 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite(max((x * c0), c1) / c2, max(x * fold(c0 / c2), fold(c1 / c2)), c0 % c2 == 0 && c2 > 0) || rewrite((x * c0 + y) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || + rewrite((x * c0 - y) / c1, x * fold(c0 / c1) - y / c0, c0 % c1 == 0 && c1 > 0) || rewrite((x * c0 - y) / c0, x + (0 - y) / c0) || rewrite((x * c1 - y) / c0, (0 - y) / c0 - x, c0 + c1 == 0) || rewrite((y + x * c0) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || @@ -131,6 +132,8 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite((w + (z + (x * c0 + y))) / c1, (y + z + w) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || rewrite((w + (z + (y + x * c0))) / c1, (y + z + w) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || + rewrite(((0 - x) - x) / 2, -x) || + /** In (x + c0) / c1, when can we pull the constant addition out of the numerator? An obvious answer is the constant is a multiple of the denominator, but diff --git a/src/Simplify_Mod.cpp b/src/Simplify_Mod.cpp index 04b082483119..649073494163 100644 --- a/src/Simplify_Mod.cpp +++ b/src/Simplify_Mod.cpp @@ -64,6 +64,10 @@ Expr Simplify::visit(const Mod *op, ExprInfo *info) { rewrite((((x * c0) - y) - z) % c1, (-y - z) % c1, c0 % c1 == 0) || rewrite((((x * c0) - y) + z) % c1, (z - y) % c1, c0 % c1 == 0) || + rewrite((x + y + y) % 2, x % 2) || + rewrite(((y + x) + y) % 2, x % 2) || + rewrite((((z + y) + x) + y) % 2, (z + x) % 2) || + rewrite(ramp(x, c0, c2) % broadcast(c1, c2), broadcast(x, c2) % broadcast(c1, c2), (c0 % c1 == 0)) || rewrite(ramp(x, c0, lanes) % broadcast(c1, lanes), ramp(x % c1, c0, lanes), // First and last lanes are the same when... From 0bd306a0dfd8abcf72b755fa2af9481ba9cb5acf Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 3 Sep 2026 22:31:14 +0200 Subject: [PATCH 27/29] Fix simplifier rule I f*kd up earlier. --- src/Simplify_Div.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 13bef635b8bd..c042abd1f036 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -95,7 +95,7 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite(max((x * c0), c1) / c2, max(x * fold(c0 / c2), fold(c1 / c2)), c0 % c2 == 0 && c2 > 0) || rewrite((x * c0 + y) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || - rewrite((x * c0 - y) / c1, x * fold(c0 / c1) - y / c0, c0 % c1 == 0 && c1 > 0) || + rewrite((x * c0 - y) / c1, (0 - y) / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || rewrite((x * c0 - y) / c0, x + (0 - y) / c0) || rewrite((x * c1 - y) / c0, (0 - y) / c0 - x, c0 + c1 == 0) || rewrite((y + x * c0) / c1, y / c1 + x * fold(c0 / c1), c0 % c1 == 0 && c1 > 0) || From 4e4be66df08baa0b4a98dc1fde0ca4bd67239cd3 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 3 Sep 2026 22:47:58 +0200 Subject: [PATCH 28/29] Restore the aligned split logic in apply_split This is the align handling from d54fbb391, which was reverted in the working tree while trying a substitution-based rewrite of the split. That rewrite is abandoned; without this, aligned splits do not lower correctly. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- src/ApplySplit.cpp | 159 ++++++++++++++++++++++++++++++++++++++------- 1 file changed, 135 insertions(+), 24 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index ab2542f77660..14c77ae772d4 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -23,10 +23,17 @@ vector apply_split(const Split &split, const string &prefix, Expr old_max = Variable::make(Int(32), prefix + split.old_var + ".loop_max"); Expr old_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); Expr old_extent = (old_max - old_min) + 1; + Expr outer_min = Variable::make(Int(32), prefix + split.outer + ".loop_min"); dim_extent_alignment[split.inner] = split.factor; - Expr base = outer * split.factor + old_min; + Expr base; + if (split.align.defined()) { + base = outer * split.factor; + } else { + base = outer * split.factor + old_min; + } + string base_name = prefix + split.inner + ".base"; Expr base_var = Variable::make(Int(32), base_name); string old_var_name = prefix + split.old_var; @@ -38,8 +45,17 @@ vector apply_split(const Split &split, const string &prefix, internal_assert(tail != TailStrategy::Auto) << "An explicit tail strategy should exist at this point\n"; + // When align is defined, tiles are anchored to align instead of to + // old_min, so knowing that the factor divides the extent is not + // enough to prove no boundary guard is needed: we additionally need + // the tiling anchored at align to line up with the tiling anchored + // at old_min, i.e. old_min and align must be congruent mod factor. + bool alignment_matches_old_min = !split.align.defined() || + is_const_zero(simplify((old_min - split.align) % split.factor)); + if ((iter != dim_extent_alignment.end()) && - is_const_zero(simplify(iter->second % split.factor))) { + is_const_zero(simplify(iter->second % split.factor)) && + alignment_matches_old_min) { // We have proved that the split factor divides the // old extent. No need to adjust the base or add an if // statement. @@ -62,10 +78,19 @@ vector apply_split(const Split &split, const string &prefix, // condition. We'll directly tell it that the loop // variable is bounded above by the original loop max by // replacing the variable with a promise-clamped version - // of it. We don't also use the original loop min because - // it needlessly complicates the expressions and doesn't - // actually communicate anything new. - Expr guarded = promise_clamped(old_var, old_var, old_max); + // of it. + Expr guarded; + if (split.align.defined()) { + // Because the un-rebased base block can start before old_min, + // we must clamp both the minimum and maximum boundaries. + guarded = promise_clamped(old_var, old_min, old_max); + } else { + // We don't also use the original loop min because + // it needlessly complicates the expressions and doesn't + // actually communicate anything new. + guarded = promise_clamped(old_var, old_var, old_max); + } + string guarded_var_name = prefix + split.old_var + ".guarded"; Expr guarded_var = Variable::make(Int(32), guarded_var_name); @@ -76,8 +101,6 @@ vector apply_split(const Split &split, const string &prefix, predicate_type = ApplySplitResult::Predicate; break; case TailStrategy::Predicate: - // This is identical to GuardWithIf, but maybe it makes - // sense to keep it anyways? substitution_type = ApplySplitResult::Substitution; predicate_type = ApplySplitResult::Predicate; break; @@ -97,36 +120,123 @@ vector apply_split(const Split &split, const string &prefix, // for the guarded version. result.emplace_back(prefix + split.old_var, guarded_var, substitution_type); result.emplace_back(guarded_var_name, guarded, ApplySplitResult::LetStmt); - result.emplace_back(likely(old_var <= old_max), predicate_type); + + Expr guard_cond = likely(old_var <= old_max); + if (split.align.defined()) { + guard_cond = likely(old_var >= old_min && old_var <= old_max); + } + result.emplace_back(guard_cond, predicate_type); } else if (tail == TailStrategy::ShiftInwards) { // Adjust the base downwards to not compute off the // end of the realization. - // We'll only mark the base as likely (triggering a loop - // partition) if we're at or inside the innermost - // non-trivial loop. base = likely_if_innermost(base); - base = Min::make(base, old_max + (1 - split.factor)); + if (split.align.defined()) { + base = Max::make(base, old_min - split.align); + base = Min::make(base, old_max + (1 - split.factor) - split.align); + } else { + base = Min::make(base, old_max + (1 - split.factor)); + } } else if (tail == TailStrategy::ShiftInwardsAndBlend) { + // Unclamped base, saved before the Min/Max below adjust it. Used + // to figure out how much (if at all) the boundary tile got + // shifted, so we know which elements of it are redundant with a + // neighboring tile and must be masked out rather than + // recomputed (to avoid double-counting in a reduction). Expr old_base = base; base = likely(base); - base = Min::make(base, old_max + (1 - split.factor)); - // Make a mask which will be a loop invariant if inner gets - // vectorized, and apply it if we're in the tail. - Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner >= unwanted_elems; - mask = select(base == old_base, likely(const_true()), mask); + Expr mask; + if (split.align.defined()) { + // Because base is anchored to align instead of old_min, the + // boundary tile can now be shifted at either end (whereas + // without align only the max end is reachable, since base + // is structurally >= old_min already). Elements shifted in + // from the low end overlap the tile above (mask out the + // last shift_low of them); elements shifted in from the + // high end overlap the tile below (mask out the first + // shift_high of them). + Expr low_bound = old_min - split.align; + Expr high_bound = old_max + (1 - split.factor) - split.align; + Expr shift_low = low_bound - old_base; + Expr shift_high = old_base - high_bound; + base = Max::make(base, low_bound); + base = Min::make(base, high_bound); + Expr mask_low = inner < split.factor - shift_low; + Expr mask_high = inner >= shift_high; + mask = select(old_base < low_bound, mask_low, + select(old_base > high_bound, mask_high, likely(const_true()))); + } else { + // Without align, base is structurally >= old_min (outer + // starts at 0), so only the max end can ever be shifted. + base = Min::make(base, old_max + (1 - split.factor)); + Expr unwanted_elems = (-old_extent) % split.factor; + mask = inner >= unwanted_elems; + mask = select(base == old_base, likely(const_true()), mask); + } result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { - Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner < split.factor - unwanted_elems; - mask = select(outer < outer_max, likely(const_true()), mask); + Expr mask; + if (split.align.defined()) { + // Unlike ShiftInwardsAndBlend, the max end is intentionally + // left unclamped here (RoundUp relies on padding, not on + // shifting, to handle overrun at the max end) -- but the min + // end still needs clamping: align can make the min-end tile + // start before old_min, and unlike ShiftInwards/blend at the + // max end, there's no padding below old_min to absorb an + // underrun into, so it has to be prevented outright. + // + // The mask below compares old_base (the unclamped base) + // against low_bound/high_bound directly, rather than + // comparing outer against outer_min/outer_max: the latter + // needs loop partitioning to split the loop into three + // pieces (prologue/steady-state/epilogue) to stay correct, + // and partition_loops doesn't reliably do that here when + // both boundaries are data-dependent, silently dropping the + // last tile. Comparing old_base against the bounds directly + // is correct regardless of how (or whether) the loop gets + // partitioned, matching the approach already proven correct + // above for ShiftInwardsAndBlend. + Expr old_base = base; + Expr low_bound = old_min - split.align; + Expr high_bound = old_max + (1 - split.factor) - split.align; + Expr shift_low = low_bound - old_base; + Expr shift_high = old_base - high_bound; + base = Max::make(likely(base), low_bound); + // The min end is clamped (shifted forward), so its overlap + // is with the tile *above* -- same geometry as + // ShiftInwardsAndBlend, mask out the trailing shift_low + // elements. The max end is left unclamped, so shift_high + // counts a genuine overrun past old_max with no + // neighboring tile to defer to -- mask out the trailing + // shift_high elements too (the opposite convention from + // ShiftInwardsAndBlend's clamped max end, which instead + // masks out the *leading* elements of a shifted-back tile). + Expr mask_low = inner < split.factor - shift_low; + Expr mask_high = inner < split.factor - shift_high; + mask = select(old_base < low_bound, mask_low, + select(old_base > high_bound, mask_high, likely(const_true()))); + } else { + Expr unwanted_elems = (-old_extent) % split.factor; + Expr fresh_high = inner < split.factor - unwanted_elems; + mask = select(outer < outer_max, likely(const_true()), fresh_high); + } result.emplace_back(mask, ApplySplitResult::BlendProvides); } else { internal_assert(tail == TailStrategy::RoundUp); } + // Add align back in last, after all tail-strategy clamping/masking is + // done in terms of the unaligned base: this keeps align as a bare + // top-level addend in the final expressions (so e.g. it can still + // cancel algebraically against a matching subtraction elsewhere) + // rather than being smeared into a Max/Min-clamped expression, while + // letting the inner loop variable itself range over the simple, + // often-constant [0, factor) instead of [align, align + factor). + if (split.align.defined()) { + base = base + split.align; + } + // Define the original variable as the base value computed above plus the inner loop variable. result.emplace_back(old_var_name, base_var + inner, ApplySplitResult::LetStmt); result.emplace_back(base_name, base, ApplySplitResult::LetStmt); @@ -176,8 +286,9 @@ vector> compute_loop_bounds_after_split(const Split &spl let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0); let_stmts.emplace_back(prefix + split.inner + ".loop_max", split.factor - 1); if (split.align.defined()) { - Expr outer_min = (old_var_min - split.align) / split.factor; - Expr outer_max = (old_var_max - split.align) / split.factor; + Expr align = split.align; + Expr outer_min = (old_var_min - align) / split.factor; + Expr outer_max = (old_var_max - align) / split.factor; let_stmts.emplace_back(prefix + split.outer + ".loop_min", outer_min); let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_max); } else { From 49501ce4b5e49e0d52d99d33e70f9cc0cb47a80f Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 3 Sep 2026 23:04:39 +0200 Subject: [PATCH 29/29] Collect a repeated term across a nested sum in the Add simplifier x + x already becomes x*2, but only when the two terms are adjacent. In a nested sum they never meet, so x + y + y stays as written and both the rewrite rules and modulus_remainder lose the fact that it is even in y. Found while investigating why an aligned split's alignment fails to cancel against a matching subtraction at the use site: the sum there is of the form (even + off) + off, whose evenness is exactly what the mux index needs. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- src/Simplify_Add.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index 4a0e5e5cd19c..8ab0fb2c61c0 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -80,6 +80,16 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite((c0 - x) + y, (y - x) + c0) || rewrite(max(x, y * c0 + z) + (u - y) * c0, max(x - y * c0, z) + u * c0) || + // Collect a repeated term across a nested sum, so that facts about + // it become visible to the rules and to modulus_remainder. For + // example, x + y + y is even in y, but only once written as x + y*2. + rewrite((x + y) + y, x + y * 2) || + rewrite((y + x) + y, x + y * 2) || + rewrite((x + (y + z)) + z, z * 2 + (x + y)) || + rewrite((x + (z + y)) + z, z * 2 + (x + y)) || + rewrite(((y + z) + x) + z, z * 2 + (x + y)) || + rewrite(((z + y) + x) + z, z * 2 + (x + y)) || + rewrite((x - y) + y, x) || rewrite(x + (y - x), y) ||