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/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")) 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/ApplySplit.cpp b/src/ApplySplit.cpp index ddb9bc1098c5..6df71d634b9e 100644 --- a/src/ApplySplit.cpp +++ b/src/ApplySplit.cpp @@ -23,10 +23,17 @@ vector apply_split(const Split &split, const string &prefix, Expr old_max = Variable::make(Int(32), prefix + split.old_var + ".loop_max"); Expr old_min = Variable::make(Int(32), prefix + split.old_var + ".loop_min"); Expr old_extent = (old_max - old_min) + 1; + Expr outer_min = Variable::make(Int(32), prefix + split.outer + ".loop_min"); dim_extent_alignment[split.inner] = split.factor; - Expr base = outer * split.factor + old_min; + Expr base; + if (split.align.defined()) { + base = outer * split.factor; + } else { + base = outer * split.factor + old_min; + } + string base_name = prefix + split.inner + ".base"; Expr base_var = Variable::make(Int(32), base_name); string old_var_name = prefix + split.old_var; @@ -38,8 +45,17 @@ vector apply_split(const Split &split, const string &prefix, internal_assert(tail != TailStrategy::Auto) << "An explicit tail strategy should exist at this point\n"; + // When align is defined, tiles are anchored to align instead of to + // old_min, so knowing that the factor divides the extent is not + // enough to prove no boundary guard is needed: we additionally need + // the tiling anchored at align to line up with the tiling anchored + // at old_min, i.e. old_min and align must be congruent mod factor. + bool alignment_matches_old_min = !split.align.defined() || + is_const_zero(simplify((old_min - split.align) % split.factor)); + if ((iter != dim_extent_alignment.end()) && - is_const_zero(simplify(iter->second % split.factor))) { + is_const_zero(simplify(iter->second % split.factor)) && + alignment_matches_old_min) { // We have proved that the split factor divides the // old extent. No need to adjust the base or add an if // statement. @@ -58,14 +74,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 +94,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,31 +113,109 @@ vector apply_split(const Split &split, const string &prefix, // for the guarded version. result.emplace_back(prefix + split.old_var, guarded_var, substitution_type); result.emplace_back(guarded_var_name, guarded, ApplySplitResult::LetStmt); - result.emplace_back(likely(old_var <= old_max), predicate_type); + + Expr guard_cond = likely(old_var <= old_max); + if (split.align.defined()) { + guard_cond = likely(old_var >= old_min && old_var <= old_max); + } + result.emplace_back(guard_cond, predicate_type); } else if (tail == TailStrategy::ShiftInwards) { // Adjust the base downwards to not compute off the // end of the realization. - // We'll only mark the base as likely (triggering a loop - // partition) if we're at or inside the innermost - // non-trivial loop. base = likely_if_innermost(base); - base = Min::make(base, old_max + (1 - split.factor)); + if (split.align.defined()) { + base = Max::make(base, old_min - split.align); + base = Min::make(base, old_max + (1 - split.factor) - split.align); + } else { + base = Min::make(base, old_max + (1 - split.factor)); + } } else if (tail == TailStrategy::ShiftInwardsAndBlend) { + // Unclamped base, saved before the Min/Max below adjust it. Used + // to figure out how much (if at all) the boundary tile got + // shifted, so we know which elements of it are redundant with a + // neighboring tile and must be masked out rather than + // recomputed (to avoid double-counting in a reduction). Expr old_base = base; base = likely(base); - base = Min::make(base, old_max + (1 - split.factor)); - // Make a mask which will be a loop invariant if inner gets - // vectorized, and apply it if we're in the tail. - Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner >= unwanted_elems; - mask = select(base == old_base, likely(const_true()), mask); + Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; + Expr mask; + if (split.align.defined()) { + // Because base is anchored to align instead of old_min, the + // 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); + } result.emplace_back(mask, ApplySplitResult::BlendProvides); } else if (tail == TailStrategy::RoundUpAndBlend) { - Expr unwanted_elems = (-old_extent) % split.factor; - Expr mask = inner < split.factor - unwanted_elems; - mask = select(outer < outer_max, likely(const_true()), mask); + Expr zero_based_inner = split.align.defined() ? (inner - split.align) : inner; + 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); @@ -173,12 +267,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/BoundConstantExtentLoops.cpp b/src/BoundConstantExtentLoops.cpp index ebc41007e6bf..29232242b028 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 "IRPrinter.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; @@ -47,6 +40,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 +50,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 << "\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..b59240f14edc 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,8 @@ class BoundSmallAllocations : public IRMutator { for (const Expr &e : op->extents) { total_extent *= e; } - Expr bound = find_constant_bound(total_extent, Direction::Upper, scope); + 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)) { user_assert(op->memory_type != MemoryType::Register) 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 new file mode 100644 index 000000000000..bfccac9c39f3 --- /dev/null +++ b/src/BoundsTracker.cpp @@ -0,0 +1,190 @@ +#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(); + } +} + +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}); + 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); + 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 + // 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); + + // 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; +} + +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::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); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/BoundsTracker.h b/src/BoundsTracker.h new file mode 100644 index 000000000000..ac222989ad09 --- /dev/null +++ b/src/BoundsTracker.h @@ -0,0 +1,176 @@ +#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. 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; + + /** 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..3bdb9424e809 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 @@ -547,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/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/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..2de86684a8e7 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,35 @@ 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 * 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/IRMatch.h b/src/IRMatch.h index 6fa4cadc4eae..b62d4892d74c 100644 --- a/src/IRMatch.h +++ b/src/IRMatch.h @@ -2554,7 +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, {}); - 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; @@ -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/Lower.cpp b/src/Lower.cpp index 753dadb5f6ec..21acfff8df2d 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); @@ -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"; 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/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/src/Schedule.h b/src/Schedule.h index ba3d1eea5ca3..7bfa92981ac1 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -334,6 +334,8 @@ struct ReductionVariable; struct Split { std::string old_var, outer, inner; Expr factor; + 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. 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/Simplify.cpp b/src/Simplify.cpp index 75b908ce8202..2ac9081dd4d1 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()) { @@ -345,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); @@ -370,6 +430,39 @@ 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) { + if (can_prove_depth >= max_can_prove_depth) { + // 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); +} + +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/SimplifyCorrelatedDifferences.cpp b/src/SimplifyCorrelatedDifferences.cpp index dc90aa6d7cef..be06029801ad 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 { @@ -52,11 +54,25 @@ 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) || + // 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) { + debug(4) << "Rewrote " << Expr(op) << " as " << rewrite.result << "\n"; return rewrite.result; } } @@ -79,19 +95,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 +142,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 +157,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 +186,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 +207,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 +278,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/Simplify_Add.cpp b/src/Simplify_Add.cpp index a07ad1b4464b..632d91d80d74 100644 --- a/src/Simplify_Add.cpp +++ b/src/Simplify_Add.cpp @@ -201,6 +201,10 @@ Expr Simplify::visit(const Add *op, ExprInfo *info) { rewrite(x + ((c0 - x) / c1) * c1, c0 - ((c0 - x) % c1), c1 > 0) || rewrite(x + ((c0 - x) / c1 + y) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || rewrite(x + (y + (c0 - x) / c1) * c1, y * c1 - ((c0 - x) % c1) + c0, c1 > 0) || + rewrite(((0 - x) / c0) + ((x % c0 + c1) / c0), (c1 / c0) - (x / c0), c0 > 0 && (c1 + 1) % c0 == 0) || + + // rewrite(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/Simplify_Div.cpp b/src/Simplify_Div.cpp index 4098f6f027e7..d7839340fb99 100644 --- a/src/Simplify_Div.cpp +++ b/src/Simplify_Div.cpp @@ -84,6 +84,18 @@ 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) && // Fold repeated division (rewrite((x / c0) / c2, x / fold(c0 * c2), c0 > 0 && c2 > 0 && !overflows(c0 * c2)) || 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/src/Simplify_Internal.h b/src/Simplify_Internal.h index 94a50bebc644..bf45a9d3977f 100644 --- a/src/Simplify_Internal.h +++ b/src/Simplify_Internal.h @@ -441,6 +441,35 @@ class Simplify : public VariadicVisitor { std::set truths, falsehoods; + // How deeply are we nested inside the conditions of can_prove predicates? + // 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 = 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 !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); + + // Simplify the condition of a can_prove predicate in a rewrite rule, using + // 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 88d3ce2cbf5e..c1c161d6cdc8 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, 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 5203a0c14166..880f2a4b890d 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, 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/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/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) 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 { 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/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 d88c13fa177f..ccb2230844e1 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -339,10 +339,17 @@ tests( specialize_to_gpu.cpp specialize_trim_condition.cpp spirv_ir.cpp + split_aligned.cpp + 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 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 @@ -481,6 +488,10 @@ tests( random.cpp reorder_rvars.cpp 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 thread_safety.cpp 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); 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_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; +} 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; +} diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 981a00e6f0ee..c3dd0b7f8c1c 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), @@ -2377,6 +2414,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 +2454,95 @@ 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}); + + // 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}); + + // 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(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. + 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}); + + // 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}); + + // 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}); + + // 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: + 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); + + // 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 +2555,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"); 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..99ac9743d8c9 --- /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 != 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; +} diff --git a/test/correctness/split_aligned_2d_3x3.cpp b/test/correctness/split_aligned_2d_3x3.cpp new file mode 100644 index 000000000000..d93d80493373 --- /dev/null +++ b/test/correctness/split_aligned_2d_3x3.cpp @@ -0,0 +1,122 @@ +#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"}; + 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); + }; + 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() + .align_bounds(x, 3, offset_x) + .align_bounds(y, 3, offset_y); + } + + 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_nested.cpp b/test/correctness/split_aligned_nested.cpp new file mode 100644 index 000000000000..59b411c614f6 --- /dev/null +++ b/test/correctness/split_aligned_nested.cpp @@ -0,0 +1,100 @@ +#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. +// +// 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; + +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 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); + + 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, ts1) + .split(xo, xoo, xoi, 3, p2, ts2) + .unroll(xi); + + 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; + } + } + } + } + } + } + + printf("Success!\n"); + return 0; +} 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; +} 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; +} 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; +}