From 5f684b397707f301efc42856331a3d9de3ad1fb1 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 25 Aug 2026 18:14:54 +0200 Subject: [PATCH 01/36] WIP: shared BoundsTracker for BoundConstantExtentLoops/BoundSmallAllocations/AllocationBoundsInference Introduces BoundsTracker, a struct that accumulates enclosing pure LetStmt/Let bindings and dominating facts while a mutator descends a Stmt tree, and uses them to find constant bounds far more reliably than a bare find_constant_bound() call. In addition to a Scope fast path, find_constant_bound_aggressive()/find_constant_bounds_aggressive() fall back to wrapping an expression in all pending pure lets, inlining them with substitute_in_all_lets, and re-simplifying under the dominating facts -- generalizing the trick bound_constant_extent_loops has always used to find constant loop extents. Migrates all three targeted passes onto it: - BoundConstantExtentLoops: same two-tier (exact vs guarded upper bound) unroll/vectorize logic, now expressed via find_constant_bounds_aggressive()'s interval collapse-to-a-point check instead of a separate ad hoc IntImm check. - BoundSmallAllocations: Frame/visit_let chain now binds through tracker.push_let(); find_constant_bound() call sites upgraded to find_constant_bound_aggressive() so allocation/realize extents get the same aggressive treatment. - AllocationBoundsInference: gains LetStmt/For tracking it never had before, and runs the box_touched() result through tracker.simplify_with_context() before CSE, so a Realize's per- dimension min/max can be simplified using enclosing let context that box_touched (called with an empty scope) can't see on its own. KNOWN ISSUE (not yet resolved): correctness_unroll_loop_with_implied_constant_bounds segfaults via infinite recursion inside Simplify's fact/var_info substitution machinery, triggered from BoundConstantExtentLoops's aggressive fallback when two dominating facts (a bounds-query check and a 4-way equality conjunction "three_channels") are both active for the same simplify() call. Confirmed via debug instrumentation that BoundsTracker builds the same wrapped expression and fact list the original hand-rolled implementation would have; the crash is inside the Simplify engine's own var_info replacement logic (Simplify_Exprs.cpp:270-282), not in BoundsTracker's bookkeeping. Root cause not yet isolated -- needs a minimal standalone repro against Simplify() directly (bypassing BoundsTracker/Lower.cpp entirely) to determine whether this is a latent pre-existing Simplify bug that BoundConstantExtentLoops previously never triggered, or a subtle behavioral difference between BoundsTracker's fact accumulation and the original vector-based one. All other targeted correctness tests (bounds, bound_small_allocations, unroll, vectorize, realize, extern, sliding_window, partition_loops, split, etc.) pass. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_019EGMmdqNC6mTcMSCDBFbwV --- src/AllocationBoundsInference.cpp | 24 +++++- src/BoundConstantExtentLoops.cpp | 76 +++++++++--------- src/BoundSmallAllocations.cpp | 24 +++--- src/BoundsTracker.cpp | 110 ++++++++++++++++++++++++++ src/BoundsTracker.h | 126 ++++++++++++++++++++++++++++++ src/CMakeLists.txt | 2 + 6 files changed, 303 insertions(+), 59 deletions(-) create mode 100644 src/BoundsTracker.cpp create mode 100644 src/BoundsTracker.h diff --git a/src/AllocationBoundsInference.cpp b/src/AllocationBoundsInference.cpp index 9aaaf1c7e661..187aad4f1f6f 100644 --- a/src/AllocationBoundsInference.cpp +++ b/src/AllocationBoundsInference.cpp @@ -1,5 +1,6 @@ #include "AllocationBoundsInference.h" #include "Bounds.h" +#include "BoundsTracker.h" #include "CSE.h" #include "ExternFuncArgument.h" #include "Function.h" @@ -19,10 +20,6 @@ using std::vector; namespace { -Expr cse_and_simplify(const Expr &x) { - return simplify(common_subexpression_elimination(x)); -} - // Figure out the region touched of each buffer, and deposit them as // let statements outside of each realize node, or at the top level if // they're not internal allocations. @@ -34,6 +31,25 @@ class AllocationInference : public IRMutator { const FuncValueBounds &func_bounds; set touched_by_extern; + // Tracks the enclosing pure lets and for loops, so that box_touched's + // result -- computed with no knowledge of anything outside op->body -- + // can be simplified as thoroughly as if it had been. + BoundsTracker tracker; + + Expr cse_and_simplify(const Expr &x) { + return simplify(common_subexpression_elimination(tracker.simplify_with_context(x))); + } + + Stmt visit(const LetStmt *op) override { + auto binding = tracker.push_let(op->name, op->value); + return IRMutator::visit(op); + } + + Stmt visit(const For *op) override { + auto binding = tracker.push_for(op->name, op->min, op->max); + return IRMutator::visit(op); + } + Stmt visit(const Realize *op) override { map::const_iterator iter = env.find(op->name); internal_assert(iter != env.end()); diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index ebc41007e6bf..5691ee575df6 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -1,11 +1,10 @@ #include "BoundConstantExtentLoops.h" #include "Bounds.h" -#include "CSE.h" +#include "BoundsTracker.h" #include "IRMutator.h" #include "IROperator.h" #include "Simplify.h" -#include "SimplifyCorrelatedDifferences.h" -#include "Substitute.h" +#include "Util.h" namespace Halide { namespace Internal { @@ -15,29 +14,23 @@ class BoundLoops : public IRMutator { protected: using IRMutator::visit; - std::vector> lets; + BoundsTracker tracker; Stmt visit(const LetStmt *op) override { - if (is_pure(op->value)) { - lets.emplace_back(op->name, op->value); - Stmt s = IRMutator::visit(op); - lets.pop_back(); - return s; - } else { - return IRMutator::visit(op); - } + auto binding = tracker.push_let(op->name, op->value); + return IRMutator::visit(op); } - std::vector facts; Stmt visit(const IfThenElse *op) override { - facts.push_back(op->condition); - Stmt then_case = mutate(op->then_case); - Stmt else_case; + Stmt then_case, else_case; + { + auto fact = tracker.push_fact(op->condition); + then_case = mutate(op->then_case); + } if (op->else_case.defined()) { - facts.back() = simplify(!op->condition); + auto fact = tracker.push_fact(simplify(!op->condition)); else_case = mutate(op->else_case); } - facts.pop_back(); if (then_case.same_as(op->then_case) && else_case.same_as(op->else_case)) { return op; @@ -59,33 +52,36 @@ class BoundLoops : public IRMutator { Stmt body = op->body; const IntImm *e = extent.as(); - if (e == nullptr) { - // We're about to hard fail. Get really aggressive - // with the simplifier. - extent = rewrap_used_lets(extent, lets); - extent = remove_likelies(extent); - extent = substitute_in_all_lets(extent); - extent = simplify(extent, - Scope::empty_scope(), - Scope::empty_scope(), - facts); - e = extent.as(); - } - Expr extent_upper; if (e == nullptr) { - // Still no luck. Try taking an upper bound and - // injecting an if statement around the body. - extent_upper = find_constant_bound(extent, Direction::Upper, Scope()); - if (extent_upper.defined()) { - e = extent_upper.as(); - body = - IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <= - op->max), - body); + // We're about to hard fail. Get really aggressive with the + // simplifier: inline every enclosing let and simplify under + // every dominating condition. + debug(4) << "Trying to find a constant bound for loop " << op->name << "\n" + << "Extent: " << extent << "\n"; + Interval bounds = tracker.find_constant_bounds_aggressive(extent); + debug(4) << "Bounds found: [" << bounds.min << ", " << bounds.max << "]\n"; + auto lo = bounds.has_lower_bound() ? as_const_int(bounds.min) : std::nullopt; + auto hi = bounds.has_upper_bound() ? as_const_int(bounds.max) : std::nullopt; + if (lo && hi && *lo == *hi) { + // The bound is exact. + e = bounds.max.as(); + } else if (hi) { + extent_upper = bounds.max; } } + if (e == nullptr && extent_upper.defined()) { + // Still no luck getting an exact extent. Take the upper + // bound instead and guard the body with an if statement. + debug(4) << "Found an upper bound instead: " << extent_upper << "\n"; + e = extent_upper.as(); + body = + IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <= + op->max), + body); + } + if (e == nullptr && permit_failed_unroll && op->for_type == ForType::Unrolled) { // Still no luck, but we're allowed to fail. Rewrite // to a serial loop. diff --git a/src/BoundSmallAllocations.cpp b/src/BoundSmallAllocations.cpp index c8683b08a4ce..9c987c04dd6e 100644 --- a/src/BoundSmallAllocations.cpp +++ b/src/BoundSmallAllocations.cpp @@ -1,9 +1,8 @@ #include "BoundSmallAllocations.h" -#include "Bounds.h" +#include "BoundsTracker.h" #include "CodeGen_Internal.h" #include "IRMutator.h" #include "IROperator.h" -#include "Simplify.h" namespace Halide { namespace Internal { @@ -15,17 +14,17 @@ class BoundSmallAllocations : public IRMutator { using IRMutator::visit; // Track constant bounds - Scope scope; + BoundsTracker tracker; template auto visit_let(const LetOrLetStmt *op) -> decltype(op->body) { // Visit an entire chain of lets in a single method to conserve stack space. struct Frame { const LetOrLetStmt *op; - ScopedBinding binding; - Frame(const LetOrLetStmt *op, Scope &scope) + BoundsTracker::Binding binding; + Frame(const LetOrLetStmt *op, BoundsTracker &tracker) : op(op), - binding(scope, op->name, find_constant_bounds(op->value, scope)) { + binding(tracker.push_let(op->name, op->value)) { } }; std::vector frames; @@ -33,7 +32,7 @@ class BoundSmallAllocations : public IRMutator { do { result = op->body; - frames.emplace_back(op, scope); + frames.emplace_back(op, tracker); } while ((op = result.template as())); result = mutate(result); @@ -58,12 +57,7 @@ class BoundSmallAllocations : public IRMutator { DeviceAPI device_api = DeviceAPI::None; Stmt visit(const For *op) override { - Interval min_bounds = find_constant_bounds(op->min, scope); - Interval max_bounds = find_constant_bounds(op->max, scope); - Interval b = Interval::make_union(min_bounds, max_bounds); - b.min = simplify(b.min); - b.max = simplify(b.max); - ScopedBinding bind(scope, op->name, b); + auto binding = tracker.push_for(op->name, op->min, op->max); bool new_in_thread_loop = in_thread_loop || op->for_type == ForType::GPUThread; ScopedValue old_in_thread_loop(in_thread_loop, new_in_thread_loop); @@ -86,7 +80,7 @@ class BoundSmallAllocations : public IRMutator { bool changed = false; bool found_non_constant_extent = false; for (Range &r : region) { - Expr bound = find_constant_bound(r.extent, Direction::Upper, scope); + Expr bound = tracker.find_constant_bound_aggressive(r.extent, Direction::Upper); // We can allow non-constant extents for now, as long as all // remaining dimensions are 1 (so the stride is unused, which // will be non-constant). @@ -116,7 +110,7 @@ class BoundSmallAllocations : public IRMutator { for (const Expr &e : op->extents) { total_extent *= e; } - Expr bound = find_constant_bound(total_extent, Direction::Upper, scope); + Expr bound = tracker.find_constant_bound_aggressive(total_extent, Direction::Upper); if (!bound.defined() && must_be_constant(op->memory_type)) { user_assert(op->memory_type != MemoryType::Register) diff --git a/src/BoundsTracker.cpp b/src/BoundsTracker.cpp new file mode 100644 index 000000000000..d4d155072f36 --- /dev/null +++ b/src/BoundsTracker.cpp @@ -0,0 +1,110 @@ +#include "BoundsTracker.h" + +#include "IR.h" +#include "IROperator.h" +#include "Simplify.h" +#include "Substitute.h" + +namespace Halide { +namespace Internal { + +BoundsTracker::Binding::Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let) + : tracker(tracker), scope_binding(std::move(scope_binding)), recorded_let(recorded_let) { +} + +BoundsTracker::Binding::Binding(Binding &&other) noexcept + : tracker(other.tracker), + scope_binding(std::move(other.scope_binding)), + recorded_let(other.recorded_let) { + other.recorded_let = false; +} + +BoundsTracker::Binding::~Binding() { + if (recorded_let) { + tracker->lets.pop_back(); + } +} + +BoundsTracker::Binding BoundsTracker::push_for(const std::string &name, const Expr &min, const Expr &max) { + Interval min_bounds = find_constant_bounds(min); + Interval max_bounds = find_constant_bounds(max); + Interval b = Interval::make_union(min_bounds, max_bounds); + b.min = simplify(b.min); + b.max = simplify(b.max); + return Binding(this, ScopedBinding(scope, name, b), false); +} + +BoundsTracker::Binding BoundsTracker::push_let(const std::string &name, const Expr &value) { + bool pure = is_pure(value); + if (pure) { + lets.emplace_back(name, value); + } + return Binding(this, ScopedBinding(scope, name, find_constant_bounds(value)), pure); +} + +BoundsTracker::FactGuard::FactGuard(BoundsTracker *tracker) + : tracker(tracker) { +} + +BoundsTracker::FactGuard::FactGuard(FactGuard &&other) noexcept + : tracker(other.tracker) { + other.tracker = nullptr; +} + +BoundsTracker::FactGuard::~FactGuard() { + if (tracker) { + tracker->facts.pop_back(); + } +} + +BoundsTracker::FactGuard BoundsTracker::push_fact(const Expr &condition) { + facts.push_back(condition); + return FactGuard(this); +} + +Expr BoundsTracker::find_constant_bound(const Expr &e, Direction d) const { + return Halide::Internal::find_constant_bound(e, d, scope); +} + +Interval BoundsTracker::find_constant_bounds(const Expr &e) const { + return Halide::Internal::find_constant_bounds(e, scope); +} + +Expr BoundsTracker::simplify_with_context(const Expr &e) const { + Expr wrapped = e; + for (const auto &[name, value] : reverse_view(lets)) { + wrapped = Let::make(name, value, wrapped); + } + wrapped = remove_likelies(wrapped); + wrapped = substitute_in_all_lets(wrapped); + // Deliberately pass an empty bounds scope here, not `scope`: mixing a + // bounds scope with equality facts can make the simplifier represent a + // variable by its (wide) interval instead of substituting the exact + // value an equality fact implies, which weakens the very reasoning this + // call exists to do. Any tightening from `scope` happens afterward, once + // this expression is no longer being asked to prove an equality. + wrapped = Halide::Internal::simplify(wrapped, Scope::empty_scope(), + Scope::empty_scope(), facts); + return wrapped; +} + +Expr BoundsTracker::find_constant_bound_aggressive(const Expr &e, Direction d) const { + Expr bound = find_constant_bound(e, d); + if (bound.defined()) { + return bound; + } + Expr wrapped = simplify_with_context(e); + return Halide::Internal::find_constant_bound(wrapped, d, scope); +} + +Interval BoundsTracker::find_constant_bounds_aggressive(const Expr &e) const { + Interval interval = find_constant_bounds(e); + if (interval.has_lower_bound() && interval.has_upper_bound()) { + return interval; + } + Expr wrapped = simplify_with_context(e); + return Halide::Internal::find_constant_bounds(wrapped, scope); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h new file mode 100644 index 000000000000..f5af577b8055 --- /dev/null +++ b/src/BoundsTracker.h @@ -0,0 +1,126 @@ +#ifndef HALIDE_BOUNDS_TRACKER_H +#define HALIDE_BOUNDS_TRACKER_H + +/** \file + * A utility for finding constant bounds of expressions at some point inside + * a Stmt tree. + */ + +#include +#include +#include + +#include "Bounds.h" +#include "Expr.h" +#include "Scope.h" + +namespace Halide { +namespace Internal { + +/** Accumulates the bounds-relevant context available at some point inside a + * Stmt tree -- enclosing pure LetStmt/Let bindings and For loop ranges -- as + * a mutator or visitor descends, and uses it to find constant bounds for + * expressions at that point far more reliably than a bare + * find_constant_bound() call. + * + * In addition to the usual scope-based lookup (cheap, but only sees a bound + * if every intermediate variable it passes through was itself pushed with an + * already-constant bound), find_constant_bound_aggressive() falls back to + * literally wrapping an expression in all enclosing pure lets, inlining + * them, and re-simplifying. This is the trick bound_constant_extent_loops + * has always used to find constant loop extents, generalized so other + * passes that infer constant bounds (e.g. BoundSmallAllocations, + * AllocationBoundsInference) can use it too. + */ +class BoundsTracker { +public: + /** An RAII binding produced by push_for/push_let. Pops everything it + * pushed when destroyed. */ + class Binding { + public: + Binding() = default; + Binding(const Binding &) = delete; + Binding &operator=(const Binding &) = delete; + Binding(Binding &&other) noexcept; + ~Binding(); + + private: + friend class BoundsTracker; + Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let); + + BoundsTracker *tracker = nullptr; + ScopedBinding scope_binding; + bool recorded_let = false; + }; + + /** Push the bounds of a for loop variable: the envelope [lower bound of + * min, upper bound of max], each resolved against everything pushed so + * far. */ + Binding push_for(const std::string &name, const Expr &min, const Expr &max); + + /** Push a let binding. Always updates the fast-path scope with a + * constant-bounds estimate of the value. Additionally records the + * syntactic binding for find_constant_bound_aggressive()'s slow path, + * but only if the value is pure -- substituting an impure expression + * into multiple places would change its meaning. */ + Binding push_let(const std::string &name, const Expr &value); + + /** An RAII guard produced by push_fact. Pops the fact when destroyed. */ + class FactGuard { + public: + FactGuard() = default; + FactGuard(const FactGuard &) = delete; + FactGuard &operator=(const FactGuard &) = delete; + FactGuard(FactGuard &&other) noexcept; + ~FactGuard(); + + private: + friend class BoundsTracker; + explicit FactGuard(BoundsTracker *tracker); + + BoundsTracker *tracker = nullptr; + }; + + /** Push a condition known to be true at this point (e.g. because we're + * in the then-case of an IfThenElse that tests it, or it's the + * condition of a dominating assert). Used as a simplifier assumption by + * the slow path in find_constant_bound_aggressive(). Doesn't affect the + * fast-path scope. */ + FactGuard push_fact(const Expr &condition); + + /** Fast path only: find a constant bound using the current scope. See + * find_constant_bound() in Bounds.h. */ + Expr find_constant_bound(const Expr &e, Direction d) const; + Interval find_constant_bounds(const Expr &e) const; + + /** Fast path first; on failure, wrap e in all pending pure lets + * (producing a self-contained copy with no free references to enclosing + * lets), inline them with substitute_in_all_lets, and simplify before + * retrying against the resulting expression and the scope. More + * expensive than find_constant_bound(), but succeeds far more often, + * because the simplifier can cancel terms across let boundaries that + * interval arithmetic through opaque variable lookups cannot. */ + Expr find_constant_bound_aggressive(const Expr &e, Direction d) const; + Interval find_constant_bounds_aggressive(const Expr &e) const; + + /** Wrap e in all pending pure lets, inline them with + * substitute_in_all_lets, and simplify under the current scope and + * dominating facts. Unlike find_constant_bound_aggressive(), the result + * need not be a constant -- this is just a plain simplify() call that + * can see context (enclosing let values, dominating conditions) that a + * caller holding only a bare Expr has no way to pass in. Useful for + * expressions built by context-free helpers (e.g. box_touched) whose + * result ends up referencing variables bound by lets enclosing the + * point the helper was called from. */ + Expr simplify_with_context(const Expr &e) const; + +private: + Scope scope; + std::vector> lets; + std::vector facts; +}; + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 02215600ef8e..407eafccca8b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -71,6 +71,7 @@ target_sources( Bounds.h BoundsInference.h BoundSmallAllocations.h + BoundsTracker.h Buffer.h Callable.h CanonicalizeGPUVars.h @@ -255,6 +256,7 @@ target_sources( Bounds.cpp BoundsInference.cpp BoundSmallAllocations.cpp + BoundsTracker.cpp Buffer.cpp Callable.cpp CanonicalizeGPUVars.cpp From 29f127512c51ca9d846dc13fdc937c4a5b76b7f8 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 25 Aug 2026 20:46:32 +0200 Subject: [PATCH 02/36] Fix use-after-free in BoundConstantExtentLoops's aggressive extent lookup BoundLoops::visit(For*) took bounds.max.as() as a raw pointer while `bounds` was a stack-local Interval about to go out of scope. If that Interval's Expr was the only thing keeping the underlying IntImm node's refcount alive, the node could be freed as soon as `bounds` was destroyed, leaving `e` dangling. The freed memory would typically get reused shortly after (while unwinding through further LetStmt/IfThenElse processing), corrupting the IntImm embedded in the constructed For loop and manifesting later as an infinite Add/Sub/Variable recursion inside Simplify -- reported by the user as a segfault in correctness_unroll_loop_with_implied_constant_bounds, reproduced and fixed with their help using an ASan build. Fixed by copying the Expr into `extent_upper` (already a function-scoped local used for the guarded-upper-bound case) before extracting the raw IntImm pointer from it, so the node stays referenced for the rest of the function regardless of which branch is taken. Also fixes the CMake issue that blocked building an ASan config in the first place: Halide_initmod (the object library holding the runtime's embedded bitcode blobs, linked into the shared Halide target) never had POSITION_INDEPENDENT_CODE set, unlike the Halide target itself. This happened to link fine in optimized builds, where x86-64 codegen tends to use RIP-relative addressing regardless, but failed with absolute 32-bit relocation errors in unoptimized/Debug/ASan builds. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_019EGMmdqNC6mTcMSCDBFbwV --- src/BoundConstantExtentLoops.cpp | 15 ++++++++++----- src/runtime/CMakeLists.txt | 13 ++++++++++++- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index 5691ee575df6..17d896bf3739 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -1,5 +1,4 @@ #include "BoundConstantExtentLoops.h" -#include "Bounds.h" #include "BoundsTracker.h" #include "IRMutator.h" #include "IROperator.h" @@ -63,11 +62,17 @@ class BoundLoops : public IRMutator { debug(4) << "Bounds found: [" << bounds.min << ", " << bounds.max << "]\n"; auto lo = bounds.has_lower_bound() ? as_const_int(bounds.min) : std::nullopt; auto hi = bounds.has_upper_bound() ? as_const_int(bounds.max) : std::nullopt; - if (lo && hi && *lo == *hi) { - // The bound is exact. - e = bounds.max.as(); - } else if (hi) { + if (hi) { + // Copy the Expr out of `bounds` before it goes out of + // scope below -- otherwise e, taken as a raw pointer via + // as(), would be left dangling into a node whose + // only reference was owned by this soon-to-be-destroyed + // Interval. extent_upper = bounds.max; + if (lo && *lo == *hi) { + // The bound is exact: no guard needed. + e = extent_upper.as(); + } } } diff --git a/src/runtime/CMakeLists.txt b/src/runtime/CMakeLists.txt index cac63a328b52..74beb59978c6 100644 --- a/src/runtime/CMakeLists.txt +++ b/src/runtime/CMakeLists.txt @@ -155,7 +155,18 @@ add_library(Halide_initmod OBJECT) add_library(Halide::initmod ALIAS Halide_initmod) # All these are binary2cpp-generated files, so no need to export compile commands for them. -set_target_properties(Halide_initmod PROPERTIES EXPORT_COMPILE_COMMANDS NO) +# POSITION_INDEPENDENT_CODE must be set explicitly: this object library is +# consumed by the Halide target (see POSITION_INDEPENDENT_CODE there), but as +# an OBJECT library its own sources don't automatically inherit that -- without +# it, non-optimized builds (e.g. Debug, ASan) can fail to link into a shared +# Halide with relocation errors on the large embedded byte arrays generated +# here, even though optimized builds usually happen to produce PIC-compatible +# code anyway on x86-64 and don't show the problem. +set_target_properties(Halide_initmod + PROPERTIES + EXPORT_COMPILE_COMMANDS NO + POSITION_INDEPENDENT_CODE ON +) # Note: ensure that these flags match the flags in the Makefile. # Note: this always uses Clang-from-LLVM for compilation, so none of these flags should need conditionalization. From 681d1f15854b7be4afcb8594c90fc6085fa56f60 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 25 Aug 2026 23:38:43 +0200 Subject: [PATCH 03/36] Migrate LowerWarpShuffles onto BoundsTracker Both DetermineAllocStride and LowerWarpShuffles maintained their own Scope bounds, populated on For loops only when the loop's min and max were literal constants (is_const(op->min) && is_const(op->max)), and never updated by LetStmt/Let at all -- so any allocation size or stride computation that depended on a let-bound intermediate value (very common after earlier lowering passes hoist bounds calculations into lets) had no way to resolve to a constant. All uses of `bounds` in this file only ever feed it into simplify() or reduce_expr() (itself simplify()-based), never bounds_of_expr_in_scope() directly -- and Simplify's own internal bounds representation is constant-only anyway (it converts via as_const_int at ingestion), so BoundsTracker's constant-collapsing scope loses nothing here, unlike SlidingWindow/HexagonOptimize which need genuinely symbolic interval tracking BoundsTracker doesn't provide (left unmigrated). Adds two small BoundsTracker capabilities needed by this pass: - interval_scope(): exposes the underlying Scope for passes that feed it directly to simplify()/similar rather than going through find_constant_bound(). - push_interval(): pushes an already-computed Interval directly, for LowerWarpShuffles::visit(IfThenElse*)'s lane-masking case, which narrows an existing binding rather than deriving a new one. Verified with a full correctness suite run under an actual CUDA JIT target (HL_TARGET/HL_JIT_TARGET=host-cuda, reconfigured via -DHalide_TARGET=host-cuda since ctest bakes the target into each test's ENVIRONMENT property at configure time rather than inheriting it from the shell). One unrelated pre-existing failure (correctness_gpu_register_at_block_level, in PromoteGPURegisters/ MultiRamp, which runs before LowerWarpShuffles in the pipeline) was confirmed to reproduce identically against an unmodified origin/main build with the same target override, so it predates this change. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_019EGMmdqNC6mTcMSCDBFbwV --- src/BoundsTracker.cpp | 4 +++ src/BoundsTracker.h | 20 +++++++++++++++ src/LowerWarpShuffles.cpp | 53 +++++++++++++++++++++++---------------- 3 files changed, 55 insertions(+), 22 deletions(-) diff --git a/src/BoundsTracker.cpp b/src/BoundsTracker.cpp index d4d155072f36..63aca5219588 100644 --- a/src/BoundsTracker.cpp +++ b/src/BoundsTracker.cpp @@ -34,6 +34,10 @@ BoundsTracker::Binding BoundsTracker::push_for(const std::string &name, const Ex return Binding(this, ScopedBinding(scope, name, b), false); } +BoundsTracker::Binding BoundsTracker::push_interval(const std::string &name, const Interval &interval) { + return Binding(this, ScopedBinding(scope, name, interval), false); +} + BoundsTracker::Binding BoundsTracker::push_let(const std::string &name, const Expr &value) { bool pure = is_pure(value); if (pure) { diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h index f5af577b8055..44c578530c02 100644 --- a/src/BoundsTracker.h +++ b/src/BoundsTracker.h @@ -58,6 +58,13 @@ class BoundsTracker { * far. */ Binding push_for(const std::string &name, const Expr &min, const Expr &max); + /** Push an already-computed Interval directly, bypassing derivation. + * Useful when a caller has proven a tighter bound for a variable already + * in scope (e.g. because a dominating conditional narrows it) and wants + * to temporarily refine it. Does not participate in + * find_constant_bound_aggressive()'s let-substitution. */ + Binding push_interval(const std::string &name, const Interval &interval); + /** Push a let binding. Always updates the fast-path scope with a * constant-bounds estimate of the value. Additionally records the * syntactic binding for find_constant_bound_aggressive()'s slow path, @@ -114,6 +121,19 @@ class BoundsTracker { * point the helper was called from. */ Expr simplify_with_context(const Expr &e) const; + /** The current fast-path scope of constant bounds, for passes that need + * to feed it directly into simplify() or a similar helper that accepts a + * Scope of assumptions, rather than going through + * find_constant_bound(). Note that Simplify's own internal bounds + * representation is constant-only anyway (it converts via as_const_int + * at ingestion), so this scope -- itself always constant-or-unbounded -- + * loses nothing for that use case. It is not, however, suitable for + * general symbolic interval arithmetic (e.g. bounds_of_expr_in_scope) + * where a non-constant symbolic bound would otherwise be useful. */ + const Scope &interval_scope() const { + return scope; + } + private: Scope scope; std::vector> lets; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index fb800c3f37e4..de0f0e7d9715 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -1,5 +1,6 @@ #include "LowerWarpShuffles.h" +#include "BoundsTracker.h" #include "ExprUsesVar.h" #include "IREquality.h" #include "IRMatch.h" @@ -139,7 +140,7 @@ class DetermineAllocStride : public IRVisitor { // be assumed to be zero. Scope dependent_vars; - Scope bounds; + BoundsTracker tracker; // Get the derivative of an integer expression w.r.t the warp // lane. Returns an undefined Expr if the result is non-trivial. @@ -189,11 +190,13 @@ class DetermineAllocStride : public IRVisitor { void visit(const Let *op) override { ScopedBinding bind(dependent_vars, op->name, warp_stride(op->value)); + auto bounds_bind = tracker.push_let(op->name, op->value); IRVisitor::visit(op); } void visit(const LetStmt *op) override { ScopedBinding bind(dependent_vars, op->name, warp_stride(op->value)); + auto bounds_bind = tracker.push_let(op->name, op->value); IRVisitor::visit(op); } @@ -233,9 +236,7 @@ class DetermineAllocStride : public IRVisitor { } void visit(const For *op) override { - ScopedBinding - bind_bounds_if(is_const(op->min) && is_const(op->max), - bounds, op->name, Interval(op->min, op->max)); + auto bounds_bind = tracker.push_for(op->name, op->min, op->max); ScopedBinding bound_dependent_if((expr_uses_vars(op->min, dependent_vars) || expr_uses_vars(op->max, dependent_vars)), @@ -285,7 +286,7 @@ class DetermineAllocStride : public IRVisitor { // A version of can_prove which exploits the constant bounds we've been tracking bool can_prove(const Expr &e) { - return is_const_one(simplify(e, bounds)); + return is_const_one(simplify(e, tracker.interval_scope())); } Expr get_stride() { @@ -307,7 +308,7 @@ class DetermineAllocStride : public IRVisitor { // any already discovered on previous stores. bool this_ok = (s.defined() && (can_prove(stride == s) && - can_prove(reduce_expr(e / stride - var, warp_size, bounds) == 0))); + can_prove(reduce_expr(e / stride - var, warp_size, tracker.interval_scope()) == 0))); internal_assert(stride.defined()); @@ -331,7 +332,7 @@ class DetermineAllocStride : public IRVisitor { for (const Expr &e : single_stores) { // If only thread zero was active for the store, that makes the proof simpler. Expr simpler = substitute(lane_var, 0, e); - bool this_ok = can_prove(reduce_expr(simpler / stride, warp_size, bounds) == 0); + bool this_ok = can_prove(reduce_expr(simpler / stride, warp_size, tracker.interval_scope()) == 0); if (!this_ok) { bad.push_back(e); } @@ -367,13 +368,21 @@ class LowerWarpShuffles : public IRMutator { Expr stride; }; Scope allocation_info; - Scope bounds; + BoundsTracker tracker; int cuda_cap; + Expr visit(const Let *op) override { + auto binding = tracker.push_let(op->name, op->value); + return IRMutator::visit(op); + } + + Stmt visit(const LetStmt *op) override { + auto binding = tracker.push_let(op->name, op->value); + return IRMutator::visit(op); + } + Stmt visit(const For *op) override { - ScopedBinding - bind_if(is_const(op->min) && is_const(op->max), - bounds, op->name, Interval(op->min, op->max)); + auto bounds_bind = tracker.push_for(op->name, op->min, op->max); if (!this_lane.defined() && op->for_type == ForType::GPULane) { bool should_mask = false; @@ -411,8 +420,8 @@ class LowerWarpShuffles : public IRMutator { // the number of lanes (rounded up). Expr extent = op->extent(); Expr new_size = (alloc->extents[0] + extent - 1) / extent; - new_size = simplify(new_size, bounds); - new_size = find_constant_bound(new_size, Direction::Upper, bounds); + new_size = simplify(new_size, tracker.interval_scope()); + new_size = tracker.find_constant_bound_aggressive(new_size, Direction::Upper); auto sz = as_const_int(new_size); user_assert(sz) << "Warp-level allocation with non-constant size: " << alloc->extents[0] << ". Use Func::bound_extent."; @@ -472,11 +481,11 @@ class LowerWarpShuffles : public IRMutator { if ((lt && equal(lt->a, this_lane) && is_const(lt->b)) || (le && equal(le->a, this_lane) && is_const(le->b))) { Expr condition = mutate(op->condition); - const Interval *in = bounds.find(this_lane_name); + const Interval *in = tracker.interval_scope().find(this_lane_name); internal_assert(in); Interval interval = *in; interval.max = lt ? simplify(lt->b - 1) : le->b; - ScopedBinding bind(bounds, this_lane_name, interval); + auto bind = tracker.push_interval(this_lane_name, interval); Stmt then_case = mutate(op->then_case); Stmt else_case = mutate(op->else_case); return IfThenElse::make(condition, then_case, else_case); @@ -506,7 +515,7 @@ class LowerWarpShuffles : public IRMutator { // of the index and shifting the high bits down to cover // them. Reassembling the result into a flat address gives // the expression below. - Expr in_warp_idx = simplify((idx / (warp_size * stride)) * stride + reduce_expr(idx, stride, bounds), bounds); + Expr in_warp_idx = simplify((idx / (warp_size * stride)) * stride + reduce_expr(idx, stride, tracker.interval_scope()), tracker.interval_scope()); return op->with(value, in_warp_idx, op->predicate, ModulusRemainder()); } else { return IRMutator::visit(op); @@ -531,7 +540,7 @@ class LowerWarpShuffles : public IRMutator { // Load the right lanes from stripe number i equiv = select(idx >= i, make_warp_load(type, name, make_const(idx.type(), i), lane), equiv); } - return simplify(equiv, bounds); + return simplify(equiv, tracker.interval_scope()); } // Load the value to be shuffled @@ -600,7 +609,7 @@ class LowerWarpShuffles : public IRMutator { } else if (expr_match((this_lane + wild) % wild, lane, result) && (bits = is_const_power_of_two_integer(result[1])) && *bits <= 5) { - result[0] = simplify(result[0] % result[1], bounds); + result[0] = simplify(result[0] % result[1], tracker.interval_scope()); // Rotate. Mux a shuffle up and a shuffle down. Uses fewer // intermediate registers than using a general gather for // this. @@ -611,7 +620,7 @@ class LowerWarpShuffles : public IRMutator { shfl_args({membermask, base_val, (1 << *bits) - result[0], 0}), Call::PureExtern); Expr cond = (this_lane >= (1 << *bits) - result[0]); Expr equiv = select(cond, up, down); - shuffled = simplify(equiv, bounds); + shuffled = simplify(equiv, tracker.interval_scope()); } else { // The format of the mask is a pain. The high bits tell // you how large the a warp is for this instruction @@ -641,10 +650,10 @@ class LowerWarpShuffles : public IRMutator { Expr stride = alloc->stride; // Break the index into lane and stripe components - Expr lane = simplify(reduce_expr(idx / stride, warp_size, bounds), bounds); - idx = simplify((idx / (warp_size * stride)) * stride + reduce_expr(idx, stride, bounds), bounds); + Expr lane = simplify(reduce_expr(idx / stride, warp_size, tracker.interval_scope()), tracker.interval_scope()); + idx = simplify((idx / (warp_size * stride)) * stride + reduce_expr(idx, stride, tracker.interval_scope()), tracker.interval_scope()); // We don't want the idx to depend on the lane var, so try to eliminate it - idx = simplify(solve_expression(idx, this_lane_name).result, bounds); + idx = simplify(solve_expression(idx, this_lane_name).result, tracker.interval_scope()); return make_warp_load(op->type, op->name, idx, lane); } else { return IRMutator::visit(op); From 9b952413071f63f4107aff2a373c20b356cab1e7 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 00:24:21 +0200 Subject: [PATCH 04/36] Gather Let/Assert context in SimplifyCorrelatedDifferences via BoundsTracker SimplifyCorrelatedDifferences doesn't just back find_constant_bounds() (via the exported bound_correlated_differences() on a single Expr) -- simplify_correlated_differences() is also run directly as a whole-tree lowering pass in Lower.cpp, several times. Give it a BoundsTracker so it gathers the same constant-bounds context the other migrated passes do, and use it in cancel_correlated_subexpression()'s final simplify() call. This complements rather than replaces the pass's existing `lets` tracking (used to wrap terms for CSE before solve_expression): `lets` deliberately excludes pure lets that are constant w.r.t. the current loop_var, since the monotonicity analysis doesn't need them, but the final simplify() can still benefit from resolving them, and from dominating assert conditions this pass previously never looked at at all (new visit(Block*) override, peeling leading asserts the same way BoundConstantExtentLoops peels dominating if-conditions). Deliberately uses interval_scope()/known_facts() fed straight into simplify(), not find_constant_bound_aggressive()'s more powerful wrap-every-pending-let-and-resimplify path: this pass is already documented as quadratic in loop nesting depth and runs across the whole tree multiple times, so paying that cost on every correlated-difference site would be a real compile-time risk. (Chased what looked like a ~9x compile-time regression down this path during development -- it turned out to be an unrelated Debug-vs-RelWithDebInfo build type mismatch between the two binaries being compared, not anything caused by this change or the earlier migrations; a same-build-type comparison confirms no regression.) Verified with a full correctness suite run (CUDA JIT target enabled): 462/463 passed, with the one failure being the same pre-existing, unrelated GPU register-allocation issue already confirmed to reproduce against an unmodified origin/main build. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_019EGMmdqNC6mTcMSCDBFbwV --- src/BoundsTracker.h | 9 +++++ src/SimplifyCorrelatedDifferences.cpp | 58 +++++++++++++++++++++++---- 2 files changed, 59 insertions(+), 8 deletions(-) diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h index 44c578530c02..a24b4fedc8b4 100644 --- a/src/BoundsTracker.h +++ b/src/BoundsTracker.h @@ -134,6 +134,15 @@ class BoundsTracker { return scope; } + /** The dominating conditions currently known to hold (see push_fact()), + * for passes that want to feed them directly into simplify() as + * assumptions alongside interval_scope(), without paying for the more + * expensive wrap-in-every-pending-let-and-resimplify path that + * find_constant_bound_aggressive()/simplify_with_context() use. */ + const std::vector &known_facts() const { + return facts; + } + private: Scope scope; std::vector> lets; diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index dc90aa6d7cef..d712f8458aee 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -1,5 +1,6 @@ #include "SimplifyCorrelatedDifferences.h" +#include "BoundsTracker.h" #include "CSE.h" #include "ExprUsesVar.h" #include "IRMatch.h" @@ -79,19 +80,34 @@ class SimplifyCorrelatedDifferences : public IRMutator { }; vector lets; + // Tracks constant bounds implied by enclosing lets and dominating + // asserts. This is a whole-tree lowering pass (see + // simplify_correlated_differences() below), not just a helper called + // from find_constant_bounds() on one Expr at a time, so it's worth + // gathering this context as we descend: it lets the final simplify() + // in cancel_correlated_subexpression() resolve free variables that + // `lets` above doesn't bother tracking (lets that are pure and + // constant w.r.t. loop_var are deliberately excluded from `lets`, + // since the monotonicity analysis doesn't need them) and lets it use + // dominating assert conditions it otherwise never sees. + BoundsTracker tracker; + template auto visit_let(const LetStmtOrLet *op) -> decltype(op->body) { // Visit an entire chain of lets in a single method to conserve stack space. struct Frame { const LetStmtOrLet *op; ScopedBinding binding; + BoundsTracker::Binding tracker_binding; Expr new_value; - Frame(const LetStmtOrLet *op, const string &loop_var, Scope &scope) + Frame(const LetStmtOrLet *op, const string &loop_var, Scope &scope, BoundsTracker &tracker) : op(op), - binding(scope, op->name, derivative_bounds(op->value, loop_var, scope)) { + binding(scope, op->name, derivative_bounds(op->value, loop_var, scope)), + tracker_binding(tracker.push_let(op->name, op->value)) { } - Frame(const LetStmtOrLet *op) - : op(op) { + Frame(const LetStmtOrLet *op, BoundsTracker &tracker) + : op(op), + tracker_binding(tracker.push_let(op->name, op->value)) { } }; std::vector frames; @@ -111,13 +127,13 @@ class SimplifyCorrelatedDifferences : public IRMutator { do { result = op->body; if (loop_var.empty()) { - frames.emplace_back(op); + frames.emplace_back(op, tracker); continue; } bool pure = is_pure(op->value); if (!pure || expr_uses_vars(op->value, monotonic) || monotonic.contains(op->name)) { - frames.emplace_back(op, loop_var, monotonic); + frames.emplace_back(op, loop_var, monotonic, tracker); Expr new_value = mutate(op->value); bool may_substitute_in = new_value.type() == Int(32) && pure; lets.emplace_back(OuterLet{op->name, new_value, may_substitute_in}); @@ -126,7 +142,7 @@ class SimplifyCorrelatedDifferences : public IRMutator { // Pure and constant w.r.t the loop var. Doesn't // shadow any outer thing already in the monotonic // scope. - frames.emplace_back(op); + frames.emplace_back(op, tracker); } } while ((op = result.template as())); @@ -155,6 +171,7 @@ class SimplifyCorrelatedDifferences : public IRMutator { } Stmt visit(const For *op) override { + auto bounds_bind = tracker.push_for(op->name, op->min, op->max); Stmt s = op; // This is unfortunately quadratic in maximum loop nesting depth if (loop_var.empty()) { @@ -175,6 +192,25 @@ class SimplifyCorrelatedDifferences : public IRMutator { return s; } + // A leading assert in a Block holds for everything after it (Halide + // blocks are left-leaning, so a chain of asserts shows up as nested + // Blocks, each with an AssertStmt as `first`; recursing into `rest` + // naturally walks the whole chain). + Stmt visit(const Block *op) override { + if (const AssertStmt *a = op->first.as()) { + Stmt first = mutate(op->first); + auto fact = tracker.push_fact(a->condition); + Stmt rest = mutate(op->rest); + if (first.same_as(op->first) && rest.same_as(op->rest)) { + return op; + } else { + return Block::make(first, rest); + } + } else { + return IRMutator::visit(op); + } + } + // Add the names of any free variables in an expr to the provided set void track_free_vars(const Expr &e, std::set *vars) { class TrackFreeVars : public IRVisitor { @@ -227,7 +263,13 @@ class SimplifyCorrelatedDifferences : public IRMutator { e = common_subexpression_elimination(e); e = solve_expression(e, loop_var).result; e = PartiallyCancelDifferences()(e); - e = simplify(e); + // Cheaper than find_constant_bound_aggressive()'s + // wrap-every-pending-let-and-resimplify path (this pass is + // already quadratic in loop nesting depth and runs across the + // whole tree several times, so that would be too expensive + // here): just hand the already-computed constant scope and + // dominating facts to the simplifier directly. + e = simplify(e, tracker.interval_scope(), Scope::empty_scope(), tracker.known_facts()); debug(1) << [&]() -> std::string { if (is_monotonic(e, loop_var) != Monotonic::Unknown) { From 4d7cd442024395c6695fedd63099e3b84cde7c1d Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 12:05:46 +0200 Subject: [PATCH 05/36] Bound a min-clamped ceil-div minus its unclamped multiple A producer compute_at a plain (non-aligned) split tile of its consumer, with the tail handled by PredicateStores, still needs bounds inference to find a compile-time-constant bound for the tile's extent in order to unroll it -- PredicateStores only predicates the consumer's store, not the loads that feed it, so the producer's required region for a boundary tile stays tied to the consumer's declared extent rather than becoming an unconditional full tile. Combined with align_bounds() rounding that region's extent up to a multiple of the tile factor, the resulting extent expression is a ceiling-divide of a min-clamped quantity minus the unclamped multiple the clamp is anchored to: (min(x*c0 + c1, y) + c2)/c0 - x. No existing rule covered it, so bounds inference found no bound at all and unrolling failed outright, even though the region provably fits in one tile. Add that as an exact identity (not just a bound) to SimplifyCorrelatedDifferences's PartiallyCancelDifferences: c0 > 0 means both "+c2" and "/c0" distribute over min, so it reduces to min((y+c2)/c0 - x, (c1+c2)/c0) unconditionally, not just when c1/c2 happen to be multiples of c0. Also push the enclosing loop's own range into BoundConstantExtentLoops' BoundsTracker before recursing into its body, and resimplify the extent with that context before giving up -- matching the pattern BoundSmallAllocations, AllocationBoundsInference, and LowerWarpShuffles already use. Not load-bearing for the new test (the SimplifyCorrelatedDifferences rule alone already finds the bound), but the same class of gap for any nested extent that only resolves once an enclosing loop's bound is in scope. --- src/BoundConstantExtentLoops.cpp | 2 + src/SimplifyCorrelatedDifferences.cpp | 12 +++ test/correctness/CMakeLists.txt | 1 + .../split_predicate_stores_compute_at.cpp | 102 ++++++++++++++++++ 4 files changed, 117 insertions(+) create mode 100644 test/correctness/split_predicate_stores_compute_at.cpp diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index 17d896bf3739..f893e63f9eaa 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -39,6 +39,7 @@ class BoundLoops : public IRMutator { } Stmt visit(const For *op) override { + auto bind = tracker.push_for(op->name, op->min, op->max); Expr extent = simplify(op->extent()); if (is_const(extent)) { // Nothing needs to be done @@ -48,6 +49,7 @@ class BoundLoops : public IRMutator { if (op->for_type == ForType::Unrolled || op->for_type == ForType::Vectorized) { // Give it one last chance to simplify to an int + extent = tracker.simplify_with_context(extent); Stmt body = op->body; const IntImm *e = extent.as(); diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index d712f8458aee..2fdb85a0ec75 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -29,6 +29,7 @@ class PartiallyCancelDifferences : public IRMutator { IRMatcher::Wild<2> z; IRMatcher::WildConst<0> c0; IRMatcher::WildConst<1> c1; + IRMatcher::WildConst<2> c2; Expr visit(const Sub *op) override { @@ -57,6 +58,17 @@ class PartiallyCancelDifferences : public IRMutator { 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) || + // A ceiling-divide of a quantity clamped above by an exact + // multiple of the divisor (x*c0 + c1), minus the unclamped + // multiple (x) that the clamp is anchored to. Comes up when + // rounding a min-clamped region's extent up to a multiple of + // c0: (min(x*c0 + c1, y) + c2)/c0 - x/1 == min((y+c2)/c0 - x, + // (c1+c2)/c0), exactly, for any c1, c2 -- not just when they + // happen to be multiples of c0 -- because c0 > 0 means both + // "+c2" and "/c0" distribute over min. + rewrite((min(x * c0 + c1, y) + c2) / c0 - x, min((y + c2) / c0 - x, fold((c1 + c2) / c0)), c0 > 0) || + rewrite((min(y, x * c0 + c1) + c2) / c0 - x, min((y + c2) / c0 - x, fold((c1 + c2) / c0)), c0 > 0) || + false) { return rewrite.result; } diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d88c13fa177f..24f2487697b9 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -343,6 +343,7 @@ tests( split_factor_type.cpp split_fuse_rvar.cpp split_predicate.cpp + split_predicate_stores_compute_at.cpp split_reuse_inner_name_bug.cpp split_store_compute.cpp stable_realization_order.cpp diff --git a/test/correctness/split_predicate_stores_compute_at.cpp b/test/correctness/split_predicate_stores_compute_at.cpp new file mode 100644 index 000000000000..883d9e07dfd4 --- /dev/null +++ b/test/correctness/split_predicate_stores_compute_at.cpp @@ -0,0 +1,102 @@ +#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; +} + +// Counts IfThenElse nodes reached while inside the named Func's produce node. +class CountIfsInProduce : public IRVisitor { + using IRVisitor::visit; + + std::string name; + int depth = 0; + + void visit(const ProducerConsumer *op) override { + if (op->is_producer && op->name == name) { + depth++; + IRVisitor::visit(op); + depth--; + } else { + IRVisitor::visit(op); + } + } + + void visit(const IfThenElse *op) override { + if (depth > 0) { + count++; + } + IRVisitor::visit(op); + } + +public: + explicit CountIfsInProduce(std::string n) + : name(std::move(n)) { + } + int count = 0; +}; + +} // namespace + +int main(int argc, char **argv) { + // A producer compute_at a plain (non-aligned) split tile of its + // consumer, with the consumer's tail handled by PredicateStores rather + // than GuardWithIf. PredicateStores only predicates the consumer's + // store, not the loads that feed it, so the producer's required region + // for a boundary tile still comes out tied to the consumer's declared + // extent rather than as an unconditional full tile -- bounds inference + // needs to find a compile-time-constant *upper* bound (the split + // factor) for that region's extent to unroll it at all, and fold all + // the per-position validity checks into a single guard around the + // whole tile rather than one nested check per unrolled position. + Var x{"x"}, xo{"xo"}, xi{"xi"}; + Func g{"g"}, f{"f"}; + + g(x) = x * 2; + f(x) = g(x) + 1; + f.output_buffer().dim(0).set_min(0); + + f.split(x, xo, xi, 5, TailStrategy::PredicateStores).never_partition_all(); + g.compute_at(f, xo).align_bounds(x, 5).unroll(x); + + Module m = f.compile_to_module({}); + + CountIfsInProduce checker("g"); + for (const LoweredFunc &lf : m.functions()) { + lf.body.accept(&checker); + } + + // Before the fix, bounds inference couldn't find any bound for the + // tile's extent at all, and unrolling g failed outright. After the + // fix it finds at least an upper bound (the split factor), which + // BoundConstantExtentLoops uses as the unrolled extent, guarded by at + // most one check around the whole tile -- not one check per position. + if (!check(checker.count <= 1, + "expected at most one guard around the whole unrolled tile")) + return 1; + + // Values must still come out right at and past the boundary, for + // several sizes that aren't a multiple of the split factor. + for (int w : {1, 4, 5, 6, 9, 11, 23}) { + Buffer out = f.realize({w}); + for (int i = 0; i < w; i++) { + int expected = i * 2 + 1; + if (out(i) != expected) { + printf("out(%d) = %d instead of %d (w = %d)\n", i, out(i), expected, w); + return 1; + } + } + } + + printf("Success!\n"); + return 0; +} From 2137977b05d06face705312f092ee196acc1ffb5 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 14:27:08 +0200 Subject: [PATCH 06/36] Use loop monotonicity to bound an unrollable tile exactly An unrolled producer tile inside a PredicateStores split gets an extent of the form (min(x*c + c, y) + c)/c*c - x*c, where the enclosing tile loop's own max (a ceiling-divide of y) is exactly what bounds y from below and makes the ceiling-divide exact. BoundConstantExtentLoops could only find the upper bound, so it unrolled to the split factor and wrapped the body in a guard that is always true. Two gaps, both on BoundsTracker's deliberately-expensive slow path: simplify_with_context inlines the enclosing lets, which is what makes the loop variable appear on both sides of the extent's subtraction -- but nothing cancelled it back out. Run bound_correlated_differences and re-simplify, keeping the result only when it actually shrank (it can grow the expression). That turns the extent into min((y + c)/c - x, 1)*c. That form is monotonic in the loop variable, so its extremes over the loop are reached at the ends of the loop's range -- but push_for only recorded a constants-only Interval, which drops the fact that the loop's max mentions y too. Record the range symbolically as well, and have find_constant_bounds_aggressive substitute the endpoints into an expression is_monotonic() says is monotonic in that variable. Substitution keeps y correlated between the expression and the loop bound where per-node interval arithmetic can't, so the extent comes out as exactly [c, c]. This reuses Monotonic.h's existing analysis rather than teaching the simplifier a new kind of fact, and costs nothing until the cheap paths have already failed to find a bound. split_predicate_stores_compute_at now checks for no guard at all inside the unrolled tile, rather than at most one. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_018SvYp54SUAXaJ49MizJg63 --- src/BoundsTracker.cpp | 85 +++++++++++++++++-- src/BoundsTracker.h | 25 +++++- .../split_predicate_stores_compute_at.cpp | 15 ++-- 3 files changed, 111 insertions(+), 14 deletions(-) diff --git a/src/BoundsTracker.cpp b/src/BoundsTracker.cpp index 63aca5219588..15a97e0129c3 100644 --- a/src/BoundsTracker.cpp +++ b/src/BoundsTracker.cpp @@ -1,28 +1,38 @@ #include "BoundsTracker.h" +#include "ExprUsesVar.h" #include "IR.h" #include "IROperator.h" +#include "Monotonic.h" #include "Simplify.h" +#include "SimplifyCorrelatedDifferences.h" #include "Substitute.h" namespace Halide { namespace Internal { -BoundsTracker::Binding::Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let) - : tracker(tracker), scope_binding(std::move(scope_binding)), recorded_let(recorded_let) { +BoundsTracker::Binding::Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let, + bool recorded_loop) + : tracker(tracker), scope_binding(std::move(scope_binding)), recorded_let(recorded_let), + recorded_loop(recorded_loop) { } BoundsTracker::Binding::Binding(Binding &&other) noexcept : tracker(other.tracker), scope_binding(std::move(other.scope_binding)), - recorded_let(other.recorded_let) { + recorded_let(other.recorded_let), + recorded_loop(other.recorded_loop) { other.recorded_let = false; + other.recorded_loop = false; } BoundsTracker::Binding::~Binding() { if (recorded_let) { tracker->lets.pop_back(); } + if (recorded_loop) { + tracker->loops.pop_back(); + } } BoundsTracker::Binding BoundsTracker::push_for(const std::string &name, const Expr &min, const Expr &max) { @@ -31,7 +41,17 @@ BoundsTracker::Binding BoundsTracker::push_for(const std::string &name, const Ex Interval b = Interval::make_union(min_bounds, max_bounds); b.min = simplify(b.min); b.max = simplify(b.max); - return Binding(this, ScopedBinding(scope, name, b), false); + + // Also record the range symbolically. The scope above can only hold + // constants, so it drops any relationship between the loop variable and a + // symbol appearing in its min or max (e.g. the tile index of a split being + // bounded by a ceiling-divide of the extent being split). + bool recorded_loop = false; + if (min.type() == Int(32) && max.type() == Int(32) && is_pure(min) && is_pure(max)) { + loops.push_back(LoopRange{name, min, max}); + recorded_loop = true; + } + return Binding(this, ScopedBinding(scope, name, b), false, recorded_loop); } BoundsTracker::Binding BoundsTracker::push_interval(const std::string &name, const Interval &interval) { @@ -89,6 +109,21 @@ Expr BoundsTracker::simplify_with_context(const Expr &e) const { // this expression is no longer being asked to prove an equality. wrapped = Halide::Internal::simplify(wrapped, Scope::empty_scope(), Scope::empty_scope(), facts); + + // Now that the lets are inlined, a dependence on an enclosing loop + // variable may appear on both sides of a subtraction. Cancel those out + // and simplify again -- the facts recorded by push_for can only relate + // the loop variable to the rest of the expression once the expression + // mentions it directly. This can grow the expression, so only keep the + // result if it actually bought us something. + Expr cancelled = bound_correlated_differences(wrapped); + if (!cancelled.same_as(wrapped)) { + cancelled = Halide::Internal::simplify(cancelled, Scope::empty_scope(), + Scope::empty_scope(), facts); + if (is_const(cancelled) || cancelled.node_type() < wrapped.node_type()) { + wrapped = cancelled; + } + } return wrapped; } @@ -101,13 +136,53 @@ Expr BoundsTracker::find_constant_bound_aggressive(const Expr &e, Direction d) c return Halide::Internal::find_constant_bound(wrapped, d, scope); } +Interval BoundsTracker::tighten_using_loop_monotonicity(const Expr &e, Interval interval) const { + if (e.type() != Int(32)) { + return interval; + } + // Innermost first: the tightest correlation is usually with the nearest + // enclosing loop. + for (const LoopRange &loop : reverse_view(loops)) { + if (interval.has_lower_bound() && interval.has_upper_bound()) { + break; + } + if (!expr_uses_var(e, loop.name)) { + continue; + } + Monotonic m = is_monotonic(e, loop.name); + Expr at_lower, at_upper; + if (m == Monotonic::Increasing || m == Monotonic::Constant) { + at_lower = loop.min; + at_upper = loop.max; + } else if (m == Monotonic::Decreasing) { + at_lower = loop.max; + at_upper = loop.min; + } else { + continue; + } + if (!interval.has_lower_bound()) { + Expr lo = simplify(substitute(loop.name, at_lower, e)); + interval.min = Halide::Internal::find_constant_bounds(lo, scope).min; + } + if (!interval.has_upper_bound()) { + Expr hi = simplify(substitute(loop.name, at_upper, e)); + interval.max = Halide::Internal::find_constant_bounds(hi, scope).max; + } + } + return interval; +} + Interval BoundsTracker::find_constant_bounds_aggressive(const Expr &e) const { Interval interval = find_constant_bounds(e); if (interval.has_lower_bound() && interval.has_upper_bound()) { return interval; } Expr wrapped = simplify_with_context(e); - return Halide::Internal::find_constant_bounds(wrapped, scope); + interval = Halide::Internal::find_constant_bounds(wrapped, scope); + if (interval.has_lower_bound() && interval.has_upper_bound()) { + return interval; + } + return tighten_using_loop_monotonicity(wrapped, interval); } } // namespace Internal diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h index a24b4fedc8b4..ac222989ad09 100644 --- a/src/BoundsTracker.h +++ b/src/BoundsTracker.h @@ -46,16 +46,23 @@ class BoundsTracker { private: friend class BoundsTracker; - Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let); + Binding(BoundsTracker *tracker, ScopedBinding scope_binding, bool recorded_let, + bool recorded_loop = false); BoundsTracker *tracker = nullptr; ScopedBinding scope_binding; bool recorded_let = false; + bool recorded_loop = false; }; /** Push the bounds of a for loop variable: the envelope [lower bound of * min, upper bound of max], each resolved against everything pushed so - * far. */ + * far. Additionally records the range symbolically, so that + * find_constant_bounds_aggressive() can substitute the endpoints into an + * expression that turns out to be monotonic in the loop variable. That + * recovers bounds the constants-only scope can't represent, because a + * loop's min or max may itself mention a symbol the expression also + * mentions. */ Binding push_for(const std::string &name, const Expr &min, const Expr &max); /** Push an already-computed Interval directly, bypassing derivation. @@ -144,9 +151,23 @@ class BoundsTracker { } private: + /** Tighten an interval by exploiting monotonicity in an enclosing loop + * variable: if e is monotonic in it, e's extremes over the loop are + * reached at the ends of the loop's range, so substituting the symbolic + * endpoints and constant-bounding the results can succeed where + * per-node interval arithmetic can't, because substitution keeps a + * symbol shared by e and the loop's range correlated. */ + Interval tighten_using_loop_monotonicity(const Expr &e, Interval interval) const; + + struct LoopRange { + std::string name; + Expr min, max; + }; + Scope scope; std::vector> lets; std::vector facts; + std::vector loops; }; } // namespace Internal diff --git a/test/correctness/split_predicate_stores_compute_at.cpp b/test/correctness/split_predicate_stores_compute_at.cpp index 883d9e07dfd4..2cee901c960f 100644 --- a/test/correctness/split_predicate_stores_compute_at.cpp +++ b/test/correctness/split_predicate_stores_compute_at.cpp @@ -75,13 +75,14 @@ int main(int argc, char **argv) { lf.body.accept(&checker); } - // Before the fix, bounds inference couldn't find any bound for the - // tile's extent at all, and unrolling g failed outright. After the - // fix it finds at least an upper bound (the split factor), which - // BoundConstantExtentLoops uses as the unrolled extent, guarded by at - // most one check around the whole tile -- not one check per position. - if (!check(checker.count <= 1, - "expected at most one guard around the whole unrolled tile")) + // The tile's extent is exactly the split factor: the enclosing tile + // loop's own max bounds the consumer's extent from below, so the + // ceiling-divide that rounds the region up to a multiple of the factor + // is exact. BoundConstantExtentLoops must find that as an exact + // constant, not just an upper bound, so the unrolled body needs no + // guard at all. + if (!check(checker.count == 0, + "expected no guard inside the unrolled tile")) return 1; // Values must still come out right at and past the boundary, for From a5d88de33f86a4f4ca2d6fca3e5e0d4dea6cd182 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 17 Aug 2026 20:53:07 +0200 Subject: [PATCH 07/36] Move one simplification call in lowering. Co-authored-by: Andrew Adams --- src/Lower.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Lower.cpp b/src/Lower.cpp index 753dadb5f6ec..3e44fa6ad7a6 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -450,13 +450,13 @@ void lower_impl(const vector &output_funcs, if (t.has_feature(Target::Profile) || t.has_feature(Target::ProfileByTimer)) { debug(1) << "Injecting profiling...\n"; s = inject_profiling(s, pipeline_name, env, t); - s = simplify(s); log("Lowering after injecting profiling:", s); } debug(1) << "Finding intrinsics...\n"; // Must be run after the last simplification, because it turns // divisions into shifts, which the simplifier reverses. + s = simplify(s); s = find_intrinsics(s); log("Lowering after finding intrinsics:", s); From cf3e0227a1423fa32588af9aeddd4b64446ccd5a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 16:11:00 +0200 Subject: [PATCH 08/36] Move printing of the final simplification and call it 'after reaching conceptual stmt' --- src/Lower.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/Lower.cpp b/src/Lower.cpp index 3e44fa6ad7a6..21acfff8df2d 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -470,9 +470,6 @@ void lower_impl(const vector &output_funcs, log("Lowering after stripping asserts:", s); } - debug(1) << "Lowering after final simplification:\n" - << s << "\n\n"; - if (!custom_passes.empty()) { for (size_t i = 0; i < custom_passes.size(); i++) { debug(1) << "Running custom lowering pass " << i << "...\n"; @@ -484,6 +481,8 @@ void lower_impl(const vector &output_funcs, // Make a copy of the Stmt code, before we lower anything to less human-readable code. result_module.set_conceptual_code_stmt(s); + debug(1) << "Lowering after reaching conceptual Stmt:\n" + << s << "\n\n"; if (t.arch != Target::Hexagon && t.has_feature(Target::HVX)) { debug(1) << "Splitting off Hexagon offload...\n"; From d9788802306264a510bdb140d43bed086b9b113d Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 17 Aug 2026 20:51:50 +0200 Subject: [PATCH 09/36] 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 24f2487697b9..4f518ad4cbc1 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 891e8e45f706f8b655ef32715c0dd7fe58ebb732 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 16:04:29 +0200 Subject: [PATCH 10/36] Add enable backtraces to async_copy_chain as I have found it to be stalled/deadlocked several times but couldn't debug it. --- test/correctness/async_copy_chain.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/correctness/async_copy_chain.cpp b/test/correctness/async_copy_chain.cpp index 238b1eb8821b..efab90f153cd 100644 --- a/test/correctness/async_copy_chain.cpp +++ b/test/correctness/async_copy_chain.cpp @@ -5,7 +5,8 @@ using namespace Halide; Var x, y; void check(Func f) { - Buffer out = f.realize({256, 256}); + Target target = get_jit_target_from_environment().with_feature(Target::EnableBacktraces); + Buffer out = f.realize({256, 256}, target); out.for_each_element([&](int x, int y) { if (out(x, y) != x + y) { printf("out(%d, %d) = %d instead of %d\n", x, y, out(x, y), x + y); From cc301254e98b45a2531b42459da25d4e2686cee8 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 19:57:15 +0200 Subject: [PATCH 11/36] 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 c903c8c6f6680fd4de1b700e07898b6526a46856 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 21:48:45 +0200 Subject: [PATCH 12/36] 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 fad568a0b3f72917ae3af5ff22afee6f493a6217 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 21:49:04 +0200 Subject: [PATCH 13/36] 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 973eecb72b8c624c63b18768ff5cd110d54c2c84 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:48:57 +0200 Subject: [PATCH 14/36] 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 4c3dd9715961cd2c1d2a9ea72d4ba319b06f5d5b Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:51:44 +0200 Subject: [PATCH 15/36] 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 4f518ad4cbc1..2c5df5528413 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -484,6 +484,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 6b38cffdfb672f52c867115f8ee8db2b517a94b6 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:54:25 +0200 Subject: [PATCH 16/36] 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 2c5df5528413..e22fdba12521 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 41d3858e6a97010285a2126208ef3a36d6584ad1 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 22:58:52 +0200 Subject: [PATCH 17/36] 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 7cece554806b53a7da21b27f4ecd6053c4ac4dd5 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 23:36:57 +0200 Subject: [PATCH 18/36] 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 c929d18a3bd5d4a3e932d85ab0e7a6678027a656 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 20 Aug 2026 23:47:20 +0200 Subject: [PATCH 19/36] 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 e22fdba12521..7b01b5122b1a 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -487,6 +487,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 7bfc6f7dee87c375d6bc10191f1bab2307b154b5 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 19 Aug 2026 17:00:59 +0200 Subject: [PATCH 20/36] 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 7d1867cd569ff50e1574f5773ee712b55249bb69 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Fri, 21 Aug 2026 00:06:08 +0200 Subject: [PATCH 21/36] 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 7b01b5122b1a..9bedfe0ed17f 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 bf05e6c9f92050fc547b1187e6643737ddbcd9c0 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 24 Aug 2026 15:04:12 +0200 Subject: [PATCH 22/36] WIP checkpoint: aligned-split debug tracing, simplifier rules, 6x6 compute_at test Checkpoint of in-progress work before investigating the surviving mux under compute_at with aligned splits. --- .gitignore | 1 + src/BoundConstantExtentLoops.cpp | 13 +-- src/CMakeLists.txt | 2 +- src/SimplifyCorrelatedDifferences.cpp | 3 + src/Simplify_Add.cpp | 3 + src/runtime/thread_pool_common.h | 102 ++------------------- test/correctness/CMakeLists.txt | 1 + test/correctness/split_aligned_2d_6x6.cpp | 105 ++++++++++++++++++++++ 8 files changed, 122 insertions(+), 108 deletions(-) create mode 100644 test/correctness/split_aligned_2d_6x6.cpp diff --git a/.gitignore b/.gitignore index 4a013d25b796..b198de54cd0b 100644 --- a/.gitignore +++ b/.gitignore @@ -242,6 +242,7 @@ xcuserdata # NeoVim + clangd .cache +.ccls-cache # Emacs tags diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index f893e63f9eaa..0162dcd17a80 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -61,7 +61,7 @@ class BoundLoops : public IRMutator { debug(4) << "Trying to find a constant bound for loop " << op->name << "\n" << "Extent: " << extent << "\n"; Interval bounds = tracker.find_constant_bounds_aggressive(extent); - debug(4) << "Bounds found: [" << bounds.min << ", " << bounds.max << "]\n"; + debug(4) << "Bounds found: " << bounds << "\n"; auto lo = bounds.has_lower_bound() ? as_const_int(bounds.min) : std::nullopt; auto hi = bounds.has_upper_bound() ? as_const_int(bounds.max) : std::nullopt; if (hi) { @@ -78,17 +78,6 @@ class BoundLoops : public IRMutator { } } - if (e == nullptr && extent_upper.defined()) { - // Still no luck getting an exact extent. Take the upper - // bound instead and guard the body with an if statement. - debug(4) << "Found an upper bound instead: " << extent_upper << "\n"; - e = extent_upper.as(); - body = - IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <= - op->max), - body); - } - if (e == nullptr && permit_failed_unroll && op->for_type == ForType::Unrolled) { // Still no luck, but we're allowed to fail. Rewrite // to a serial loop. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 407eafccca8b..3bdb9424e809 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -549,7 +549,7 @@ if (WITH_SERIALIZATION) find_package(FlatBuffers REQUIRED) _Halide_pkgdep(FlatBuffers NAMESPACE flatbuffers) - target_link_libraries(Halide PRIVATE flatbuffers::flatbuffers) + target_link_libraries(Halide PRIVATE flatbuffers) set(fb_def "${CMAKE_CURRENT_SOURCE_DIR}/halide_ir.fbs") set(fb_dir "${Halide_BINARY_DIR}/include/flatc") diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index 2fdb85a0ec75..be06029801ad 100644 --- a/src/SimplifyCorrelatedDifferences.cpp +++ b/src/SimplifyCorrelatedDifferences.cpp @@ -54,6 +54,8 @@ class PartiallyCancelDifferences : public IRMutator { rewrite(max(x, c0) - min(x, c1), max(max(c0 - x, x - c1), fold(max(0, c0 - c1)))) || 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)) || + rewrite(min(x + c0, y) - max(x, z), min(min(x + c0, y) - max(x, z), c0)) || + rewrite(max(x + c0, y) - min(x, z), max(max(x + c0, y) - min(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) || @@ -70,6 +72,7 @@ class PartiallyCancelDifferences : public IRMutator { rewrite((min(y, x * c0 + c1) + c2) / c0 - x, min((y + c2) / c0 - x, fold((c1 + c2) / c0)), c0 > 0) || false) { + debug(4) << "Rewrote " << Expr(op) << " as " << rewrite.result << "\n"; return rewrite.result; } } diff --git a/src/Simplify_Add.cpp b/src/Simplify_Add.cpp index a2298c2e019c..632d91d80d74 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -203,6 +203,9 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { 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(min(x, 0) + min(-x + c0, y), min(min(min(c0, y), x + y), c0 - x)) || + rewrite(min(x - z, 0) + min((z - x) + c0, y), min(min(min(c0, y), x - z + y), c0 - z - x)) || + false)))) { return mutate(rewrite.result, info); } diff --git a/src/runtime/thread_pool_common.h b/src/runtime/thread_pool_common.h index f8ffaa1f64a7..51c6bcf34b27 100644 --- a/src/runtime/thread_pool_common.h +++ b/src/runtime/thread_pool_common.h @@ -109,26 +109,6 @@ struct work { } }; -// A thread that stalls on a job it owns may start tasks belonging to other -// parallel regions rather than idle, but doing so nests a new owned job inside -// the one it is already waiting on, and that stack frame can't unwind until the -// new job completes. Bounding how many jobs a thread may have stalled at once -// lets it start one more instance of an outer loop it is already inside, but -// not an unbounded number of them. Descending without stalling is already -// bounded, because it can only go one level deeper into the loop nest. -constexpr int max_stalled_jobs = 2; - -// Which entry of work_queue.stalled_jobs belongs to the calling thread. Thread -// ids are dense on some platforms and multiples of four on others, so mix them -// before taking the low bits. Two threads may share an entry, which can only -// cause one of them to decline to start a job it could have run. -ALWAYS_INLINE int stalled_jobs_slot() { - static_assert(MAX_THREADS <= 256 && (MAX_THREADS & (MAX_THREADS - 1)) == 0, - "MAX_THREADS must be a power of two no greater than 256."); - const uint32_t id = (uint32_t)halide_current_thread_id(); - return (int)(((id * (uint32_t)2654435761) >> 24) & (MAX_THREADS - 1)); -} - ALWAYS_INLINE int clamp_num_threads(int threads) { if (threads > MAX_THREADS) { return MAX_THREADS; @@ -181,24 +161,10 @@ struct work_queue_t { // must signal or broadcast the appropriate condition variable. halide_cond_with_spinning wake_a_team, wake_b_team, wake_owners; - // A separate channel for workers that found a job they could run but - // for an unavailable semaphore. This is distinct from the A/B teams, - // which model idle capacity (no runnable work): a worker here is - // blocked on an external event, not idle. Keeping it separate lets a - // semaphore release wake exactly the workers waiting on a semaphore, - // without a thundering herd of genuinely-idle workers waking only to - // rescan and go back to sleep. - halide_cond_with_spinning wake_from_semaphore; - // The number of sleeping workers and owners. An over-estimate - a // waking-up thread may not have decremented this yet. int workers_sleeping, owners_sleeping; - // The number of workers parked on wake_from_semaphore. A subset of - // workers_sleeping (those threads are also counted there, so the A/B - // team bookkeeping is undisturbed). - int workers_parked_on_semaphore; - // Keep track of threads so they can be joined at shutdown halide_thread *threads[MAX_THREADS]; @@ -212,12 +178,6 @@ struct work_queue_t { // to prevent deadlock due to oversubscription of threads. int threads_reserved; - // For each thread, how many of the jobs it owns have stalled, indexed by - // stalled_jobs_slot(). Only maintained for jobs that stall, because - // finding the index can cost a syscall, and a stalled job has nothing - // better to do. - uint8_t stalled_jobs[MAX_THREADS]; - ALWAYS_INLINE bool running() const { return !shutdown; } @@ -301,31 +261,11 @@ WEAK void worker_thread_idle() { work_queue.workers_sleeping--; } -// A worker that found runnable work blocked only on an unavailable -// semaphore. Unlike an idle worker, it must be woken by a semaphore -// release, so it waits on its own channel rather than the A/B teams. -WEAK void worker_thread_blocked_on_semaphore() { - work_queue.workers_sleeping++; - work_queue.workers_parked_on_semaphore++; - work_queue.wake_from_semaphore.wait(&work_queue.mutex); - work_queue.workers_parked_on_semaphore--; - work_queue.workers_sleeping--; -} - WEAK void worker_thread_already_locked(work *owned_job) { - // Set the first time this job stalls. Threads that don't own a job are - // free to work on anything, so they never need it. - int slot = -1; - while (owned_job ? owned_job->running() : !work_queue.shutdown) { work *job = work_queue.jobs; work **prev_ptr = &work_queue.jobs; - // Did we pass over a job that we could otherwise run, but for an - // unavailable semaphore? If so, a future semaphore release (not - // just newly-enqueued work) can make us runnable. - bool blocked_on_semaphore = false; - if (owned_job) { if (owned_job->exit_status != halide_error_code_success) { if (owned_job->active_workers == 0) { @@ -376,19 +316,7 @@ WEAK void worker_thread_already_locked(work *owned_job) { if (!enough_threads) { log_message("Not enough threads for job " << job->task.name << " available: " << threads_available << " min_threads: " << job->task.min_threads); } - // Starting a job from another parallel region leaves the job we - // own waiting on our stack, so only do it for jobs that can't - // block, and only if we aren't already holding one that way. - bool can_use_this_thread_stack = - !owned_job || (job->siblings == owned_job->siblings); - if (!can_use_this_thread_stack && job->task.min_threads == 0) { - if (slot < 0) { - slot = stalled_jobs_slot(); - work_queue.stalled_jobs[slot]++; - } - can_use_this_thread_stack = - work_queue.stalled_jobs[slot] < max_stalled_jobs; - } + bool can_use_this_thread_stack = !owned_job || (job->siblings == owned_job->siblings) || job->task.min_threads == 0; if (!can_use_this_thread_stack) { log_message("Cannot run job " << job->task.name << " on this thread."); } @@ -402,7 +330,6 @@ WEAK void worker_thread_already_locked(work *owned_job) { break; } else { log_message("Cannot acquire semaphores for " << job->task.name); - blocked_on_semaphore = true; } } prev_ptr = &(job->next_job); @@ -416,8 +343,6 @@ WEAK void worker_thread_already_locked(work *owned_job) { // is very informative when profiling. if (owned_job) { worker_thread_stall(owned_job); - } else if (blocked_on_semaphore) { - worker_thread_blocked_on_semaphore(); } else { worker_thread_idle(); } @@ -542,10 +467,6 @@ WEAK void worker_thread_already_locked(work *owned_job) { work_queue.wake_owners.broadcast(); } } - - if (slot >= 0) { - work_queue.stalled_jobs[slot]--; - } } WEAK void worker_thread(void *arg) { @@ -673,13 +594,6 @@ WEAK void enqueue_work_already_locked(int num_jobs, work *jobs, work *task_paren } } - // Workers blocked on a semaphore wait on their own channel, so the - // broadcasts above don't reach them. Wake them too: they may be able - // to steal this newly-enqueued work. - if (work_queue.workers_parked_on_semaphore) { - work_queue.wake_from_semaphore.broadcast(); - } - if (job_has_acquires || job_may_block) { if (task_parent != nullptr) { task_parent->threads_reserved--; @@ -828,7 +742,6 @@ WEAK void halide_shutdown_thread_pool() { work_queue.wake_owners.broadcast(); work_queue.wake_a_team.broadcast(); work_queue.wake_b_team.broadcast(); - work_queue.wake_from_semaphore.broadcast(); halide_mutex_unlock(&work_queue.mutex); // Wait until they leave @@ -854,14 +767,13 @@ WEAK int halide_default_semaphore_init(halide_semaphore_t *s, int n) { WEAK int halide_default_semaphore_release(halide_semaphore_t *s, int n) { halide_semaphore_impl_t *sem = (halide_semaphore_impl_t *)s; int old_val = Halide::Runtime::Internal::Synchronization::atomic_fetch_add_acquire_release(&sem->value, n); - if (n != 0) { + // TODO(abadams|zvookin): Is this correct if an acquire can be for say count of 2 and the releases are 1 each? + if (old_val == 0 && n != 0) { // Don't wake if nothing released. + // We may have just made a job runnable halide_mutex_lock(&work_queue.mutex); - if (work_queue.workers_parked_on_semaphore) { - work_queue.wake_from_semaphore.broadcast(); - } - if (work_queue.owners_sleeping) { - work_queue.wake_owners.broadcast(); - } + work_queue.wake_a_team.broadcast(); + work_queue.wake_b_team.broadcast(); + work_queue.wake_owners.broadcast(); halide_mutex_unlock(&work_queue.mutex); } return old_val + n; diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 9bedfe0ed17f..836808d533c7 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_6x6.cpp split_aligned_nested.cpp split_aligned_reduction.cpp split_by_non_factor.cpp diff --git a/test/correctness/split_aligned_2d_6x6.cpp b/test/correctness/split_aligned_2d_6x6.cpp new file mode 100644 index 000000000000..8261ff6b9a6d --- /dev/null +++ b/test/correctness/split_aligned_2d_6x6.cpp @@ -0,0 +1,105 @@ +#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 % 6) * 6 + y * (i % 6); +} + +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"}; + // offset_x.set_range(0, 5); + // offset_y.set_range(0, 5); + auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { + return (6 * ((y - offset_y) % 6)) + ((x - offset_x) % 6); + }; + std::vector ways; + ways.reserve(36); + for (int i = 0; i < 36; ++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); + + f + .split(x, xo, xi, 6, offset_x, Halide::TailStrategy::GuardWithIf) + .split(y, yo, yi, 6, offset_y, Halide::TailStrategy::GuardWithIf) + .never_partition_all() + .reorder(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 &f : module.functions()) { + f.body.accept(&checker); + } + + for (int i = 0; i < 4; i++) { + printf("Testing runtime alignment: x=%d y=%d\n", i / 6, i % 6); + offset_x.set(i / 6); + offset_y.set(i % 6); + Buffer im = f.realize({32, 32, 3}); + f.realize(im, get_target_from_environment()); + + for (int y = 0; y < 32; y++) { + for (int x = 0; x < 32; x++) { + int selector = idx(6 + x, 6 + y, offset_x.get(), offset_y.get()); + int expected = produce_mux_argument(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 != 0) { + std::printf("Expected 0 muxes: %d\n", checker.mux_count); + return 1; + } + if (checker.for_count != 1) { + std::printf("Expected 3 for loops: %d\n", checker.for_count); + return 1; + } + + printf("Success!\n"); + return 0; +} From 3823f186d5e79f6ba5722a8998e5f8d88a36dfb2 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 24 Aug 2026 15:05:40 +0200 Subject: [PATCH 23/36] 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 | 2 +- test/correctness/split_aligned_2d_3x3.cpp | 115 ++++++++++++++++++++++ test/correctness/split_aligned_2d_6x6.cpp | 105 -------------------- 3 files changed, 116 insertions(+), 106 deletions(-) create mode 100644 test/correctness/split_aligned_2d_3x3.cpp delete mode 100644 test/correctness/split_aligned_2d_6x6.cpp diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 836808d533c7..c0509796813f 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -341,7 +341,7 @@ tests( spirv_ir.cpp split_aligned.cpp split_aligned_2d.cpp - split_aligned_2d_6x6.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; +} diff --git a/test/correctness/split_aligned_2d_6x6.cpp b/test/correctness/split_aligned_2d_6x6.cpp deleted file mode 100644 index 8261ff6b9a6d..000000000000 --- a/test/correctness/split_aligned_2d_6x6.cpp +++ /dev/null @@ -1,105 +0,0 @@ -#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 % 6) * 6 + y * (i % 6); -} - -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"}; - // offset_x.set_range(0, 5); - // offset_y.set_range(0, 5); - auto idx = [](const auto &x, const auto &y, const auto &offset_x, const auto &offset_y) { - return (6 * ((y - offset_y) % 6)) + ((x - offset_x) % 6); - }; - std::vector ways; - ways.reserve(36); - for (int i = 0; i < 36; ++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); - - f - .split(x, xo, xi, 6, offset_x, Halide::TailStrategy::GuardWithIf) - .split(y, yo, yi, 6, offset_y, Halide::TailStrategy::GuardWithIf) - .never_partition_all() - .reorder(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 &f : module.functions()) { - f.body.accept(&checker); - } - - for (int i = 0; i < 4; i++) { - printf("Testing runtime alignment: x=%d y=%d\n", i / 6, i % 6); - offset_x.set(i / 6); - offset_y.set(i % 6); - Buffer im = f.realize({32, 32, 3}); - f.realize(im, get_target_from_environment()); - - for (int y = 0; y < 32; y++) { - for (int x = 0; x < 32; x++) { - int selector = idx(6 + x, 6 + y, offset_x.get(), offset_y.get()); - int expected = produce_mux_argument(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 != 0) { - std::printf("Expected 0 muxes: %d\n", checker.mux_count); - return 1; - } - if (checker.for_count != 1) { - std::printf("Expected 3 for loops: %d\n", checker.for_count); - return 1; - } - - printf("Success!\n"); - return 0; -} From 9d5a33f6056f5a0fdf6b3f17e977f612272f3ca3 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 24 Aug 2026 15:34:03 +0200 Subject: [PATCH 24/36] Give producers a tile-aligned region inside an aligned split An aligned split iterates over whole tiles anchored at align, but the provides it makes are clamped to the Func's own bounds. Bounds inference derived the producing stage's region at each loop level from those provides, so a Func computed inside the split got a region starting at max(outer * factor + align, old_min) -- no longer congruent to align modulo factor. That defeats the point of an aligned split. A producer indexing on old_var % factor keeps a non-constant index after unrolling, so a mux over the tile never folds to the single way it selects. Loop partitioning would otherwise carve out a steady-state region where the clamp is provably dead, so the problem only showed up under never_partition_all(), and only when the producer was scheduled with compute_at rather than inlined. Record the tile alongside the clamped promise and prefer it when defining the producing stage's bounds, in the same style as the existing .guarded lookup. The provide keeps its [old_min, old_max] promise, which is what confines the stores and keeps the output buffer's required region correct. This trades a little compute for the alignment: boundary tiles are now computed in full, so a producer reading an input buffer requires up to factor-1 more of it on each side, as it would under RoundUp. Fixes correctness_split_aligned_2d_3x3, which now sees zero muxes. --- src/ApplySplit.cpp | 20 ++++++++++++++++++++ src/BoundsInference.cpp | 20 ++++++++++++++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index 6df71d634b9e..fb595e950a20 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -79,6 +79,26 @@ vector apply_split(const Split &split, const string &prefix, // 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); + + // The clamp above is what the *stores* of this stage are + // confined to, but it is not the region this stage iterates + // over: an aligned split iterates whole tiles anchored at + // align, and the boundary tiles hang off the ends of + // [old_min, old_max]. Record the tile so that bounds inference + // can give anything computed inside this split a tile-aligned + // region. + // + // Without it, a producer's region starts at + // max(outer * factor + align, old_min), which is no longer + // congruent to align modulo factor. A producer indexing on + // `old_var % factor` -- the whole point of an aligned split -- + // then keeps a non-constant index after unrolling, so e.g. a + // mux over the tile never folds away. + result.emplace_back(prefix + split.old_var + ".aligned_min", + base_var + split.align, ApplySplitResult::LetStmt); + result.emplace_back(prefix + split.old_var + ".aligned_max", + base_var + split.align + split.factor - 1, + ApplySplitResult::LetStmt); } else { // Legacy: structurally guaranteed to be >= old_min guarded = promise_clamped(old_var, old_var, old_max); diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index ba8266883e4b..90d47a78613c 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -1177,13 +1177,29 @@ class BoundsInference : public IRMutator { internal_assert(box[i].is_bounded()); string var = b.first + "." + f_args[i]; - if (box[i].is_single_point()) { + // An aligned split iterates over whole tiles, but the + // provides it makes are clamped to this Func's own + // bounds, so box[i] describes the clamped region. Use + // the tile instead where one was recorded, so that + // producers computed inside the split get a + // tile-aligned region and can fold away indexing that + // is periodic in the tile. + Expr aligned_min, aligned_max; + if (let_vars_in_scope.contains(var + ".aligned_min") && + let_vars_in_scope.contains(var + ".aligned_max")) { + aligned_min = Variable::make(Int(32), var + ".aligned_min"); + aligned_max = Variable::make(Int(32), var + ".aligned_max"); + } + + if (aligned_max.defined()) { + body = LetStmt::make(var + ".max", aligned_max, body); + } else if (box[i].is_single_point()) { body = LetStmt::make(var + ".max", Variable::make(Int(32), var + ".min"), body); } else { body = LetStmt::make(var + ".max", box[i].max, body); } - body = LetStmt::make(var + ".min", box[i].min, body); + body = LetStmt::make(var + ".min", aligned_min.defined() ? aligned_min : box[i].min, body); } } } From f0c047b39b550a8ad500884da62fd05013a06d30 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Mon, 24 Aug 2026 16:19:37 +0200 Subject: [PATCH 25/36] Starting point. --- src/ApplySplit.cpp | 20 -------------------- src/BoundsInference.cpp | 20 ++------------------ test/correctness/split_aligned_2d_3x3.cpp | 9 ++++++++- 3 files changed, 10 insertions(+), 39 deletions(-) diff --git a/src/ApplySplit.cpp b/src/ApplySplit.cpp index fb595e950a20..6df71d634b9e 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -79,26 +79,6 @@ vector apply_split(const Split &split, const string &prefix, // 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); - - // The clamp above is what the *stores* of this stage are - // confined to, but it is not the region this stage iterates - // over: an aligned split iterates whole tiles anchored at - // align, and the boundary tiles hang off the ends of - // [old_min, old_max]. Record the tile so that bounds inference - // can give anything computed inside this split a tile-aligned - // region. - // - // Without it, a producer's region starts at - // max(outer * factor + align, old_min), which is no longer - // congruent to align modulo factor. A producer indexing on - // `old_var % factor` -- the whole point of an aligned split -- - // then keeps a non-constant index after unrolling, so e.g. a - // mux over the tile never folds away. - result.emplace_back(prefix + split.old_var + ".aligned_min", - base_var + split.align, ApplySplitResult::LetStmt); - result.emplace_back(prefix + split.old_var + ".aligned_max", - base_var + split.align + split.factor - 1, - ApplySplitResult::LetStmt); } else { // Legacy: structurally guaranteed to be >= old_min guarded = promise_clamped(old_var, old_var, old_max); diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index 90d47a78613c..ba8266883e4b 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -1177,29 +1177,13 @@ class BoundsInference : public IRMutator { internal_assert(box[i].is_bounded()); string var = b.first + "." + f_args[i]; - // An aligned split iterates over whole tiles, but the - // provides it makes are clamped to this Func's own - // bounds, so box[i] describes the clamped region. Use - // the tile instead where one was recorded, so that - // producers computed inside the split get a - // tile-aligned region and can fold away indexing that - // is periodic in the tile. - Expr aligned_min, aligned_max; - if (let_vars_in_scope.contains(var + ".aligned_min") && - let_vars_in_scope.contains(var + ".aligned_max")) { - aligned_min = Variable::make(Int(32), var + ".aligned_min"); - aligned_max = Variable::make(Int(32), var + ".aligned_max"); - } - - if (aligned_max.defined()) { - body = LetStmt::make(var + ".max", aligned_max, body); - } else if (box[i].is_single_point()) { + if (box[i].is_single_point()) { body = LetStmt::make(var + ".max", Variable::make(Int(32), var + ".min"), body); } else { body = LetStmt::make(var + ".max", box[i].max, body); } - body = LetStmt::make(var + ".min", aligned_min.defined() ? aligned_min : box[i].min, body); + body = LetStmt::make(var + ".min", box[i].min, body); } } } diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp index f72ca58c3c9e..d93d80493373 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,12 @@ 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) + .unroll(y) + .never_partition_all() + .align_bounds(x, 3, offset_x) + .align_bounds(y, 3, offset_y); } Module module = f.compile_to_module({offset_x, offset_y}); From 9ee4682c8687d7f9c016d5b69c9720b342b94ea7 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 08:50:34 +0200 Subject: [PATCH 26/36] Fix bad merge. --- src/BoundConstantExtentLoops.cpp | 12 ++++++++++++ src/BoundSmallAllocations.cpp | 1 + src/Bounds.cpp | 6 ++++++ src/BoundsTracker.cpp | 1 + src/Simplify_Sub.cpp | 3 +++ 5 files changed, 23 insertions(+) diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index 0162dcd17a80..29232242b028 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -2,6 +2,7 @@ #include "BoundsTracker.h" #include "IRMutator.h" #include "IROperator.h" +#include "IRPrinter.h" #include "Simplify.h" #include "Util.h" @@ -78,6 +79,17 @@ class BoundLoops : public IRMutator { } } + if (e == nullptr && extent_upper.defined()) { + // Still no luck getting an exact extent. Take the upper + // bound instead and guard the body with an if statement. + debug(4) << "Found an upper bound instead: " << extent_upper << "\n"; + e = extent_upper.as(); + body = + IfThenElse::make(likely_if_innermost(Variable::make(Int(32), op->name) <= + op->max), + body); + } + if (e == nullptr && permit_failed_unroll && op->for_type == ForType::Unrolled) { // Still no luck, but we're allowed to fail. Rewrite // to a serial loop. diff --git a/src/BoundSmallAllocations.cpp b/src/BoundSmallAllocations.cpp index 9c987c04dd6e..b59240f14edc 100644 --- a/src/BoundSmallAllocations.cpp +++ b/src/BoundSmallAllocations.cpp @@ -110,6 +110,7 @@ class BoundSmallAllocations : public IRMutator { for (const Expr &e : op->extents) { total_extent *= e; } + debug(3) << "Finding constant bound for Allocation " << op->name << " with extent " << total_extent << "\n"; Expr bound = tracker.find_constant_bound_aggressive(total_extent, Direction::Upper); if (!bound.defined() && must_be_constant(op->memory_type)) { diff --git a/src/Bounds.cpp b/src/Bounds.cpp index 7aeff2f87bad..6f02de82d708 100644 --- a/src/Bounds.cpp +++ b/src/Bounds.cpp @@ -119,9 +119,15 @@ Expr find_constant_bound(const Expr &e, Direction d, const Scope &scop } Interval find_constant_bounds(const Expr &e, const Scope &scope) { + debug(4) << "find_constant_bounds for " << e << "\n"; + for (auto it = scope.cbegin(); it != scope.cend(); ++it) { + debug(5) << " with " << it.name() << " in " << it.value() << "\n"; + } Expr expr = bound_correlated_differences(simplify(remove_likelies(e))); + debug(4) << " bcd(simplify()): " << expr << "\n"; Interval interval = bounds_of_expr_in_scope(expr, scope, FuncValueBounds(), true); interval = simplify(interval); + debug(4) << " interval: " << interval << "\n"; // Note that we can get non-const but well-defined results (e.g. signed_integer_overflow); // for our purposes here, treat anything non-const as no-bound. diff --git a/src/BoundsTracker.cpp b/src/BoundsTracker.cpp index 15a97e0129c3..bfccac9c39f3 100644 --- a/src/BoundsTracker.cpp +++ b/src/BoundsTracker.cpp @@ -101,6 +101,7 @@ Expr BoundsTracker::simplify_with_context(const Expr &e) const { } wrapped = remove_likelies(wrapped); wrapped = substitute_in_all_lets(wrapped); + debug(4) << "Simplify with context: " << wrapped << "\n"; // Deliberately pass an empty bounds scope here, not `scope`: mixing a // bounds scope with equality facts can make the simplifier represent a // variable by its (wide) interval instead of substituting the exact diff --git a/src/Simplify_Sub.cpp b/src/Simplify_Sub.cpp index cb3fc09a0a13..e09543e4e821 100644 --- a/src/Simplify_Sub.cpp +++ b/src/Simplify_Sub.cpp @@ -425,6 +425,9 @@ Expr Simplify::visit(const Sub *op, ExprInfo *info) { rewrite(x / c0 - (x - y) / c0, ((y + fold(c0 - 1)) - (x % c0)) / c0, c0 > 0) || rewrite((x - y) / c0 - x / c0, ((x % c0) - y) / c0, c0 > 0) || + rewrite(y - ((((y - x) + z) / c1) * c1 + x), (((y - x) + z) % c1) - z, c1 > 0) || + rewrite(y - (((z + (y - x)) / c1) * c1 + x), ((z + (y - x)) % c1) - z, c1 > 0) || + // Simplification of bounds code for various tail // strategies requires cancellations of the form: // min(f(x), y) - g(x) From 514936e98ac4864f8d6090949fab0ca66f5b6a00 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Tue, 25 Aug 2026 10:08:54 +0200 Subject: [PATCH 27/36] 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 c0509796813f..ccb2230844e1 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 990ef0b03ce1615ad4f0e56bd286cfc19a8882f7 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 22:43:49 +0200 Subject: [PATCH 28/36] Let can_prove predicates use the simplifier's known facts The condition of a can_prove predicate in a rewrite rule was simplified on its own, without any of the facts the simplifier has learned on the way down the IR. Substitute those facts into the condition first, and store facts in the same comparison direction the simplifier produces, so that a fact stated as x > y is usable when it visits y < x. This makes fact-driven rewrite rules possible: max/min now pick a side when the facts order the operands, and a division can cancel a multiplication inside a max or min. Co-authored-by: Claude --- src/IRMatch.h | 3 +++ src/Simplify.cpp | 27 ++++++++++++++++++++ src/Simplify_Div.cpp | 7 +++++ src/Simplify_Internal.h | 12 +++++++++ src/Simplify_Max.cpp | 4 +++ src/Simplify_Min.cpp | 4 +++ test/correctness/simplify.cpp | 48 +++++++++++++++++++++++++++++++++++ 7 files changed, 105 insertions(+) diff --git a/src/IRMatch.h b/src/IRMatch.h index 6fa4cadc4eae..49f575f878be 100644 --- a/src/IRMatch.h +++ b/src/IRMatch.h @@ -2554,6 +2554,9 @@ struct CanProve { // Includes a raw call to an inlined make method, so don't inline. [[nodiscard]] HALIDE_NEVER_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const { Expr condition = a.make(state, {}); + // Inject anything the prover currently knows to be true or false into + // the condition before trying to simplify it. + condition = prover->substitute_facts(condition); condition = prover->mutate(condition, nullptr); val.u.u64 = is_const_one(condition); ty = Bool(condition.type().lanes()); diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 75b908ce8202..c28d656656de 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -85,6 +85,16 @@ void Simplify::found_buffer_reference(const string &name, size_t dimensions) { } void Simplify::ScopedFact::learn_false(const Expr &fact) { + // Canonicalize the direction of comparisons, so that facts are stored in + // the same form the simplifier produces when it visits them. + if (const GT *gt = fact.as()) { + learn_false(gt->b < gt->a); + return; + } else if (const GE *ge = fact.as()) { + learn_false(!(ge->a < ge->b)); + return; + } + Simplify::VarInfo info; info.old_uses = info.new_uses = 0; if (const Variable *v = fact.as()) { @@ -172,6 +182,16 @@ void Simplify::ScopedFact::learn_lower_bound(const Variable *v, int64_t val) { } void Simplify::ScopedFact::learn_true(const Expr &fact) { + // Canonicalize the direction of comparisons, so that facts are stored in + // the same form the simplifier produces when it visits them. + if (const GT *gt = fact.as()) { + learn_true(gt->b < gt->a); + return; + } else if (const GE *ge = fact.as()) { + learn_true(!(ge->a < ge->b)); + return; + } + Simplify::VarInfo info; info.old_uses = info.new_uses = 0; if (const Variable *v = fact.as()) { @@ -370,6 +390,13 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) { return substitute_facts_impl(s, truths, falsehoods); } +Expr Simplify::substitute_facts(const Expr &e) { + if (truths.empty() && falsehoods.empty()) { + return e; + } + return substitute_facts_impl(e, truths, falsehoods); +} + Simplify::ScopedFact::~ScopedFact() { for (const auto *v : pop_list) { simplify->var_info.pop(v->name); diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..5d9734b97faf 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -85,6 +85,13 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { (!op->type.is_float() && rewrite(x / x, select(x == 0, 0, 1))) || (no_overflow(op->type) && + // Facts learned higher up in the IR may tell us which side of a max + // or min survives the division. + (has_facts() && + (rewrite(max(x * c0, y) / c0, x, c0 > 0 && can_prove(y / c0 <= x, this)) || + rewrite(max(y, x * c0) / c0, x, c0 > 0 && can_prove(y / c0 <= x, this)) || + rewrite(min(x * c0, y) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)) || + rewrite(min(y, x * c0) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)))) || // Fold repeated division (rewrite((x / c0) / c2, x / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2)) || rewrite((x / c0 + c1) / c2, (x + fold(c1 * c0)) / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2) && !overflows(c0 * c1)) || diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 94a50bebc644..7841bc8fbe32 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -441,6 +441,18 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; + // Is there anything in the truths/falsehoods sets? Used to gate rewrite + // rules whose predicates are only ever provable from facts learned higher + // up in the IR, so that we don't pay for them in the common case. + bool has_facts() const { + return !truths.empty() || !falsehoods.empty(); + } + + // Replace exprs known to be truths or falsehoods with const_true or + // const_false. Used to inject everything currently known into the + // conditions of can_prove predicates in rewrite rules. + Expr substitute_facts(const Expr &e); + struct ScopedFact { Simplify *simplify; diff --git a/src/Simplify_Max.cpp b/src/Simplify_Max.cpp index 88d3ce2cbf5e..5c10bcde17b4 100644 --- a/src/Simplify_Max.cpp +++ b/src/Simplify_Max.cpp @@ -71,6 +71,10 @@ Expr Simplify::visit(const Max *op, ExprInfo *info) { // RHS for ExprInfo to update correctly. if (EVAL_IN_LAMBDA // (rewrite(max(x, x), a) || + // Facts learned higher up in the IR may tell us which side wins. + (has_facts() && + (rewrite(max(x, y), a, can_prove(y < x, this)) || + rewrite(max(x, y), b, can_prove(x < y, this)))) || rewrite(max(x, c0), b, is_max_value(c0)) || rewrite(max(x, c0), a, is_min_value(c0)) || rewrite(max((x / c0) * c0, x), b, c0 > 0) || diff --git a/src/Simplify_Min.cpp b/src/Simplify_Min.cpp index 5203a0c14166..55d7cac5cf16 100644 --- a/src/Simplify_Min.cpp +++ b/src/Simplify_Min.cpp @@ -70,6 +70,10 @@ Expr Simplify::visit(const Min *op, ExprInfo *info) { // RHS for ExprInfo to update correctly. if (EVAL_IN_LAMBDA // (rewrite(min(x, x), a) || + // Facts learned higher up in the IR may tell us which side wins. + (has_facts() && + (rewrite(min(x, y), a, can_prove(x < y, this)) || + rewrite(min(x, y), b, can_prove(y < x, this)))) || rewrite(min(x, c0), b, is_min_value(c0)) || rewrite(min(x, c0), a, is_max_value(c0)) || rewrite(min((x / c0) * c0, x), a, c0 > 0) || diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 981a00e6f0ee..a60bcb9a422e 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2377,6 +2377,18 @@ void check_invariant() { } } +void check_with_assumptions(const Expr &a, const Expr &b, const std::vector &assumptions) { + Expr simpler = simplify(a, Scope(), Scope(), assumptions); + if (!equal(simpler, b)) { + std::cerr + << "\nSimplification failure:\n" + << "Input: " << a << "\n" + << "Output: " << simpler << "\n" + << "Expected output: " << b << "\n"; + abort(); + } +} + void check_unreachable() { Var x("x"), y("y"); @@ -2405,6 +2417,41 @@ void check_unreachable() { Evaluate::make(0)); } +void check_facts() { + Expr x = Var("x"), y = Var("y"), z = Var("z"); + + // A fact stated in any comparison direction should let the simplifier pick + // the winning side of a max or min. + check_with_assumptions(max(x, y), x, {x > y}); + check_with_assumptions(max(x, y), x, {y < x}); + check_with_assumptions(max(x, y), y, {x < y}); + check_with_assumptions(max(x, y), y, {y > x}); + check_with_assumptions(min(x, y), y, {x > y}); + check_with_assumptions(min(x, y), x, {x < y}); + + // Facts about compound expressions work too. + check_with_assumptions(max(x + z, y * 3), x + z, {x + z > y * 3}); + check_with_assumptions(max(max(x, y), z), z, {max(x, y) < z}); + + // A fact only applies where it holds. + check(IfThenElse::make(x < y, not_no_op(max(x, y)), not_no_op(max(x, y))), + IfThenElse::make(x < y, not_no_op(y), not_no_op(max(x, y)))); + + // A division can cancel a multiplication inside a max or min when we know + // which side wins after the division. + check_with_assumptions(max(x * 8, y) / 8, x, {x >= y / 8}); + check_with_assumptions(max(y, x * 8) / 8, x, {x >= y / 8}); + check_with_assumptions(min(x * 8, y) / 8, x, {x <= y / 8}); + check_with_assumptions(min(y, x * 8) / 8, x, {x <= y / 8}); + + // Without the fact, the division stays put. + check(max(x * 8, y) / 8, max(x * 8, y) / 8); + + // Facts that don't strictly order the operands don't fire these rules. + check_with_assumptions(max(x, y), max(x, y), {x != y}); + check_with_assumptions(max(x * 8, y) / 8, max(x * 8, y) / 8, {x < y / 8}); +} + int main(int argc, char **argv) { check_invariant(); check_casts(); @@ -2417,6 +2464,7 @@ int main(int argc, char **argv) { check_bitwise(); check_lets(); check_unreachable(); + check_facts(); // Miscellaneous cases that don't fit into one of the categories above. Expr x = Var("x"), y = Var("y"); From 482af79f436447e1249824542332400d58c763c8 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 22:55:59 +0200 Subject: [PATCH 29/36] Make fact lookup aware of comparison direction and strictness Facts and the conditions of can_prove predicates are now looked up in the same canonical form: GT and GE are mapped onto LT, Not is unwrapped, and a comparison can be settled by the other strictness of the same comparison in either direction. This means it no longer matters how a fact was spelled relative to how the rule that consumes it was, and a strict fact such as x > y settles the non-strict predicate the max/min rules ask for. Those rules ask non-strictly, since a tie makes either side of a max or min an equally good answer, so a fact of x >= y is enough to pick a side. Co-authored-by: Claude --- src/Simplify.cpp | 48 ++++++++++++++++++++++++++++++++--- src/Simplify_Div.cpp | 4 +-- src/Simplify_Max.cpp | 4 +-- src/Simplify_Min.cpp | 4 +-- test/correctness/simplify.cpp | 28 ++++++++++++++++++-- 5 files changed, 76 insertions(+), 12 deletions(-) diff --git a/src/Simplify.cpp b/src/Simplify.cpp index c28d656656de..e524f1d9900d 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -365,16 +365,56 @@ void Simplify::ScopedFact::learn_true(const Expr &fact) { } namespace { +// Is a boolean Expr known to be true or false? Facts are stored in the same +// form the simplifier itself produces, so a comparison has to be canonicalized +// the same way before looking it up. +std::optional lookup_fact(const Expr &e, + const std::set &truths, + const std::set &falsehoods) { + if (const Not *n = e.as()) { + auto known = lookup_fact(n->a, truths, falsehoods); + return known ? std::make_optional(!*known) : known; + } else if (const GT *gt = e.as()) { + return lookup_fact(gt->b < gt->a, truths, falsehoods); + } else if (const GE *ge = e.as()) { + return lookup_fact(!(ge->a < ge->b), truths, falsehoods); + } + + if (truths.count(e)) { + return true; + } else if (falsehoods.count(e)) { + return false; + } + + // A comparison may also be settled by the other strictness of the same + // comparison, in either direction. + if (const LT *lt = e.as()) { + // a < b is implied by !(b <= a), and ruled out by b <= a and by b < a. + if (falsehoods.count(lt->b <= lt->a)) { + return true; + } else if (truths.count(lt->b <= lt->a) || truths.count(lt->b < lt->a)) { + return false; + } + } else if (const LE *le = e.as()) { + // a <= b is implied by a < b and by !(b < a), and ruled out by b < a. + if (truths.count(le->a < le->b) || falsehoods.count(le->b < le->a)) { + return true; + } else if (truths.count(le->b < le->a)) { + return false; + } + } + + return std::nullopt; +} + template T substitute_facts_impl(const T &t, const std::set &truths, const std::set &falsehoods) { return mutate_with(t, [&](auto *self, const Expr &e) { if (e.type().is_bool()) { - if (truths.count(e)) { - return make_one(e.type()); - } else if (falsehoods.count(e)) { - return make_zero(e.type()); + if (auto known = lookup_fact(e, truths, falsehoods)) { + return *known ? make_one(e.type()) : make_zero(e.type()); } } return self->mutate_base(e); diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 5d9734b97faf..4934e961d385 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -88,8 +88,8 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { // Facts learned higher up in the IR may tell us which side of a max // or min survives the division. (has_facts() && - (rewrite(max(x * c0, y) / c0, x, c0 > 0 && can_prove(y / c0 <= x, this)) || - rewrite(max(y, x * c0) / c0, x, c0 > 0 && can_prove(y / c0 <= x, this)) || + (rewrite(max(x * c0, y) / c0, x, c0 > 0 && can_prove(x >= y / c0, this)) || + rewrite(max(y, x * c0) / c0, x, c0 > 0 && can_prove(x >= y / c0, this)) || rewrite(min(x * c0, y) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)) || rewrite(min(y, x * c0) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)))) || // Fold repeated division diff --git a/src/Simplify_Max.cpp b/src/Simplify_Max.cpp index 5c10bcde17b4..cae81d969585 100644 --- a/src/Simplify_Max.cpp +++ b/src/Simplify_Max.cpp @@ -73,8 +73,8 @@ Expr Simplify::visit(const Max *op, ExprInfo *info) { (rewrite(max(x, x), a) || // Facts learned higher up in the IR may tell us which side wins. (has_facts() && - (rewrite(max(x, y), a, can_prove(y < x, this)) || - rewrite(max(x, y), b, can_prove(x < y, this)))) || + (rewrite(max(x, y), a, can_prove(y <= x, this)) || + rewrite(max(x, y), b, can_prove(x <= y, this)))) || rewrite(max(x, c0), b, is_max_value(c0)) || rewrite(max(x, c0), a, is_min_value(c0)) || rewrite(max((x / c0) * c0, x), b, c0 > 0) || diff --git a/src/Simplify_Min.cpp b/src/Simplify_Min.cpp index 55d7cac5cf16..3444c1dd1509 100644 --- a/src/Simplify_Min.cpp +++ b/src/Simplify_Min.cpp @@ -72,8 +72,8 @@ Expr Simplify::visit(const Min *op, ExprInfo *info) { (rewrite(min(x, x), a) || // Facts learned higher up in the IR may tell us which side wins. (has_facts() && - (rewrite(min(x, y), a, can_prove(x < y, this)) || - rewrite(min(x, y), b, can_prove(y < x, this)))) || + (rewrite(min(x, y), a, can_prove(x <= y, this)) || + rewrite(min(x, y), b, can_prove(y <= x, this)))) || rewrite(min(x, c0), b, is_min_value(c0)) || rewrite(min(x, c0), a, is_max_value(c0)) || rewrite(min((x / c0) * c0, x), a, c0 > 0) || diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index a60bcb9a422e..1eae900c8eb4 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2429,13 +2429,26 @@ void check_facts() { check_with_assumptions(min(x, y), y, {x > y}); check_with_assumptions(min(x, y), x, {x < y}); + // A non-strict fact is enough to pick a side of a max or min, and a strict + // fact implies the non-strict one. + check_with_assumptions(max(x, y), x, {x >= y}); + check_with_assumptions(max(x, y), y, {x <= y}); + check_with_assumptions(min(x, y), x, {x <= y}); + check_with_assumptions(min(x, y), y, {x >= y}); + // Facts about compound expressions work too. check_with_assumptions(max(x + z, y * 3), x + z, {x + z > y * 3}); check_with_assumptions(max(max(x, y), z), z, {max(x, y) < z}); - // A fact only applies where it holds. + // Both branches of an if learn from the condition, in opposite directions. check(IfThenElse::make(x < y, not_no_op(max(x, y)), not_no_op(max(x, y))), - IfThenElse::make(x < y, not_no_op(y), not_no_op(max(x, y)))); + IfThenElse::make(x < y, not_no_op(y), not_no_op(x))); + + // A fact only applies where it holds. + check(Block::make(not_no_op(max(x, y)), + IfThenElse::make(x < y, not_no_op(max(x, y)))), + Block::make(not_no_op(max(x, y)), + IfThenElse::make(x < y, not_no_op(y)))); // A division can cancel a multiplication inside a max or min when we know // which side wins after the division. @@ -2444,6 +2457,17 @@ void check_facts() { check_with_assumptions(min(x * 8, y) / 8, x, {x <= y / 8}); check_with_assumptions(min(y, x * 8) / 8, x, {x <= y / 8}); + // The direction in which a fact is stated doesn't matter, on either side: + // both the facts and the conditions of can_prove predicates are looked up + // in the same canonical form. + check_with_assumptions(max(x * 8, y) / 8, x, {y / 8 <= x}); + check_with_assumptions(max(x * 8, y) / 8, x, {!(x < y / 8)}); + check_with_assumptions(min(x * 8, y) / 8, x, {y / 8 >= x}); + + // A strict fact settles a non-strict predicate too. + check_with_assumptions(max(x * 8, y) / 8, x, {x > y / 8}); + check_with_assumptions(min(x * 8, y) / 8, x, {x < y / 8}); + // Without the fact, the division stays put. check(max(x * 8, y) / 8, max(x * 8, y) / 8); From caba268a4d29a930a8c05cee060995735af4f7f0 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 23:07:41 +0200 Subject: [PATCH 30/36] Don't re-enter fact-driven rewrite rules from inside a can_prove Simplifying the condition of a can_prove predicate visits the operands again, so a fact-driven rule that matches every node of its type recursed without bound on nested min/max trees. Disable those rules while inside a can_prove condition; the facts themselves are still substituted in at every level. Co-authored-by: Claude --- src/IRMatch.h | 5 +---- src/Simplify.cpp | 5 +++++ src/Simplify_Internal.h | 19 +++++++++++++++---- test/correctness/simplify.cpp | 9 +++++++++ 4 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/IRMatch.h b/src/IRMatch.h index 49f575f878be..16fdde3aa4a9 100644 --- a/src/IRMatch.h +++ b/src/IRMatch.h @@ -2554,10 +2554,7 @@ struct CanProve { // Includes a raw call to an inlined make method, so don't inline. [[nodiscard]] HALIDE_NEVER_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const { Expr condition = a.make(state, {}); - // Inject anything the prover currently knows to be true or false into - // the condition before trying to simplify it. - condition = prover->substitute_facts(condition); - condition = prover->mutate(condition, nullptr); + condition = prover->simplify_can_prove_condition(condition); val.u.u64 = is_const_one(condition); ty = Bool(condition.type().lanes()); return false; diff --git a/src/Simplify.cpp b/src/Simplify.cpp index e524f1d9900d..86b514b24dd3 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -430,6 +430,11 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) { return substitute_facts_impl(s, truths, falsehoods); } +Expr Simplify::simplify_can_prove_condition(const Expr &e) { + ScopedValue guard(in_can_prove, true); + return mutate(substitute_facts(e), nullptr); +} + Expr Simplify::substitute_facts(const Expr &e) { if (truths.empty() && falsehoods.empty()) { return e; diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 7841bc8fbe32..89d15d6ceaae 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -441,11 +441,18 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; - // Is there anything in the truths/falsehoods sets? Used to gate rewrite - // rules whose predicates are only ever provable from facts learned higher - // up in the IR, so that we don't pay for them in the common case. + // Are we already inside the simplification of the condition of a can_prove + // predicate? Fact-driven rules are disabled in there, because simplifying + // such a condition visits the operands again, and a rule that fires on + // every node of its type would recurse without bound on nested min/max. + bool in_can_prove = false; + + // Is there anything in the truths/falsehoods sets that a rewrite rule could + // use? Used to gate rules whose predicates are only ever provable from facts + // learned higher up in the IR, so that we don't pay for them in the common + // case. bool has_facts() const { - return !truths.empty() || !falsehoods.empty(); + return !in_can_prove && (!truths.empty() || !falsehoods.empty()); } // Replace exprs known to be truths or falsehoods with const_true or @@ -453,6 +460,10 @@ class Simplify : public VariadicVisitor { // conditions of can_prove predicates in rewrite rules. Expr substitute_facts(const Expr &e); + // Simplify the condition of a can_prove predicate in a rewrite rule, using + // everything currently known. + Expr simplify_can_prove_condition(const Expr &e); + struct ScopedFact { Simplify *simplify; diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 1eae900c8eb4..90b51ae5c56f 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2468,6 +2468,15 @@ void check_facts() { check_with_assumptions(max(x * 8, y) / 8, x, {x > y / 8}); check_with_assumptions(min(x * 8, y) / 8, x, {x < y / 8}); + // Deeply nested mins and maxes must not make the work of proving the + // predicates of the rules above blow up. + Expr nest = x; + for (int i = 0; i < 24; i++) { + nest = min(max(nest + i, y - i), z * i); + } + // The result isn't interesting; what matters is that we get one at all. + (void)simplify(nest, Scope(), Scope(), {x < y}); + // Without the fact, the division stays put. check(max(x * 8, y) / 8, max(x * 8, y) / 8); From 9ba560f47d86f904a8f86360fffa6d61daae406e Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Wed, 26 Aug 2026 23:13:46 +0200 Subject: [PATCH 31/36] Express the can_prove re-entry guard as a depth limit Recursing further is occasionally useful in principle, but measurably expensive: at a limit of 2, correctness_likely goes from 1.0s to 4.2s and correctness_autodiff from 3.4s to 11.4s, with no test producing a better simplification. Keep the limit at one level, but name the constant. Co-authored-by: Claude --- src/Simplify.cpp | 2 +- src/Simplify_Internal.h | 14 ++++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 86b514b24dd3..bbf9260580f2 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -431,7 +431,7 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) { } Expr Simplify::simplify_can_prove_condition(const Expr &e) { - ScopedValue guard(in_can_prove, true); + ScopedValue guard(can_prove_depth, can_prove_depth + 1); return mutate(substitute_facts(e), nullptr); } diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 89d15d6ceaae..4e178e9009f9 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -441,18 +441,20 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; - // Are we already inside the simplification of the condition of a can_prove - // predicate? Fact-driven rules are disabled in there, because simplifying - // such a condition visits the operands again, and a rule that fires on - // every node of its type would recurse without bound on nested min/max. - bool in_can_prove = false; + // How deeply are we nested inside the conditions of can_prove predicates? + // Simplifying such a condition visits the operands again, so a fact-driven + // rule that matches every node of its type recurses, and the work grows + // like the nesting depth of the expression raised to this. Bound it. + int can_prove_depth = 0; + static constexpr int max_can_prove_depth = 1; // Is there anything in the truths/falsehoods sets that a rewrite rule could // use? Used to gate rules whose predicates are only ever provable from facts // learned higher up in the IR, so that we don't pay for them in the common // case. bool has_facts() const { - return !in_can_prove && (!truths.empty() || !falsehoods.empty()); + return can_prove_depth < max_can_prove_depth && + (!truths.empty() || !falsehoods.empty()); } // Replace exprs known to be truths or falsehoods with const_true or From 55827b5e3f7cc58a510febe914cfd4b04c7baea7 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 27 Aug 2026 09:50:12 +0200 Subject: [PATCH 32/36] Add a non-recursive known_true predicate for rewrite rules can_prove as a rewrite predicate recursively invokes the simplifier on every expression matching the rule's left-hand side, so a rule whose left-hand side also matches something built while proving the predicate recurses. It is also simply expensive. known_true instead looks the condition up in the facts directly. It cannot recurse, and it is cheap enough to use on a rule that matches every node of its type. The fact-driven max, min and division rules now use it, which is enough for all of them: looking up a comparison already understands direction and strictness. Co-authored-by: Claude --- src/IRMatch.h | 40 +++++++++++++++++++++++++++++++++++ src/Simplify.cpp | 8 +++++++ src/Simplify_Div.cpp | 8 +++---- src/Simplify_Internal.h | 4 ++++ src/Simplify_Max.cpp | 4 ++-- src/Simplify_Min.cpp | 4 ++-- test/correctness/simplify.cpp | 5 +++++ 7 files changed, 65 insertions(+), 8 deletions(-) diff --git a/src/IRMatch.h b/src/IRMatch.h index 16fdde3aa4a9..b62d4892d74c 100644 --- a/src/IRMatch.h +++ b/src/IRMatch.h @@ -2573,6 +2573,46 @@ std::ostream &operator<<(std::ostream &s, const CanProve &op) { return s; } +// Like can_prove, but only looks the condition up in the facts the prover +// already knows, instead of recursively invoking it. Much cheaper, and it +// cannot recurse, so unlike can_prove it is safe in a rule whose left-hand +// side matches expressions the prover may construct while proving it. +template +struct KnownTrue { + struct pattern_tag {}; + A a; + Prover *prover; // An existing simplifying mutator + + constexpr static uint32_t binds = bindings::mask; + + // This rule is a boolean-valued predicate. Bools have type UIntImm. + constexpr static IRNodeType min_node_type = IRNodeType::UIntImm; + constexpr static IRNodeType max_node_type = IRNodeType::UIntImm; + constexpr static bool canonical = true; + + constexpr static bool foldable = true; + + // Includes a raw call to an inlined make method, so don't inline. + [[nodiscard]] HALIDE_NEVER_INLINE bool make_folded_const(halide_scalar_value_t &val, Type &ty, MatcherState &state) const { + Expr condition = a.make(state, {}); + val.u.u64 = prover->is_known_true(condition) ? 1 : 0; + ty = Bool(condition.type().lanes()); + return false; + } +}; + +template +HALIDE_ALWAYS_INLINE auto known_true(A &&a, Prover *p) noexcept -> KnownTrue { + assert_is_lvalue_if_expr(); + return {pattern_arg(a), p}; +} + +template +std::ostream &operator<<(std::ostream &s, const KnownTrue &op) { + s << "known_true(" << op.a << ")"; + return s; +} + template struct IsFloat { struct pattern_tag {}; diff --git a/src/Simplify.cpp b/src/Simplify.cpp index bbf9260580f2..7fd4990a6c03 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -430,6 +430,14 @@ Stmt Simplify::ScopedFact::substitute_facts(const Stmt &s) { return substitute_facts_impl(s, truths, falsehoods); } +bool Simplify::is_known_true(const Expr &e) { + if (truths.empty() && falsehoods.empty()) { + return false; + } + auto known = lookup_fact(e, truths, falsehoods); + return known && *known; +} + Expr Simplify::simplify_can_prove_condition(const Expr &e) { ScopedValue guard(can_prove_depth, can_prove_depth + 1); return mutate(substitute_facts(e), nullptr); diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4934e961d385..1a6e1eb51470 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -88,10 +88,10 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { // Facts learned higher up in the IR may tell us which side of a max // or min survives the division. (has_facts() && - (rewrite(max(x * c0, y) / c0, x, c0 > 0 && can_prove(x >= y / c0, this)) || - rewrite(max(y, x * c0) / c0, x, c0 > 0 && can_prove(x >= y / c0, this)) || - rewrite(min(x * c0, y) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)) || - rewrite(min(y, x * c0) / c0, x, c0 > 0 && can_prove(x <= y / c0, this)))) || + (rewrite(max(x * c0, y) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || + rewrite(max(y, x * c0) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || + rewrite(min(x * c0, y) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) || + rewrite(min(y, x * c0) / c0, x, c0 > 0 && known_true(x <= y / c0, this)))) || // Fold repeated division (rewrite((x / c0) / c2, x / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2)) || rewrite((x / c0 + c1) / c2, (x + fold(c1 * c0)) / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2) && !overflows(c0 * c1)) || diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 4e178e9009f9..1eb5d3fd95b8 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -466,6 +466,10 @@ class Simplify : public VariadicVisitor { // everything currently known. Expr simplify_can_prove_condition(const Expr &e); + // Is a boolean Expr already known to be true? Unlike can_prove this only + // looks the condition up in the facts, without simplifying anything. + bool is_known_true(const Expr &e); + struct ScopedFact { Simplify *simplify; diff --git a/src/Simplify_Max.cpp b/src/Simplify_Max.cpp index cae81d969585..c1c161d6cdc8 100644 --- a/src/Simplify_Max.cpp +++ b/src/Simplify_Max.cpp @@ -73,8 +73,8 @@ Expr Simplify::visit(const Max *op, ExprInfo *info) { (rewrite(max(x, x), a) || // Facts learned higher up in the IR may tell us which side wins. (has_facts() && - (rewrite(max(x, y), a, can_prove(y <= x, this)) || - rewrite(max(x, y), b, can_prove(x <= y, this)))) || + (rewrite(max(x, y), a, known_true(y <= x, this)) || + rewrite(max(x, y), b, known_true(x <= y, this)))) || rewrite(max(x, c0), b, is_max_value(c0)) || rewrite(max(x, c0), a, is_min_value(c0)) || rewrite(max((x / c0) * c0, x), b, c0 > 0) || diff --git a/src/Simplify_Min.cpp b/src/Simplify_Min.cpp index 3444c1dd1509..880f2a4b890d 100644 --- a/src/Simplify_Min.cpp +++ b/src/Simplify_Min.cpp @@ -72,8 +72,8 @@ Expr Simplify::visit(const Min *op, ExprInfo *info) { (rewrite(min(x, x), a) || // Facts learned higher up in the IR may tell us which side wins. (has_facts() && - (rewrite(min(x, y), a, can_prove(x <= y, this)) || - rewrite(min(x, y), b, can_prove(y <= x, this)))) || + (rewrite(min(x, y), a, known_true(x <= y, this)) || + rewrite(min(x, y), b, known_true(y <= x, this)))) || rewrite(min(x, c0), b, is_min_value(c0)) || rewrite(min(x, c0), a, is_max_value(c0)) || rewrite(min((x / c0) * c0, x), a, c0 > 0) || diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 90b51ae5c56f..ef6a107ae783 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2477,6 +2477,11 @@ void check_facts() { // The result isn't interesting; what matters is that we get one at all. (void)simplify(nest, Scope(), Scope(), {x < y}); + // The rules above look their predicates up in the facts rather than + // recursively invoking the simplifier, so a fact only settles a predicate + // it is directly comparable to. This one needs arithmetic to connect: + check_with_assumptions(max(x, y), max(x, y), {x + 1 <= y}); + // Without the fact, the division stays put. check(max(x * 8, y) / 8, max(x * 8, y) / 8); From a983f5af0a447a252dc2e2ad3c706938bc311590 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 27 Aug 2026 15:49:21 +0200 Subject: [PATCH 33/36] Fix parenthesis of Simplify_Div. --- src/Simplify_Div.cpp | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/Simplify_Div.cpp b/src/Simplify_Div.cpp index 1a6e1eb51470..d7839340fb99 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -84,14 +84,19 @@ Expr Simplify::visit(const Div *op, ExprInfo *info) { rewrite(select(x, c0, c1) / c2, select(x, fold(c0 / c2), fold(c1 / c2))) || (!op->type.is_float() && rewrite(x / x, select(x == 0, 0, 1))) || + + (no_overflow(op->type) && + // Facts learned higher up in the IR may tell us which side of a max + // or min survives the division. Test them early on to prevents rewrites below + // that would make it impossible to recognize the form. + (has_facts() && + (rewrite(max(x * c0, y) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || + rewrite(max(y, x * c0) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || + rewrite(min(x * c0, y) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) || + rewrite(min(y, x * c0) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) || + false))) || + (no_overflow(op->type) && - // Facts learned higher up in the IR may tell us which side of a max - // or min survives the division. - (has_facts() && - (rewrite(max(x * c0, y) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || - rewrite(max(y, x * c0) / c0, x, c0 > 0 && known_true(x >= y / c0, this)) || - rewrite(min(x * c0, y) / c0, x, c0 > 0 && known_true(x <= y / c0, this)) || - rewrite(min(y, x * c0) / c0, x, c0 > 0 && known_true(x <= y / c0, this)))) || // Fold repeated division (rewrite((x / c0) / c2, x / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2)) || rewrite((x / c0 + c1) / c2, (x + fold(c1 * c0)) / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2) && !overflows(c0 * c1)) || From fa4a807309c0dfaaca34b3165c59409a77e1c34a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 27 Aug 2026 17:42:13 +0200 Subject: [PATCH 34/36] Guard against can_prove recursion at its source The depth limit was checked in has_facts, which only protects rules that consult it. Checking it on entry to the condition simplification instead protects every can_prove, including the pre-existing rules and any future one, and returning the condition unsimplified is the natural way to decline: the predicate simply fails to prove anything. That also frees has_facts to be a plain check, so the non-recursive known_true rules can fire at any depth. The limit is raised to four, which restricts nothing today: instrumenting every correctness test shows the deepest can_prove nesting any of them reaches is one. Co-authored-by: Claude --- src/Simplify.cpp | 5 +++++ src/Simplify_Internal.h | 12 ++++++------ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 7fd4990a6c03..7a48e0a23de8 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -439,6 +439,11 @@ bool Simplify::is_known_true(const Expr &e) { } Expr Simplify::simplify_can_prove_condition(const Expr &e) { + if (can_prove_depth >= max_can_prove_depth) { + // Refuse to nest any deeper. Returning the condition unsimplified just + // means the predicate fails to prove anything. + return e; + } ScopedValue guard(can_prove_depth, can_prove_depth + 1); return mutate(substitute_facts(e), nullptr); } diff --git a/src/Simplify_Internal.h b/src/Simplify_Internal.h index 1eb5d3fd95b8..bf45a9d3977f 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -442,19 +442,19 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; // How deeply are we nested inside the conditions of can_prove predicates? - // Simplifying such a condition visits the operands again, so a fact-driven - // rule that matches every node of its type recurses, and the work grows - // like the nesting depth of the expression raised to this. Bound it. + // Proving such a condition recursively invokes the simplifier on it, so a + // rule whose left-hand side also matches something built while proving its + // own predicate recurses without bound. Nesting is also expensive, and no + // rule currently relies on it. Bound it. int can_prove_depth = 0; - static constexpr int max_can_prove_depth = 1; + static constexpr int max_can_prove_depth = 4; // Is there anything in the truths/falsehoods sets that a rewrite rule could // use? Used to gate rules whose predicates are only ever provable from facts // learned higher up in the IR, so that we don't pay for them in the common // case. bool has_facts() const { - return can_prove_depth < max_can_prove_depth && - (!truths.empty() || !falsehoods.empty()); + return !truths.empty() || !falsehoods.empty(); } // Replace exprs known to be truths or falsehoods with const_true or From 52a5e8733d7ad7007a9569676f7a9e2fc8016c5a Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 27 Aug 2026 18:18:53 +0200 Subject: [PATCH 35/36] Fall back to fact lookup at the can_prove depth cap Refusing to simplify the condition past the depth limit meant the predicate could never be proven there, even when the fact needed was already known. substitute_facts is a plain tree walk (mutate_with over the generic IRMutator base traversal) that never invokes a rewrite rule, so it cannot re-trigger can_prove or known_true and stays safe at any depth: use it as the fallback instead of returning the condition untouched. Added a regression test built on the pre-existing can_prove-based min/max subtraction cancellations in Simplify_Sub.cpp (the rules that motivated the depth limit in the first place, since their predicate constructs a fresh subtraction that can itself match the same rule). With the limit disabled it hangs (confirmed: 15s timeout); with it in place it completes in under a second. Co-authored-by: Claude --- src/Simplify.cpp | 8 +++++--- test/correctness/simplify.cpp | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 7a48e0a23de8..10e5dd035052 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -440,9 +440,11 @@ bool Simplify::is_known_true(const Expr &e) { Expr Simplify::simplify_can_prove_condition(const Expr &e) { if (can_prove_depth >= max_can_prove_depth) { - // Refuse to nest any deeper. Returning the condition unsimplified just - // means the predicate fails to prove anything. - return e; + // Too deep to safely recurse into the full simplifier. substitute_facts + // is a plain tree walk that never invokes a rewrite rule (it can't + // re-trigger can_prove or known_true), so it remains safe and cheap + // here: fall back to it rather than giving up on the condition. + return substitute_facts(e); } ScopedValue guard(can_prove_depth, can_prove_depth + 1); return mutate(substitute_facts(e), nullptr); diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index ef6a107ae783..63a2b9de3fec 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2477,6 +2477,22 @@ void check_facts() { // The result isn't interesting; what matters is that we get one at all. (void)simplify(nest, Scope(), Scope(), {x < y}); + // can_prove-based rules (unlike the known_true ones above) recursively + // invoke the simplifier on their own predicate, and that predicate can be + // a freshly built expression rather than a piece of the original IR (e.g. + // min(x, y) - min(z, w) -> y - w, can_prove(x - y == z - w)) constructs a + // brand new subtraction). If the operands are themselves unsimplified + // instances of the same shape, this recurses; the depth limit must bound + // the work rather than let it explode. + Expr deep = min(Var("da"), Var("db")) - min(Var("dc"), Var("dd")); + for (int i = 0; i < 10; i++) { + Expr y = Var("dy" + std::to_string(i)); + Expr z = Var("dz" + std::to_string(i)); + Expr w = Var("dw" + std::to_string(i)); + deep = min(deep, y) - min(z, w); + } + (void)simplify(deep); + // The rules above look their predicates up in the facts rather than // recursively invoking the simplifier, so a fact only settles a predicate // it is directly comparable to. This one needs arithmetic to connect: From 032f724da9d7d8ae94e930a9a07759f6489dcbb2 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Thu, 27 Aug 2026 19:00:58 +0200 Subject: [PATCH 36/36] Use a direct fact lookup at the can_prove depth cap, not a tree walk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous fallback ran substitute_facts, a full tree walk, on the condition. But the only thing the caller checks is whether the result is literally the constant true, and nothing runs afterward to fold a compound expression: an And of two individually-known-true operands stays an unfolded And, never becoming true. So substitute_facts's ability to resolve facts about pieces of a compound condition was wasted work here — it can't prove anything is_known_true on the condition itself couldn't already, since folding that partial progress into a verdict is exactly the recursive work the cap exists to avoid. Co-authored-by: Claude --- src/Simplify.cpp | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/Simplify.cpp b/src/Simplify.cpp index 10e5dd035052..2ac9081dd4d1 100644 --- a/src/Simplify.cpp +++ b/src/Simplify.cpp @@ -440,11 +440,17 @@ bool Simplify::is_known_true(const Expr &e) { Expr Simplify::simplify_can_prove_condition(const Expr &e) { if (can_prove_depth >= max_can_prove_depth) { - // Too deep to safely recurse into the full simplifier. substitute_facts - // is a plain tree walk that never invokes a rewrite rule (it can't - // re-trigger can_prove or known_true), so it remains safe and cheap - // here: fall back to it rather than giving up on the condition. - return substitute_facts(e); + // Too deep to safely recurse into the full simplifier. The only thing + // the caller does with the result is check whether it is the literal + // constant true, and nothing here can fold a compound expression (an + // And of two known-true operands stays an unfolded And, not true) -- + // that folding is exactly the recursive work we're declining to do. + // So a substitute_facts tree walk can't prove anything a direct + // lookup of the condition itself couldn't already: skip the walk. + if (is_known_true(e)) { + return const_true(e.type().lanes(), nullptr); + } + return e; } ScopedValue guard(can_prove_depth, can_prove_depth + 1); return mutate(substitute_facts(e), nullptr);