Merge repeated Func bound/align_bounds/align_extent calls per Var

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];

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);

src/Deserialization.cpp

@@ -1025,8 +1025,12 @@ FuncSchedule Deserializer::deserialize_func_schedule(const Serialize::FuncSchedu
     const std::vector<StorageDim> storage_dims =
         deserialize_vector<Serialize::StorageDim, StorageDim>(func_schedule->storage_dims(),
                                                               &Deserializer::deserialize_storage_dim);
-    const std::vector<Bound> bounds = deserialize_vector<Serialize::Bound, Bound>(func_schedule->bounds(),
-                                                                                  &Deserializer::deserialize_bound);
+    const std::vector<Bound> bounds_vec = deserialize_vector<Serialize::Bound, Bound>(func_schedule->bounds(),
+                                                                                      &Deserializer::deserialize_bound);
+    std::map<std::string, Bound> bounds;
+    for (const auto &b : bounds_vec) {
+        merge_bound(bounds, b);
+    }
     const std::vector<Bound> estimates = deserialize_vector<Serialize::Bound, Bound>(func_schedule->estimates(),
                                                                                      &Deserializer::deserialize_bound);
     const std::map<std::string, FunctionPtr> wrappers = deserialize_wrapper_refs(func_schedule->wrappers());

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;
 }
 

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 "

src/Schedule.cpp

@@ -235,7 +235,7 @@ struct FuncScheduleContents {
 
     LoopLevel store_level, compute_level, hoist_storage_level;
     std::vector<StorageDim> storage_dims;
-    std::vector<Bound> bounds;
+    std::map<std::string, Bound> bounds;
     std::vector<Bound> estimates;
     std::map<std::string, Internal::FunctionPtr> 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<StorageDim> &FuncSchedule::storage_dims() const {
     return contents->storage_dims;
 }
 
-std::vector<Bound> &FuncSchedule::bounds() {
+std::map<std::string, Bound> &FuncSchedule::bounds() {
     return contents->bounds;
 }
 
-const std::vector<Bound> &FuncSchedule::bounds() const {
+const std::map<std::string, Bound> &FuncSchedule::bounds() const {
     return contents->bounds;
 }
 
+void merge_bound(std::map<std::string, Bound> &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<Bound> &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);
         }

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<std::string, Bound> &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<Bound> &bounds() const;
-    std::vector<Bound> &bounds();
+    const std::map<std::string, Bound> &bounds() const;
+    std::map<std::string, Bound> &bounds();
     // @}
 
     /** You may explicitly specify an estimate of some of the function

src/ScheduleFunctions.cpp

@@ -206,7 +206,8 @@ Stmt build_loop_nest(
     map<string, Expr> 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<string, Function> &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";

src/Serialization.cpp

@@ -1126,8 +1126,8 @@ Offset<Serialize::FuncSchedule> Serializer::serialize_func_schedule(FlatBufferBu
         storage_dims_serialized.push_back(serialize_storage_dim(builder, storage_dim));
     }
     std::vector<Offset<Serialize::Bound>> 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<Offset<Serialize::Bound>> estimates_serialized;
     for (const auto &estimate : func_schedule.estimates()) {

src/autoschedulers/adams2019/FunctionDAG.cpp

@@ -902,12 +902,12 @@ FunctionDAG::FunctionDAG(const vector<Function> &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

src/autoschedulers/anderson2021/FunctionDAG.cpp

@@ -892,12 +892,12 @@ FunctionDAG::FunctionDAG(const vector<Function> &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

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

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 <stdio.h>
+
+using namespace Halide;
+using namespace Halide::Internal;
+
+namespace {
+
+const Bound &the_bound(const Func &f, const std::string &var_name) {
+    const std::map<std::string, Bound> &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;
+}