diff --git a/src/AllocationBoundsInference.cpp b/src/AllocationBoundsInference.cpp index 9aaaf1c7e661..97acaccfc226 100644 --- a/src/AllocationBoundsInference.cpp +++ b/src/AllocationBoundsInference.cpp @@ -57,10 +57,9 @@ class AllocationInference : public IRMutator { for (size_t i = 0; i < b.size(); i++) { // Get any applicable bound on this dimension Bound bound; - for (const auto &b : f.schedule().bounds()) { - if (f_args[i] == b.var) { - bound = b; - } + auto it = f.schedule().bounds().find(f_args[i]); + if (it != f.schedule().bounds().end()) { + bound = it->second; } string prefix = op->name + "." + f_args[i]; diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index ba8266883e4b..73554dbaa1e2 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -530,7 +530,8 @@ class BoundsInference : public IRMutator { LoopLevel compute_at = func.schedule().compute_level(); LoopLevel store_at = func.schedule().store_level(); - for (auto bound : func.schedule().bounds()) { + for (const auto &entry : func.schedule().bounds()) { + Bound bound = entry.second; string min_var = prefix + bound.var + ".min"; string max_var = prefix + bound.var + ".max"; Expr min_required = Variable::make(Int(32), min_var); diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index 7a76edc1193c..2d0ce37ede90 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -1025,8 +1025,12 @@ FuncSchedule Deserializer::deserialize_func_schedule(const Serialize::FuncSchedu const std::vector storage_dims = deserialize_vector(func_schedule->storage_dims(), &Deserializer::deserialize_storage_dim); - const std::vector bounds = deserialize_vector(func_schedule->bounds(), - &Deserializer::deserialize_bound); + const std::vector bounds_vec = deserialize_vector(func_schedule->bounds(), + &Deserializer::deserialize_bound); + std::map bounds; + for (const auto &b : bounds_vec) { + merge_bound(bounds, b); + } const std::vector estimates = deserialize_vector(func_schedule->estimates(), &Deserializer::deserialize_bound); const std::map wrappers = deserialize_wrapper_refs(func_schedule->wrappers()); diff --git a/src/Func.cpp b/src/Func.cpp index 468188530c67..f9b4cf189ea4 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -2746,7 +2746,7 @@ Func &Func::bound(const Var &var, Expr min, Expr extent) { << " is not one of the pure variables of " << name() << ".\n"; Bound b = {var.name(), min, extent, Expr(), Expr()}; - func.schedule().bounds().push_back(b); + merge_bound(func.schedule().bounds(), b); // Propagate constant bounds into estimates as well. if (!is_const(min)) { @@ -2831,7 +2831,7 @@ Func &Func::align_bounds(const Var &var, Expr modulus, Expr remainder) { << " is not one of the pure variables of " << name() << ".\n"; Bound b = {var.name(), Expr(), Expr(), modulus, remainder}; - func.schedule().bounds().push_back(b); + merge_bound(func.schedule().bounds(), b); return *this; } @@ -2851,7 +2851,7 @@ Func &Func::align_extent(const Var &var, Expr modulus) { << " is not one of the pure variables of " << name() << ".\n"; Bound b = {var.name(), Expr(), Expr(), modulus, Expr()}; - func.schedule().bounds().push_back(b); + merge_bound(func.schedule().bounds(), b); return *this; } diff --git a/src/Inline.cpp b/src/Inline.cpp index f5b5a260ec8a..0ad37e4e6b7d 100644 --- a/src/Inline.cpp +++ b/src/Inline.cpp @@ -79,7 +79,8 @@ void validate_schedule_inlined_function(Function f) { } } - for (const auto &b : func_s.bounds()) { + for (const auto &entry : func_s.bounds()) { + const Bound &b = entry.second; if (b.min.defined()) { user_warning << "It is meaningless to bound dimension " << b.var << " of function " diff --git a/src/Schedule.cpp b/src/Schedule.cpp index 948233112b7c..8d5a99e46836 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -235,7 +235,7 @@ struct FuncScheduleContents { LoopLevel store_level, compute_level, hoist_storage_level; std::vector storage_dims; - std::vector bounds; + std::map bounds; std::vector estimates; std::map wrappers; MemoryType memory_type = MemoryType::Auto; @@ -256,7 +256,8 @@ struct FuncScheduleContents { // Pass an IRMutator through to all Exprs referenced in the FuncScheduleContents void mutate(IRMutator &mutator) { - for (Bound &b : bounds) { + for (auto &entry : bounds) { + Bound &b = entry.second; if (b.min.defined()) { b.min = mutator(b.min); } @@ -450,14 +451,34 @@ const std::vector &FuncSchedule::storage_dims() const { return contents->storage_dims; } -std::vector &FuncSchedule::bounds() { +std::map &FuncSchedule::bounds() { return contents->bounds; } -const std::vector &FuncSchedule::bounds() const { +const std::map &FuncSchedule::bounds() const { return contents->bounds; } +void merge_bound(std::map &bounds, const Bound &b) { + auto [it, inserted] = bounds.try_emplace(b.var, b); + if (inserted) { + return; + } + Bound &existing = it->second; + if (b.min.defined()) { + existing.min = b.min; + } + if (b.extent.defined()) { + existing.extent = b.extent; + } + if (b.modulus.defined()) { + existing.modulus = b.modulus; + } + if (b.remainder.defined()) { + existing.remainder = b.remainder; + } +} + std::vector &FuncSchedule::estimates() { return contents->estimates; } @@ -512,7 +533,8 @@ const LoopLevel &FuncSchedule::hoist_storage_level() const { } void FuncSchedule::accept(IRVisitor *visitor) const { - for (const Bound &b : bounds()) { + for (const auto &entry : bounds()) { + const Bound &b = entry.second; if (b.min.defined()) { b.min.accept(visitor); } diff --git a/src/Schedule.h b/src/Schedule.h index ba3d1eea5ca3..b68e63ed4d31 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -508,6 +508,15 @@ struct Bound { Expr modulus, remainder; }; +/** Merge \p b into \p bounds, keyed by \p b.var. Func::bound/bound_extent/ + * align_bounds/align_extent each set one or more of a Var's four Bound + * fields (min, extent, modulus, remainder) and leave the rest undefined. + * A field that \p b doesn't set is left untouched on an existing entry for + * that Var; a field it does set overwrites whatever was there, so every + * consumer of FuncSchedule::bounds() sees at most one constraint per Var, + * with no ordering between calls left for them to get wrong. */ +void merge_bound(std::map &bounds, const Bound &b); + /** Properties of one axis of the storage of a Func */ struct StorageDim { /** The var in the pure definition corresponding to this axis */ @@ -635,10 +644,11 @@ class FuncSchedule { /** You may explicitly bound some of the dimensions of a function, * or constrain them to lie on multiples of a given factor. See - * \ref Func::bound and \ref Func::align_bounds and \ref Func::align_extent. */ + * \ref Func::bound and \ref Func::align_bounds and \ref Func::align_extent. + * At most one Bound is kept per Var, keyed by its name. */ // @{ - const std::vector &bounds() const; - std::vector &bounds(); + const std::map &bounds() const; + std::map &bounds(); // @} /** You may explicitly specify an estimate of some of the function diff --git a/src/ScheduleFunctions.cpp b/src/ScheduleFunctions.cpp index f2a31b1b6562..c0f992cb3bc9 100644 --- a/src/ScheduleFunctions.cpp +++ b/src/ScheduleFunctions.cpp @@ -206,7 +206,8 @@ Stmt build_loop_nest( map dim_extent_alignment; // First hunt through the bounds for them. - for (const Bound &i : func_s.bounds()) { + for (const auto &entry : func_s.bounds()) { + const Bound &i = entry.second; if (i.extent.defined()) { dim_extent_alignment[i.var] = i.extent; } @@ -950,7 +951,8 @@ Stmt build_extern_produce(const map &env, Function f, const Ta Stmt inject_explicit_bounds(Stmt body, Function func) { const FuncSchedule &s = func.schedule(); for (size_t stage = 0; stage <= func.updates().size(); stage++) { - for (auto b : s.bounds()) { + for (const auto &entry : s.bounds()) { + Bound b = entry.second; string prefix = func.name() + ".s" + std::to_string(stage) + "." + b.var; string min_name = prefix + ".min_unbounded"; string max_name = prefix + ".max_unbounded"; diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 2dd7bf4f33aa..128fa074fb7a 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -1126,8 +1126,8 @@ Offset Serializer::serialize_func_schedule(FlatBufferBu storage_dims_serialized.push_back(serialize_storage_dim(builder, storage_dim)); } std::vector> bounds_serialized; - for (const auto &bound : func_schedule.bounds()) { - bounds_serialized.push_back(serialize_bound(builder, bound)); + for (const auto &entry : func_schedule.bounds()) { + bounds_serialized.push_back(serialize_bound(builder, entry.second)); } std::vector> estimates_serialized; for (const auto &estimate : func_schedule.estimates()) { diff --git a/src/autoschedulers/adams2019/FunctionDAG.cpp b/src/autoschedulers/adams2019/FunctionDAG.cpp index 974182970f54..a8e09d55a292 100644 --- a/src/autoschedulers/adams2019/FunctionDAG.cpp +++ b/src/autoschedulers/adams2019/FunctionDAG.cpp @@ -902,12 +902,12 @@ FunctionDAG::FunctionDAG(const vector &outputs, const Target &target) estimates[b.var] = Span(*i_min, *i_min + *i_extent - 1, false); } } - for (const auto &b : consumer.schedule().bounds()) { + for (const auto &[var_name, b] : consumer.schedule().bounds()) { auto i_min = as_const_int(b.min); auto i_extent = as_const_int(b.extent); if (i_min && i_extent) { // It's a true bound, not just an estimate - estimates[b.var] = Span(*i_min, *i_min + *i_extent - 1, true); + estimates[var_name] = Span(*i_min, *i_min + *i_extent - 1, true); } } // Set the bounds using the estimates diff --git a/src/autoschedulers/anderson2021/FunctionDAG.cpp b/src/autoschedulers/anderson2021/FunctionDAG.cpp index 07e63280ff12..5fab98e05387 100644 --- a/src/autoschedulers/anderson2021/FunctionDAG.cpp +++ b/src/autoschedulers/anderson2021/FunctionDAG.cpp @@ -892,12 +892,12 @@ FunctionDAG::FunctionDAG(const vector &outputs, const Target &target) estimates[b.var] = Span(*i_min, *i_min + *i_extent - 1, false); } } - for (const auto &b : consumer.schedule().bounds()) { + for (const auto &[var_name, b] : consumer.schedule().bounds()) { auto i_min = as_const_int(b.min); auto i_extent = as_const_int(b.extent); if (i_min && i_extent) { // It's a true bound, not just an estimate - estimates[b.var] = Span(*i_min, *i_min + *i_extent - 1, true); + estimates[var_name] = Span(*i_min, *i_min + *i_extent - 1, true); } } // Set the bounds using the estimates diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index d88c13fa177f..8fd1c106861d 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -18,6 +18,7 @@ tests( bool_compute_root_vectorize.cpp bool_predicate_cast.cpp bound.cpp + bound_merge_order.cpp bound_small_allocations.cpp bound_storage.cpp boundary_conditions.cpp diff --git a/test/correctness/bound_merge_order.cpp b/test/correctness/bound_merge_order.cpp new file mode 100644 index 000000000000..cad9ec592cb6 --- /dev/null +++ b/test/correctness/bound_merge_order.cpp @@ -0,0 +1,119 @@ +// Func::bound()/bound_extent()/align_bounds()/align_extent() each append a +// constraint to a Var's Bound. Multiple calls for the same Var must merge +// into a single Bound record, with a result that doesn't depend on which +// order the calls were made in -- BoundsInference (and every other consumer +// of FuncSchedule::bounds()) only ever expects to see one Bound per Var. +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +namespace { + +const Bound &the_bound(const Func &f, const std::string &var_name) { + const std::map &bounds = f.function().schedule().bounds(); + auto it = bounds.find(var_name); + if (it == bounds.end()) { + printf("Expected exactly one Bound for \"%s\", found 0\n", var_name.c_str()); + exit(1); + } + return it->second; +} + +void expect_extent_bound(const Func &f, const std::string &var_name, int extent) { + const Bound &b = the_bound(f, var_name); + if (!b.extent.defined() || !is_const(simplify(b.extent), extent)) { + printf("Expected extent bound %d on \"%s\", got %s\n", + extent, var_name.c_str(), b.extent.defined() ? "a different value" : "undefined"); + exit(1); + } +} + +void expect_alignment(const Func &f, const std::string &var_name, int modulus, int remainder) { + const Bound &b = the_bound(f, var_name); + if (!b.modulus.defined() || !is_const(simplify(b.modulus), modulus)) { + printf("Expected alignment modulus %d on \"%s\", got %s\n", + modulus, var_name.c_str(), b.modulus.defined() ? "a different value" : "undefined"); + exit(1); + } + if (!b.remainder.defined() || !is_const(simplify(b.remainder), remainder)) { + printf("Expected alignment remainder %d on \"%s\", got %s\n", + remainder, var_name.c_str(), b.remainder.defined() ? "a different value" : "undefined"); + exit(1); + } +} + +} // namespace + +int main(int argc, char **argv) { + Var x("x"), y("y"); + + // bound_extent() then align_bounds(): must merge into one Bound with + // both the extent and the alignment set. + { + Func f("f"); + f(x) = x; + f.bound_extent(x, 3); + f.align_bounds(x, 3, 1); + expect_extent_bound(f, "x", 3); + expect_alignment(f, "x", 3, 1); + } + + // The same two calls, in the opposite order: same merged result. + { + Func f("f"); + f(x) = x; + f.align_bounds(x, 3, 1); + f.bound_extent(x, 3); + expect_extent_bound(f, "x", 3); + expect_alignment(f, "x", 3, 1); + } + + // bound()/align_bounds() interleaved across two Vars: each Var still + // ends up with exactly one merged Bound, unaffected by the other Var's + // calls in between. + { + Func f("f"); + f(x, y) = x + y; + f.align_bounds(x, 4, 2); + f.bound(y, 0, 5); + f.bound_extent(x, 4); + f.align_bounds(y, 5, 0); + expect_extent_bound(f, "x", 4); + expect_alignment(f, "x", 4, 2); + expect_alignment(f, "y", 5, 0); + const Bound &by = the_bound(f, "y"); + if (!by.min.defined() || !is_const(simplify(by.min), 0)) { + printf("Expected min bound 0 on \"y\"\n"); + return 1; + } + if (!by.extent.defined() || !is_const(simplify(by.extent), 5)) { + printf("Expected extent bound 5 on \"y\"\n"); + return 1; + } + } + + // align_extent() only sets modulus, so it updates just that field on an + // existing Bound and leaves a remainder set by a prior align_bounds() + // alone -- merging never clobbers a field a call didn't itself set. + { + Func f("f"); + f(x) = x; + f.align_bounds(x, 3, 1); + f.align_extent(x, 4); + const Bound &b = the_bound(f, "x"); + if (!b.modulus.defined() || !is_const(simplify(b.modulus), 4)) { + printf("Expected alignment modulus 4 on \"x\" after align_extent(), got %s\n", + b.modulus.defined() ? "a different value" : "undefined"); + return 1; + } + if (!b.remainder.defined() || !is_const(simplify(b.remainder), 1)) { + printf("Expected align_extent() to leave the remainder set by a prior align_bounds() alone\n"); + return 1; + } + } + + printf("Success!\n"); + return 0; +}