From 79f14c6262c082d901aef0ea09159f0460f645c6 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 3 Sep 2026 22:48:11 +0200 Subject: [PATCH 1/6] Add a minimal repro for aligned-split phases not cancelling through a LetStmt Distilled from NeonRAW's BilinearDemosaicAlignedGenerator, which selects an interpolation per Bayer phase with a mux over (x - bayer_offset) % 2 and is scheduled with splits aligned to that same offset so the phase is constant per tile. Twelve muxes survive lowering there; this reduces it to two. The producer is computed inside the consumer's aligned tile and then split again to vectorize. That second split binds the reconstructed loop variable to a LetStmt whose value carries the "+ off" from the aligned min. The simplifier can see through a Let, but not through a LetStmt, so at the use site it only ever sees "- off" applied to an opaque name, and (x - off) % 2 never reduces to a function of the inner loop variable alone. Confirmed by simplifying the same expression both ways: behind a LetStmt neither phase folds, while inlining the value folds both to constants -- and it still folds if the value is inlined only after the mod-2 sign flip has already rewritten "- off" into "+ off", so the LetStmt is the sole cause. The test fails at present; it documents the missed simplification. The generated code is correct, only the muxes are not resolved. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_mux_phase.cpp | 96 ++++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 test/correctness/split_aligned_mux_phase.cpp diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 93048b3bed21..d2c5a7d8cba0 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_2d_3x3.cpp + split_aligned_mux_phase.cpp split_aligned_nested.cpp split_aligned_partition.cpp split_aligned_reduction.cpp diff --git a/test/correctness/split_aligned_mux_phase.cpp b/test/correctness/split_aligned_mux_phase.cpp new file mode 100644 index 000000000000..7c15ad454ca2 --- /dev/null +++ b/test/correctness/split_aligned_mux_phase.cpp @@ -0,0 +1,96 @@ +#include "Halide.h" +#include + +// A minimal distillation of a bilinear demosaicer scheduled with aligned +// splits. A demosaicer selects which interpolation to apply from the position +// within the Bayer pattern -- a mux over (x - bayer_offset) % 2 -- and is +// scheduled with a split aligned to that same offset precisely so that the +// phase is constant across every tile, letting the mux resolve at compile +// time. +// +// The alignment cancels only if the simplifier can see the "+ off" that the +// aligned split puts into the reconstructed loop variable at the same time as +// the "- off" at the use site. When the producer is computed inside the +// consumer's tile and then split again itself (here to vectorize), the second +// split binds the reconstructed variable to a LetStmt, and the simplifier +// cannot see through a LetStmt. The "+ off" is then hidden behind an opaque +// name, (x - off) % 2 never reduces to a function of the inner loop variable +// alone, and the mux survives to the end of lowering. + +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) { + const int width = 1024; + const int tile = 128; + const int vec = 8; + + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Param off{"off"}; + off.set_range(0, 1); + + // Two "phases", selected by the position relative to the offset. + Func prod{"prod"}; + prod(x) = mux((x - off) % 2, {x + 1, x * 2}); + + Func out{"out"}; + out(x) = prod(x); + + out.output_buffer().dim(0).set_min(0); + + // The consumer is tiled with a split aligned to the same offset, so every + // tile starts at a position congruent to off. + out.split(x, xo, xi, tile, off, TailStrategy::GuardWithIf) + .vectorize(xi) + .never_partition_all(); + + // The producer is computed per tile. align_bounds keeps its min congruent + // to off as well, so its phase really is constant, and then it is split + // again to vectorize. + prod.compute_at(out, xo) + .never_partition_all() + .align_bounds(x, 2, off) + .vectorize(x, vec, TailStrategy::RoundUp); + + Pipeline p(out); + + // Every phase is known at compile time, so no mux should survive lowering. + Module m = p.compile_to_module({off}, "split_aligned_mux_phase"); + MuxCounter checker; + m.functions().front().body.accept(&checker); + if (checker.mux_count != 0) { + printf("Expected 0 muxes, got %d\n", checker.mux_count); + return 1; + } + + // And it still has to compute the right thing. + for (int o = 0; o <= 1; o++) { + off.set(o); + Buffer result = p.realize({width}); + for (int i = 0; i < width; i++) { + int correct = (((i - o) % 2) + 2) % 2 == 0 ? i + 1 : i * 2; + if (result(i) != correct) { + printf("off = %d: result(%d) = %d instead of %d\n", + o, i, result(i), correct); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} From 69063baa8781cbfcf1171f0823233f1ceeeb39e7 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 4 Sep 2026 00:11:04 +0200 Subject: [PATCH 2/6] Deinterleave ramps by strides up to seven, not four A vectorized loop over a repeating pattern only resolves its phases at compile time if it gets deinterleaved by the pattern's period, so that each resulting slice has a single phase. The stride search stopped at four, which covered a Bayer sensor's period of two but not an X-Trans sensor's six. It can't simply go to eight: deinterleaving a period of eight into eight slices displaces permutes that a target may do in one instruction, and three of simd_op_check_hvx's vdelta checks regress that way. Seven is the largest value that leaves those alone. A period of eight still resolves if the loop is unrolled rather than vectorized. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- src/Deinterleave.cpp | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/Deinterleave.cpp b/src/Deinterleave.cpp index 0fe237adb5ad..d1fcc68c9a6a 100644 --- a/src/Deinterleave.cpp +++ b/src/Deinterleave.cpp @@ -520,6 +520,17 @@ class Interleaver : public IRMutator { return visit_let(op); } + // The largest stride we'll deinterleave a ramp by. This is what decides + // whether a vectorized loop over a repeating pattern -- a sensor's colour + // filter array, say -- gets split into one slice per phase, so that each + // slice has a phase known at compile time. Bayer needs two and X-Trans + // needs six. + // + // It stops below eight because a permute with a period of eight is often + // better done as a single shuffle than as eight slices: on Hexagon, for + // instance, deinterleaving (x/8)*9 + x%8 displaces a single vdelta. + static constexpr int max_deinterleave_stride = 7; + Expr visit(const Ramp *op) override { if (op->stride.type().is_vector() && is_const_one(op->stride) && @@ -536,7 +547,7 @@ class Interleaver : public IRMutator { Expr visit(const Mod *op) override { const Ramp *r = op->a.as(); - for (int i = 2; i <= 4; ++i) { + for (int i = 2; i <= max_deinterleave_stride; ++i) { if (r && is_const(op->b, i) && (r->type.lanes() % i) == 0) { @@ -550,7 +561,7 @@ class Interleaver : public IRMutator { Expr visit(const Div *op) override { const Ramp *r = op->a.as(); - for (int i = 2; i <= 4; ++i) { + for (int i = 2; i <= max_deinterleave_stride; ++i) { if (r && is_const(op->b, i) && (r->type.lanes() % i) == 0 && From 33cd9cd64e8cbe4c46f6aab5e577c963e976839a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 4 Sep 2026 00:11:06 +0200 Subject: [PATCH 3/6] Cover more pattern periods in the aligned-split mux phase test Generalizes the test over the period of the pattern the mux selects on, and runs it for two (Bayer), three, and six (X-Trans) rather than two alone. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01JbXuRsMMLQDkqE3mwtEKMs --- test/correctness/split_aligned_mux_phase.cpp | 70 ++++++++++++-------- 1 file changed, 44 insertions(+), 26 deletions(-) diff --git a/test/correctness/split_aligned_mux_phase.cpp b/test/correctness/split_aligned_mux_phase.cpp index 7c15ad454ca2..5f3932325dc9 100644 --- a/test/correctness/split_aligned_mux_phase.cpp +++ b/test/correctness/split_aligned_mux_phase.cpp @@ -1,21 +1,19 @@ #include "Halide.h" #include -// A minimal distillation of a bilinear demosaicer scheduled with aligned -// splits. A demosaicer selects which interpolation to apply from the position -// within the Bayer pattern -- a mux over (x - bayer_offset) % 2 -- and is -// scheduled with a split aligned to that same offset precisely so that the -// phase is constant across every tile, letting the mux resolve at compile -// time. +// A minimal distillation of a demosaicer scheduled with aligned splits. A +// demosaicer selects which interpolation to apply from the position within the +// sensor's colour filter pattern -- a mux over (x - offset) % period, where the +// period is 2 for a Bayer sensor and 6 for an X-Trans one -- and is scheduled +// with a split aligned to that same offset precisely so that the phase is +// constant across every tile, letting the mux resolve at compile time. // -// The alignment cancels only if the simplifier can see the "+ off" that the -// aligned split puts into the reconstructed loop variable at the same time as -// the "- off" at the use site. When the producer is computed inside the -// consumer's tile and then split again itself (here to vectorize), the second -// split binds the reconstructed variable to a LetStmt, and the simplifier -// cannot see through a LetStmt. The "+ off" is then hidden behind an opaque -// name, (x - off) % 2 never reduces to a function of the inner loop variable -// alone, and the mux survives to the end of lowering. +// For that to happen the compiler has to see two things. The alignment the +// split adds into the reconstructed loop variable has to cancel against the +// subtraction at the use site, even though the split binds that variable to a +// LetStmt and the offset is a runtime value (see reduce_expr_modulo_symbolic). +// And the vectorized loop has to be deinterleaved by the period, so that each +// of the resulting slices has a phase of its own (see Deinterleave.cpp). using namespace Halide; using namespace Halide::Internal; @@ -34,18 +32,26 @@ class MuxCounter : public IRVisitor { int mux_count{0}; }; -int main(int argc, char **argv) { +int test(int period) { const int width = 1024; - const int tile = 128; - const int vec = 8; + // The tile has to be a whole number of periods for the phase to be the + // same in every tile, and the vector a whole number of periods for each + // deinterleaved slice to have a single phase. + const int tile = period * 64; + const int vec = period * 2; Var x{"x"}, xo{"xo"}, xi{"xi"}; Param off{"off"}; - off.set_range(0, 1); + off.set_range(0, period - 1); + + // One phase per position in the pattern. + std::vector phases; + for (int i = 0; i < period; i++) { + phases.push_back(x * (i + 1) + i); + } - // Two "phases", selected by the position relative to the offset. Func prod{"prod"}; - prod(x) = mux((x - off) % 2, {x + 1, x * 2}); + prod(x) = mux((x - off) % period, phases); Func out{"out"}; out(x) = prod(x); @@ -63,7 +69,7 @@ int main(int argc, char **argv) { // again to vectorize. prod.compute_at(out, xo) .never_partition_all() - .align_bounds(x, 2, off) + .align_bounds(x, period, off) .vectorize(x, vec, TailStrategy::RoundUp); Pipeline p(out); @@ -73,24 +79,36 @@ int main(int argc, char **argv) { MuxCounter checker; m.functions().front().body.accept(&checker); if (checker.mux_count != 0) { - printf("Expected 0 muxes, got %d\n", checker.mux_count); + printf("Period %d: expected 0 muxes, got %d\n", period, checker.mux_count); return 1; } // And it still has to compute the right thing. - for (int o = 0; o <= 1; o++) { + for (int o = 0; o < period; o++) { off.set(o); Buffer result = p.realize({width}); for (int i = 0; i < width; i++) { - int correct = (((i - o) % 2) + 2) % 2 == 0 ? i + 1 : i * 2; + int phase = (((i - o) % period) + period) % period; + int correct = i * (phase + 1) + phase; if (result(i) != correct) { - printf("off = %d: result(%d) = %d instead of %d\n", - o, i, result(i), correct); + printf("Period %d, off = %d: result(%d) = %d instead of %d\n", + period, o, i, result(i), correct); return 1; } } } + return 0; +} + +int main(int argc, char **argv) { + // Bayer, a three-phase pattern, and X-Trans. + for (int period : {2, 3, 6}) { + if (test(period) != 0) { + return 1; + } + } + printf("Success!\n"); return 0; } From bf362243392153ad4f39c5482514d31259645b00 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 4 Sep 2026 20:02:53 +0200 Subject: [PATCH 4/6] Bind a variable's value modulo a constant where the variable is defined A value is often needed only modulo some constant: x % 5, (x + 3) % 5 and (x + 2*y) % 3 all need no more of x than x % 5. Reducing the value where it is defined frequently collapses it. A loop variable reconstructed by an aligned split looks like xo * 16 + offset, which modulo two is just offset. The simplifier can't discover that at the use site. A LetStmt hides the value behind a name, and it won't substitute a whole expression back in. It will substitute a variable, so this pass binds the reduced value to one: (x - offset) % 2 becomes (x.mod.2 - offset) % 2, and with x.mod.2 bound to offset that folds to zero. This is what lets a demosaicer scheduled with aligned splits resolve its Bayer phases at compile time; twelve muxes survived lowering in the distilled test before. Only reductions that collapse to a constant or a single variable are bound, so the pass can never grow the IR, and the re-simplification afterwards is skipped when it bound nothing. On a pipeline that benefits, compiling gets faster rather than slower, because everything downstream sees smaller IR: 0.83s against 0.96s for twenty compilations of a demosaicer. On one that doesn't, it is free: 0.354s against 0.349s for a plain blur. Reducing a value looks through the lets it is built from, to a bounded depth, so that terms which cancel get the chance to meet, and it takes the conditions of dominating asserts as assumptions -- a require() on an offset is often the only thing saying it is small enough to be its own remainder. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016WxFGPpnekN4bH6sMCsQBQ --- src/CMakeLists.txt | 2 + src/InjectModuloVars.cpp | 384 ++++++++++++++++++++++++ src/InjectModuloVars.h | 34 +++ src/Lower.cpp | 11 + test/correctness/CMakeLists.txt | 1 + test/correctness/inject_modulo_vars.cpp | 111 +++++++ 6 files changed, 543 insertions(+) create mode 100644 src/InjectModuloVars.cpp create mode 100644 src/InjectModuloVars.h create mode 100644 test/correctness/inject_modulo_vars.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 02215600ef8e..26add5d97e41 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -136,6 +136,7 @@ target_sources( ImageParam.h InferArguments.h InjectHostDevBufferCopies.h + InjectModuloVars.h Inline.h InlineReductions.h IntegerDivisionTable.h @@ -319,6 +320,7 @@ target_sources( ImageParam.cpp InferArguments.cpp InjectHostDevBufferCopies.cpp + InjectModuloVars.cpp Inline.cpp InlineReductions.cpp IntegerDivisionTable.cpp diff --git a/src/InjectModuloVars.cpp b/src/InjectModuloVars.cpp new file mode 100644 index 000000000000..fd4e23d4c976 --- /dev/null +++ b/src/InjectModuloVars.cpp @@ -0,0 +1,384 @@ +#include "InjectModuloVars.h" + +#include "IREquality.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "IRVisitor.h" +#include "Scope.h" +#include "Simplify.h" + +#include +#include +#include +#include +#include + +namespace Halide { +namespace Internal { + +namespace { + +// Reducing modulo m is only exact on integer types where overflow is +// undefined. On a type that wraps, a value and its reduction agree modulo m +// only if m divides the wrap-around, which isn't worth chasing. +bool suitable_type(Type t) { + return t.is_int() && t.bits() >= 32; +} + +// Note the check against zero: x % 0 is defined to be 0 rather than x, so it +// tells us nothing about x. +bool is_multiple_of(const Expr &e, int64_t m) { + auto c = as_const_int(e); + return c && *c != 0 && mod_imp(*c, m) == 0; +} + +// Call f on every Variable sitting in a position where only its value modulo +// m matters. Addition, subtraction and multiplication all commute with +// reduction, a ramp's base and stride each contribute linearly, and a value +// already reduced modulo a multiple of m is congruent to the original. +// Division is not on the list, and anything else is opaque. +template +void for_each_congruent_var(const Expr &e, int64_t m, Fn f) { + if (const Add *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Sub *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Mul *op = e.as()) { + for_each_congruent_var(op->a, m, f); + for_each_congruent_var(op->b, m, f); + } else if (const Ramp *op = e.as()) { + for_each_congruent_var(op->base, m, f); + for_each_congruent_var(op->stride, m, f); + } else if (const Broadcast *op = e.as()) { + for_each_congruent_var(op->value, m, f); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + for_each_congruent_var(op->a, m, f); + } + } else if (const Variable *op = e.as()) { + f(op); + } +} + +// Rebuild e with the variables in those same positions replaced. +Expr substitute_congruent_vars(const Expr &e, int64_t m, + const std::map &replacements) { + if (const Add *op = e.as()) { + return Add::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Sub *op = e.as()) { + return Sub::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Mul *op = e.as()) { + return Mul::make(substitute_congruent_vars(op->a, m, replacements), + substitute_congruent_vars(op->b, m, replacements)); + } else if (const Ramp *op = e.as()) { + return Ramp::make(substitute_congruent_vars(op->base, m, replacements), + substitute_congruent_vars(op->stride, m, replacements), + op->lanes); + } else if (const Broadcast *op = e.as()) { + return Broadcast::make(substitute_congruent_vars(op->value, m, replacements), + op->lanes); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + return Mod::make(substitute_congruent_vars(op->a, m, replacements), op->b); + } + } else if (const Variable *op = e.as()) { + auto it = replacements.find(op->name); + if (it != replacements.end() && it->second.type() == op->type) { + return it->second; + } + } + return e; +} + +template +const V *find(const std::map &m, const std::string &key) { + auto it = m.find(key); + return it == m.end() ? nullptr : &it->second; +} + +// Which let-bound variables are needed modulo which constants, and what the +// lets bind them to. +class FindModuloUses : public IRVisitor { + using IRVisitor::visit; + + void note(const Expr &e, int64_t m) { + for_each_congruent_var(e, m, [&](const Variable *v) { + if (let_values.count(v->name)) { + required[v->name].insert(m); + } + }); + } + + void visit(const Mod *op) override { + if (suitable_type(op->type.element_of())) { + if (auto m = as_const_int(op->b); m && *m > 1) { + note(op->a, *m); + } + } + IRVisitor::visit(op); + } + + void visit(const Let *op) override { + if (suitable_type(op->value.type().element_of())) { + let_values[op->name] = op->value; + } + IRVisitor::visit(op); + } + + void visit(const LetStmt *op) override { + if (suitable_type(op->value.type().element_of())) { + let_values[op->name] = op->value; + } + IRVisitor::visit(op); + } + +public: + std::map> required; + std::map let_values; + + // A value that must be known modulo m needs the things it is built from + // known modulo m too. Chase that to a fixed point. + void close_over_let_values() { + std::vector> worklist; + for (const auto &[name, moduli] : required) { + for (int64_t m : moduli) { + worklist.emplace_back(name, m); + } + } + while (!worklist.empty()) { + auto [name, m] = worklist.back(); + worklist.pop_back(); + for_each_congruent_var(let_values[name], m, [&](const Variable *v) { + if (let_values.count(v->name) && required[v->name].insert(m).second) { + worklist.emplace_back(v->name, m); + } + }); + } + } +}; + +// The bounds a scalar Parameter is constrained to. add_parameter_checks +// asserts these on entry, so the simplifier may rely on them, but by this +// point in lowering the constrained copy of the parameter is long gone. +class FindParamBounds : public IRVisitor { + using IRVisitor::visit; + + void visit(const Variable *op) override { + if (op->param.defined() && !op->param.is_buffer() && + !bounds.contains(op->name)) { + Expr lo = op->param.min_value(), hi = op->param.max_value(); + if (lo.defined() || hi.defined()) { + bounds.push(op->name, Interval(lo.defined() ? lo : Interval::neg_inf(), + hi.defined() ? hi : Interval::pos_inf())); + } + } + } + +public: + Scope bounds; +}; + +class InjectModuloVars : public IRMutator { + using IRMutator::visit; + + const std::map> &required; + const std::map &let_values; + const Scope ¶m_bounds; + + // How deep we're willing to look through lets while building a reduced + // value, and how big we'll let the result get before giving up. Both only + // bound work done to decide whether a reduction collapses; nothing this + // large is ever added to the IR. + static constexpr int max_inline_depth = 4; + static constexpr int max_nodes = 1000; + int nodes_built = 0; + + // Conditions asserted by statements that dominate the one we're in. + // A require() on a parameter lands here, which is often what makes a + // reduction collapse: nothing else tells the simplifier that an offset + // is small enough to be its own remainder. + std::vector assumptions; + + // For a variable needed modulo m, the constant or variable it is congruent + // to, keyed by the name of the variable we bound that to. Reductions that + // didn't collapse aren't in here, and aren't bound at all. + std::map> congruent_value; + std::map> congruent_var; + + static std::string mod_var_name(const std::string &name, int64_t m) { + return name + ".mod." + std::to_string(m); + } + + // Build something congruent to e modulo m, in the hope that it collapses. + // A variable in a congruence-preserving position becomes whatever we + // already know it to be congruent to, or failing that its own let value, + // so that terms which cancel get the chance to meet. + Expr reduce(const Expr &e, int64_t m, int depth) { + if (nodes_built++ > max_nodes) { + return e; + } + if (const Add *op = e.as()) { + return Add::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Sub *op = e.as()) { + return Sub::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Mul *op = e.as()) { + return Mul::make(reduce(op->a, m, depth), reduce(op->b, m, depth)); + } else if (const Ramp *op = e.as()) { + return Ramp::make(reduce(op->base, m, depth), reduce(op->stride, m, depth), op->lanes); + } else if (const Broadcast *op = e.as()) { + return Broadcast::make(reduce(op->value, m, depth), op->lanes); + } else if (const Mod *op = e.as()) { + if (is_multiple_of(op->b, m)) { + return Mod::make(reduce(op->a, m, depth), op->b); + } + } else if (const Variable *op = e.as()) { + if (const Expr *known = find(congruent_value[m], op->name)) { + return *known; + } + if (depth < max_inline_depth) { + if (const Expr *value = find(let_values, op->name)) { + if (value->type() == op->type) { + return reduce(*value, m, depth + 1); + } + } + } + } + return e; + } + + // The reductions to bind just inside a let of this name and value, in + // order. Records what they are congruent to so later reductions and use + // sites can pick them up. + std::vector> reductions_for(const std::string &name, + const Expr &value) { + std::vector> result; + auto it = required.find(name); + if (it == required.end()) { + return result; + } + for (int64_t m : it->second) { + nodes_built = 0; + Expr reduced = reduce(value, m, 0); + reduced = simplify(Mod::make(reduced, make_const(value.type(), m)), + param_bounds, Scope::empty_scope(), assumptions); + // Only worth binding if it collapsed to something the simplifier + // will substitute back in at the use site. + if (!is_const(reduced) && !reduced.as()) { + continue; + } + std::string new_name = mod_var_name(name, m); + result.emplace_back(new_name, reduced); + congruent_value[m][name] = reduced; + congruent_var[m][name] = Variable::make(value.type(), new_name); + } + return result; + } + + void forget(const std::string &name, + const std::vector> &reductions) { + for (const auto &[new_name, value] : reductions) { + // Recover the modulus from the name we just made. + for (auto &[m, names] : congruent_value) { + if (mod_var_name(name, m) == new_name) { + names.erase(name); + congruent_var[m].erase(name); + } + } + } + } + + Stmt visit(const Block *op) override { + Stmt first = mutate(op->first); + bool pushed = false; + if (const AssertStmt *a = first.as()) { + if (is_pure(a->condition)) { + assumptions.push_back(a->condition); + pushed = true; + } + } + Stmt rest = mutate(op->rest); + if (pushed) { + assumptions.pop_back(); + } + if (first.same_as(op->first) && rest.same_as(op->rest)) { + return op; + } + return Block::make(first, rest); + } + + Expr visit(const Mod *op) override { + Expr a = mutate(op->a), b = mutate(op->b); + if (suitable_type(op->type.element_of())) { + if (auto m = as_const_int(b); m && *m > 1) { + auto it = congruent_var.find(*m); + if (it != congruent_var.end()) { + a = substitute_congruent_vars(a, *m, it->second); + } + } + } + if (a.same_as(op->a) && b.same_as(op->b)) { + return op; + } + return Mod::make(a, b); + } + + Expr visit(const Let *op) override { + Expr value = mutate(op->value); + auto reductions = reductions_for(op->name, value); + Expr body = mutate(op->body); + forget(op->name, reductions); + for (auto it = reductions.rbegin(); it != reductions.rend(); it++) { + body = Let::make(it->first, it->second, body); + } + if (value.same_as(op->value) && body.same_as(op->body)) { + return op; + } + return Let::make(op->name, value, body); + } + + Stmt visit(const LetStmt *op) override { + Expr value = mutate(op->value); + auto reductions = reductions_for(op->name, value); + Stmt body = mutate(op->body); + forget(op->name, reductions); + for (auto it = reductions.rbegin(); it != reductions.rend(); it++) { + body = LetStmt::make(it->first, it->second, body); + } + if (value.same_as(op->value) && body.same_as(op->body)) { + return op; + } + return LetStmt::make(op->name, value, body); + } + +public: + InjectModuloVars(const std::map> &required, + const std::map &let_values, + const Scope ¶m_bounds) + : required(required), let_values(let_values), param_bounds(param_bounds) { + } + + using IRMutator::mutate; +}; + +} // namespace + +Stmt inject_modulo_vars(const Stmt &s) { + FindModuloUses finder; + s.accept(&finder); + if (finder.required.empty()) { + return s; + } + finder.close_over_let_values(); + + FindParamBounds params; + s.accept(¶ms); + + return InjectModuloVars(finder.required, finder.let_values, params.bounds).mutate(s); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/InjectModuloVars.h b/src/InjectModuloVars.h new file mode 100644 index 000000000000..346865d017c3 --- /dev/null +++ b/src/InjectModuloVars.h @@ -0,0 +1,34 @@ +#ifndef HALIDE_INJECT_MODULO_VARS_H +#define HALIDE_INJECT_MODULO_VARS_H + +/** \file + * Defines a pass that makes the value of a variable modulo a constant visible + * at the places where only that much of it is needed. + */ + +#include "Expr.h" + +namespace Halide { +namespace Internal { + +/** Find variables used only modulo a constant -- x % 5, (x + 3) % 5, and + * (x + 2*y) % 3 all need no more of x than x % 5 -- and bind that reduced + * value to a new variable next to the let that defines the original. Uses of + * the original in those positions are then replaced by the new variable. + * + * Reducing a value where it is defined often collapses it. A loop variable + * reconstructed by an aligned split looks like xo * 16 + offset, which modulo + * two is just offset. The simplifier can't discover that at the use site: a + * LetStmt hides the value behind a name, and it won't substitute a whole + * expression back in. It will substitute a variable, so giving the use site a + * name bound to the reduced value is enough. (x - offset) % 2 becomes + * (x.mod.2 - offset) % 2, and with x.mod.2 bound to offset that folds to zero. + * + * Only reductions that collapse to a constant or a single variable are kept, + * so this never grows the IR. */ +Stmt inject_modulo_vars(const Stmt &s); + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/Lower.cpp b/src/Lower.cpp index 21acfff8df2d..82f8b5419323 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -41,6 +41,7 @@ #include "IRPrinter.h" #include "InferArguments.h" #include "InjectHostDevBufferCopies.h" +#include "InjectModuloVars.h" #include "Inline.h" #include "LICM.h" #include "LoopCarry.h" @@ -377,6 +378,16 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after rewriting vector interleavings:", s); + debug(1) << "Injecting variables for values needed modulo a constant...\n"; + { + // Only worth re-simplifying if it bound anything. + Stmt t = inject_modulo_vars(s); + if (!t.same_as(s)) { + s = simplify(t); + log("Lowering after injecting modulo variables:", s); + } + } + debug(1) << "Partitioning loops to simplify boundary conditions...\n"; s = partition_loops(s); s = simplify(s); diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d2c5a7d8cba0..b73879a5d543 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -207,6 +207,7 @@ tests( in_place.cpp indexing_access_undef.cpp infer_arguments.cpp + inject_modulo_vars.cpp inline_reduction.cpp inlined_generator.cpp input_func_name_unique.cpp diff --git a/test/correctness/inject_modulo_vars.cpp b/test/correctness/inject_modulo_vars.cpp new file mode 100644 index 000000000000..edb7cdc88d04 --- /dev/null +++ b/test/correctness/inject_modulo_vars.cpp @@ -0,0 +1,111 @@ +#include "Halide.h" +#include + +// Tests the lowering pass that binds a variable's value modulo a constant next +// to the let that defines it. See InjectModuloVars.h. + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +// Pull the argument back out of the sink() call the statements below wrap +// around the expression under test. +class FindSinkArg : public IRVisitor { + using IRVisitor::visit; + + void visit(const Call *op) override { + IRVisitor::visit(op); + if (op->name == "sink") { + arg = op->args[0]; + } + } + +public: + Expr arg; +}; + +Expr run(const Expr &e, const std::vector &asserts, const Expr &let_value) { + Expr x = Variable::make(Int(32), "x"); + // Mention x a second time so the simplifier won't just inline the let and + // make the whole exercise moot -- that is exactly the situation this pass + // exists for. + Stmt s = Evaluate::make(Call::make(Int(32), "sink", {e, x}, Call::Extern)); + s = LetStmt::make("x", let_value, s); + for (auto it = asserts.rbegin(); it != asserts.rend(); it++) { + s = Block::make(AssertStmt::make(*it, 0), s); + } + s = simplify(inject_modulo_vars(s)); + FindSinkArg finder; + s.accept(&finder); + return finder.arg; +} + +int check(const char *name, const Expr &got, const Expr &expected) { + if (!equal(got, expected)) { + std::cerr << name << ": got " << got << ", expected " << expected << "\n"; + return 1; + } + return 0; +} + +} // namespace + +int main(int argc, char **argv) { + Expr x = Variable::make(Int(32), "x"); + Expr xo = Variable::make(Int(32), "xo"); + Expr off = Variable::make(Int(32), "off"); + std::vector off_is_small{0 <= off && off <= 1}; + + // An aligned split's loop variable, used with the alignment subtracted + // back off. x is congruent to off modulo two, so this is zero. The extra + // term keeps the simplifier from peeling the "+ off" off the let by + // itself, which is what it does when the value ends in a bare variable. + Expr yo = Variable::make(Int(32), "yo"); + Expr aligned = xo * 16 + (yo * 2 + off); + int result = 0; + result |= check("cancel", run((x - off) % 2, off_is_small, aligned), 0); + // The same thing written the other way round, and with an offset. + result |= check("cancel, added", run((x + off) % 2, off_is_small, aligned), 0); + result |= check("cancel, +1", run((x - off + 1) % 2, off_is_small, aligned), 1); + + // The variable reached through a multiply is still only needed mod 2. + result |= check("through a multiply", run((3 * x - 3 * off) % 2, off_is_small, aligned), 0); + + // Without the assert there's nothing to say off is its own remainder, so + // this must be left alone rather than rewritten into something wrong. + { + Expr got = run((x - off) % 2, {}, aligned); + if (is_const(got)) { + std::cerr << "unbounded offset: should not have folded, got " << got << "\n"; + result = 1; + } + } + + // A congruence modulo 2 says nothing about the value modulo 3 here, since + // 16 is not a multiple of 3. + { + Expr got = run((x - off) % 3, off_is_small, aligned); + if (is_const(got)) { + std::cerr << "wrong modulus: should not have folded, got " << got << "\n"; + result = 1; + } + } + + // Division is not congruence-preserving, so nothing may be substituted + // through it. + { + Expr got = run((x / 2) % 2, off_is_small, aligned); + if (is_const(got)) { + std::cerr << "through a divide: should not have folded, got " << got << "\n"; + result = 1; + } + } + + if (result) { + return 1; + } + + printf("Success!\n"); + return 0; +} From b3e0958399a7d7966dc9327fbd86562708f6a007 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 4 Sep 2026 21:32:44 +0200 Subject: [PATCH 5/6] Add InjectModuloVars to the Makefile's file lists The correctness test needs no entry: the Makefile globs test/correctness. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016WxFGPpnekN4bH6sMCsQBQ --- Makefile | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Makefile b/Makefile index c66ac327781f..d3680426ae4e 100644 --- a/Makefile +++ b/Makefile @@ -524,6 +524,7 @@ SOURCE_FILES = \ ImageParam.cpp \ InferArguments.cpp \ InjectHostDevBufferCopies.cpp \ + InjectModuloVars.cpp \ Inline.cpp \ InlineReductions.cpp \ IntegerDivisionTable.cpp \ @@ -730,6 +731,7 @@ HEADER_FILES = \ ImageParam.h \ InferArguments.h \ InjectHostDevBufferCopies.h \ + InjectModuloVars.h \ Inline.h \ InlineReductions.h \ IntegerDivisionTable.h \ From bf39e80959278665c00cc8bed0247ef888056a96 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sat, 5 Sep 2026 15:34:47 +0200 Subject: [PATCH 6/6] Fix C++20 extension not allowed. Co-authored-by: Claude --- src/InjectModuloVars.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/InjectModuloVars.cpp b/src/InjectModuloVars.cpp index fd4e23d4c976..22a1d73a6282 100644 --- a/src/InjectModuloVars.cpp +++ b/src/InjectModuloVars.cpp @@ -150,9 +150,12 @@ class FindModuloUses : public IRVisitor { } } while (!worklist.empty()) { - auto [name, m] = worklist.back(); + // Not a structured binding: the lambda below captures the modulus, + // and capturing a structured binding is C++20. + const std::pair item = worklist.back(); worklist.pop_back(); - for_each_congruent_var(let_values[name], m, [&](const Variable *v) { + const int64_t m = item.second; + for_each_congruent_var(let_values[item.first], m, [&](const Variable *v) { if (let_values.count(v->name) && required[v->name].insert(m).second) { worklist.emplace_back(v->name, m); }