diff --git a/Makefile b/Makefile index c66ac327781f..b69a101ab510 100644 --- a/Makefile +++ b/Makefile @@ -460,6 +460,7 @@ SOURCE_FILES = \ Bounds.cpp \ BoundsInference.cpp \ BoundSmallAllocations.cpp \ + BoundsTracker.cpp \ Buffer.cpp \ Callable.cpp \ CanonicalizeGPUVars.cpp \ @@ -665,6 +666,7 @@ HEADER_FILES = \ Bounds.h \ BoundsInference.h \ BoundSmallAllocations.h \ + BoundsTracker.h \ Buffer.h \ Callable.h \ CanonicalizeGPUVars.h \ diff --git a/src/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index ebc41007e6bf..f893e63f9eaa 100644 --- a/src/BoundConstantExtentLoops.cpp +++ b/src/BoundConstantExtentLoops.cpp @@ -1,11 +1,9 @@ #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 +13,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; @@ -47,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 @@ -56,36 +49,46 @@ 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(); - 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 (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(); + } } } + 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..dfeda2fd0b38 --- /dev/null +++ b/src/BoundsTracker.cpp @@ -0,0 +1,198 @@ +#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, + 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_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(); + tracker->facts.pop_back(); + tracker->facts.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); + + // 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}); + Expr loop_var = Variable::make(Int(32), name); + facts.push_back(loop_var >= min); + facts.push_back(loop_var <= 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) { + 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) { + 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. + debug(4) << "Simplify with context: " << wrapped << "\n"; + for (const auto &fact : facts) { + debug(4) << " [fact] " << fact << "\n"; + } + 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; +} + +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); + 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); +} + +Expr BoundsTracker::find_constant_bound_aggressive(const Expr &e, Direction d) const { + Interval interval = find_constant_bounds_aggressive(e); + if (d == Direction::Lower) { + return interval.has_lower_bound() ? interval.min : Expr(); + } else { + return interval.has_upper_bound() ? interval.max : Expr(); + } +} + +} // namespace Internal +} // namespace Halide diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h new file mode 100644 index 000000000000..9a5149c6437a --- /dev/null +++ b/src/BoundsTracker.h @@ -0,0 +1,181 @@ +#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, + 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. 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. + * 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, + * 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. Any endpoint + * still missing after that is attempted once more by substituting the + * range of an enclosing loop e is monotonic in. 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. + * + * The single-Direction form is a wrapper around the Interval form; it + * returns an undefined Expr if no bound in that direction was found. */ + Interval find_constant_bounds_aggressive(const Expr &e) const; + Expr find_constant_bound_aggressive(const Expr &e, Direction d) 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; + + /** 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; + } + + /** 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: + /** 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 +} // 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 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); diff --git a/src/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index dc90aa6d7cef..2fdb85a0ec75 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" @@ -28,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 { @@ -56,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; } @@ -79,19 +92,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 +139,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 +154,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 +183,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 +204,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 +275,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) { 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. 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..2cee901c960f --- /dev/null +++ b/test/correctness/split_predicate_stores_compute_at.cpp @@ -0,0 +1,103 @@ +#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); + } + + // 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 + // 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; +}