From ae8efcdd41d4b70a7f40287ebb7f21a1aa04e358 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 14 Sep 2026 10:08:20 -0700 Subject: [PATCH 1/4] Fix quadratic bounds blow-up at the source Bounds inference scales quadratically in the number of update definitions of a Func. This is because all stages are considered to depend on all earlier stages. This makes bgu's bounds inference stage very slow. The root cause is the quadratic loop over all (consumer, producer) pairs. This can't be made linear, because there could genuinely be a quadratic number of relationships between stages. But we can skip or neuter iterations of this loop to at least not send quadratic amounts of IR downstream. The general problem is cases where A depends on B, B depends on C, and A also directly depends on C, and the transitive dependence of A on C via B is equivalent to the direct dependence. This PR skips once instance of this. It skips the dependence of one update stage on earlier update stages along a particular axis if the very next update stage would have the same dependence on earlier update stages along that axis because the var is pure in both. I.e. this situation: f(x, y) = ... ... some number of update defs ... f(x, 0) += 3; // update def 37 f(x, 2) += 4; // update def 38 There's no need to consider update def 37's dependence on earlier stages, because it's going to be the same as update def 38's. This is more powerful than the existing optimization because it still kicks in if most but not all of the update defs are pure in x. --- src/Bounds.cpp | 7 ++ src/BoundsInference.cpp | 165 ++++++++++++++++++++++------------------ 2 files changed, 97 insertions(+), 75 deletions(-) diff --git a/src/Bounds.cpp b/src/Bounds.cpp index 7aeff2f87bad..6641c038f7d8 100644 --- a/src/Bounds.cpp +++ b/src/Bounds.cpp @@ -1857,6 +1857,13 @@ void merge_boxes(Box &a, const Box &b) { (equal(a.used, !b.used) || equal(!a.used, b.used)); for (size_t i = 0; i < a.size(); i++) { + if (b[i].is_empty()) { + continue; + } + if (a[i].is_empty()) { + a[i] = b[i]; + continue; + } if (!a[i].min.same_as(b[i].min)) { if (a[i].has_lower_bound() && b[i].has_lower_bound()) { if (a_maybe_unused && b_maybe_unused) { diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index ba8266883e4b..7eaf913ab83e 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -55,6 +55,23 @@ bool depends_on_bounds_inference(const Expr &e) { return result; } +// Check if the dimension at index 'dim_idx' is always pure (i.e. equal to 'dim') +// in the definition (including in its specializations) +bool is_dim_always_pure(const Definition &def, const string &dim, int dim_idx) { + const Variable *var = def.args()[dim_idx].as(); + if ((!var) || (var->name != dim)) { + return false; + } + + for (const Specialization &s : def.specializations()) { + bool pure = is_dim_always_pure(s.definition, dim, dim_idx); + if (!pure) { + return false; + } + } + return true; +} + /** Compute the bounds of the value of some variable defined by an * inner let stmt or for loop. E.g. for the stmt: * @@ -360,23 +377,6 @@ class BoundsInference : public IRMutator { } } - // Check if the dimension at index 'dim_idx' is always pure (i.e. equal to 'dim') - // in the definition (including in its specializations) - bool is_dim_always_pure(const Definition &def, const string &dim, int dim_idx) { - const Variable *var = def.args()[dim_idx].as(); - if ((!var) || (var->name != dim)) { - return false; - } - - for (const Specialization &s : def.specializations()) { - bool pure = is_dim_always_pure(s.definition, dim, dim_idx); - if (!pure) { - return false; - } - } - return true; - } - // Wrap a statement in let stmts defining the box Stmt define_bounds(Stmt s, const Function &producing_func, @@ -413,36 +413,6 @@ class BoundsInference : public IRMutator { internal_assert(b.empty() || b.size() == func_args.size()); - if (!b.empty()) { - // Optimization: If a dimension is pure in every update - // step of a func, then there exists a single bound for - // that dimension, instead of one bound per stage. Let's - // figure out what those dimensions are, and just have all - // stages but the last use the bounds for the last stage. - vector always_pure_dims(func_args.size(), true); - for (const Definition &def : func.updates()) { - for (size_t j = 0; j < always_pure_dims.size(); j++) { - bool pure = is_dim_always_pure(def, func_args[j], j); - if (!pure) { - always_pure_dims[j] = false; - } - } - } - - if (stage < func.updates().size()) { - size_t stages = func.updates().size(); - string last_stage = func.name() + ".s" + std::to_string(stages) + "."; - for (size_t i = 0; i < always_pure_dims.size(); i++) { - if (always_pure_dims[i]) { - const string &dim = func_args[i]; - Expr min = Variable::make(Int(32), last_stage + dim + ".min"); - Expr max = Variable::make(Int(32), last_stage + dim + ".max"); - b[i] = Interval(min, max); - } - } - } - } - if (func.has_extern_definition() && !func.extern_definition_proxy_expr().defined()) { // After we define our bounds required, we need to @@ -922,45 +892,90 @@ class BoundsInference : public IRMutator { } } + // For update defs, figure out dimensions where the var is pure, and + // there's another update def after this one where the var is also + // pure. This is used to skip some relationships along some axes + // below. + vector masked_dims; + if (consumer.stage > 0 && + i + 1 < stages.size() && + stages[i + 1].func.same_as(consumer.func)) { + const Definition &this_def = consumer.func.updates()[consumer.stage - 1]; + const Definition &next_def = consumer.func.updates()[consumer.stage]; + for (int k = 0; k < consumer.func.dimensions(); k++) { + const string &arg = consumer.func.args()[k]; + if (is_dim_always_pure(this_def, arg, k) && + is_dim_always_pure(next_def, arg, k)) { + masked_dims.push_back(k); + } + } + } + // Expand the bounds required of all the producers found // (and we are checking until i, because stages are topologically sorted). for (size_t j = 0; j < i; j++) { Stage &producer = stages[j]; + // A consumer depends on *all* stages of a producer, not just the last one. - const Box &b = boxes[producer.func.name()]; - - if (!b.empty()) { - // Check for unboundedness - for (size_t k = 0; k < b.size(); k++) { - if (!b[k].is_bounded()) { - std::ostringstream err; - if (consumer.stage == 0) { - err << "The pure definition "; - } else { - err << "Update definition number " << (consumer.stage - 1); - } - err << " of Function " << consumer.name - << " calls function " << producer.name - << " in an unbounded way in dimension " << k << "\n"; - user_error << err.str(); + auto it = boxes.find(producer.func.name()); + if (it == boxes.end() || + it->second.empty()) { + // No dependence + continue; + } + + if (producer.func.same_as(consumer.func) && + masked_dims.size() == (size_t)producer.func.dimensions()) { + // This self-bounds relationship is completely masked by + // another one, so just skip it. + continue; + } + + Box b = it->second; + + // Check for unboundedness + for (size_t k = 0; k < b.size(); k++) { + if (!b[k].is_bounded()) { + std::ostringstream err; + if (consumer.stage == 0) { + err << "The pure definition "; + } else { + err << "Update definition number " << (consumer.stage - 1); } + err << " of Function " << consumer.name + << " calls function " << producer.name + << " in an unbounded way in dimension " << k << "\n"; + user_error << err.str(); } + } - // Dump out the region required of each stage for debugging. - /* - debug(0) << "Box required of " << producer.name - << " by " << consumer.name - << " stage " << consumer.stage << ":\n" - << " used: " << b.used << "\n"; - for (size_t k = 0; k < b.size(); k++) { - debug(0) << " " << b[k].min << " ... " << b[k].max << "\n"; + // If producer = consumer, and a dim is pure, and there's + // another update def after this consumer where the dim is + // also pure, forget the dependence along this axis - it's + // redundant with the next stage. This avoids quadratic + // blow-up of bounds expressions for Funcs with lots of + // update stages where long runs of them share pure vars. + if (producer.func.same_as(consumer.func)) { + for (int k : masked_dims) { + b[k] = Interval::nothing(); } - debug(0) << "\n"; - */ + } - producer.bounds[{consumer.name, consumer.stage}] = b; - producer.consumers.push_back((int)i); + /* + // Dump out the region required of each stage for debugging. + debug(0) << "Box required of " << producer.name + << " stage " << producer.stage << ":\n" + << " by " << consumer.name + << " stage " << consumer.stage << ":\n" + << " used: " << b.used << "\n"; + for (size_t k = 0; k < b.size(); k++) { + debug(0) << " " << b[k].min << " ... " << b[k].max << "\n"; } + debug(0) << "\n"; + */ + + producer.bounds[{consumer.name, consumer.stage}] = std::move(b); + producer.consumers.push_back((int)i); } } From 4fe5f0c6c7f3a18718e4cdaf6888a0efcb58a195 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 14 Sep 2026 10:19:28 -0700 Subject: [PATCH 2/4] uncomment debug --- src/BoundsInference.cpp | 40 ++++++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index 7eaf913ab83e..b18e38c89c71 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -834,12 +834,14 @@ class BoundsInference : public IRMutator { } // Dump the stages post-inlining for debugging - /* - debug(0) << "Bounds inference stages after inlining: \n"; - for (size_t i = 0; i < stages.size(); i++) { - debug(0) << " " << i << ") " << stages[i].name << "\n"; - } - */ + debug(4) << [&] { + std::ostringstream s; + s << "Bounds inference stages after inlining: \n"; + for (size_t i = 0; i < stages.size(); i++) { + s << " " << i << ") " << stages[i].name << "\n"; + } + return s.str(); + }(); // Then compute relationships between them. for (size_t i = 0; i < stages.size(); i++) { @@ -961,18 +963,20 @@ class BoundsInference : public IRMutator { } } - /* - // Dump out the region required of each stage for debugging. - debug(0) << "Box required of " << producer.name - << " stage " << producer.stage << ":\n" - << " by " << consumer.name - << " stage " << consumer.stage << ":\n" - << " used: " << b.used << "\n"; - for (size_t k = 0; k < b.size(); k++) { - debug(0) << " " << b[k].min << " ... " << b[k].max << "\n"; - } - debug(0) << "\n"; - */ + debug(4) << [&] { + std::ostringstream s; + // Dump out the region required of each stage for debugging. + s << "Box required of " << producer.name + << " stage " << producer.stage << ":\n" + << " by " << consumer.name + << " stage " << consumer.stage << ":\n" + << " used: " << b.used << "\n"; + for (size_t k = 0; k < b.size(); k++) { + s << " " << b[k].min << " ... " << b[k].max << "\n"; + } + s << "\n"; + return s.str(); + }(); producer.bounds[{consumer.name, consumer.stage}] = std::move(b); producer.consumers.push_back((int)i); From 3cb716ce66f6fcdcabe85ec84208787bcf87ad38 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Mon, 14 Sep 2026 10:20:59 -0700 Subject: [PATCH 3/4] Move comment --- src/BoundsInference.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index b18e38c89c71..6b01472dead4 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -963,9 +963,9 @@ class BoundsInference : public IRMutator { } } + // Dump out the region required of each stage for debugging. debug(4) << [&] { std::ostringstream s; - // Dump out the region required of each stage for debugging. s << "Box required of " << producer.name << " stage " << producer.stage << ":\n" << " by " << consumer.name From 5896dd82059e03b53dd966709d09e562aaa0fc62 Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Tue, 15 Sep 2026 10:08:18 -0700 Subject: [PATCH 4/4] Skip more kinds of relationships --- src/BoundsInference.cpp | 127 ++++++++++++++++++++++------------------ 1 file changed, 70 insertions(+), 57 deletions(-) diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index 6b01472dead4..03071487abf5 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -55,23 +55,6 @@ bool depends_on_bounds_inference(const Expr &e) { return result; } -// Check if the dimension at index 'dim_idx' is always pure (i.e. equal to 'dim') -// in the definition (including in its specializations) -bool is_dim_always_pure(const Definition &def, const string &dim, int dim_idx) { - const Variable *var = def.args()[dim_idx].as(); - if ((!var) || (var->name != dim)) { - return false; - } - - for (const Specialization &s : def.specializations()) { - bool pure = is_dim_always_pure(s.definition, dim, dim_idx); - if (!pure) { - return false; - } - } - return true; -} - /** Compute the bounds of the value of some variable defined by an * inner let stmt or for loop. E.g. for the stmt: * @@ -241,6 +224,49 @@ class BoundsInference : public IRMutator { string stage_prefix; size_t fused_group_index; Inliner *inliner; + vector pure_dims_shared_with_next_stage; + + Stage(const Function &func, size_t stage, size_t fused_group_index, Inliner *inliner) + : func(func), + stage(stage), + name(func.name()), + stage_prefix(func.name() + ".s" + std::to_string(stage) + "."), + fused_group_index(fused_group_index), + inliner(inliner) { + + compute_exprs(); + + const Definition &this_def = stage == 0 ? func.definition() : func.updates()[stage - 1]; + if (stage < func.updates().size()) { + const Definition &next_def = func.updates()[stage]; + for (int k = 0; k < func.dimensions(); k++) { + const string &arg = func.args()[k]; + if (is_dim_always_pure(this_def, arg, k) && + is_dim_always_pure(next_def, arg, k)) { + pure_dims_shared_with_next_stage.push_back(k); + } + } + } + } + + // Check if the dimension at index 'dim_idx' is always pure (i.e. equal to 'dim') + // in the definition (including in its specializations) + static bool is_dim_always_pure(const Definition &def, const string &dim, int dim_idx) { + internal_assert(def.defined()); + internal_assert((size_t)dim_idx < def.args().size()); + const Variable *var = def.args()[dim_idx].as(); + if ((!var) || (var->name != dim)) { + return false; + } + + for (const Specialization &s : def.specializations()) { + bool pure = is_dim_always_pure(s.definition, dim, dim_idx); + if (!pure) { + return false; + } + } + return true; + } // Computed expressions on the left and right-hand sides. // Note that a function definition might have different LHS or reduction domain @@ -807,21 +833,10 @@ class BoundsInference : public IRMutator { continue; } - Stage s; - s.func = f[i]; - s.stage = 0; - s.name = s.func.name(); - s.fused_group_index = find_fused_group_index(s.func, fused_groups); - s.compute_exprs(); - s.stage_prefix = s.name + ".s0."; - s.inliner = &inliner; - stages.push_back(s); - + int fused_group_index = find_fused_group_index(f[i], fused_groups); + stages.emplace_back(f[i], 0, fused_group_index, &inliner); for (size_t j = 0; j < f[i].updates().size(); j++) { - s.stage = (int)(j + 1); - s.stage_prefix = s.name + ".s" + std::to_string(s.stage) + "."; - s.compute_exprs(); - stages.push_back(s); + stages.emplace_back(f[i], (int)(j + 1), fused_group_index, &inliner); } } @@ -894,25 +909,6 @@ class BoundsInference : public IRMutator { } } - // For update defs, figure out dimensions where the var is pure, and - // there's another update def after this one where the var is also - // pure. This is used to skip some relationships along some axes - // below. - vector masked_dims; - if (consumer.stage > 0 && - i + 1 < stages.size() && - stages[i + 1].func.same_as(consumer.func)) { - const Definition &this_def = consumer.func.updates()[consumer.stage - 1]; - const Definition &next_def = consumer.func.updates()[consumer.stage]; - for (int k = 0; k < consumer.func.dimensions(); k++) { - const string &arg = consumer.func.args()[k]; - if (is_dim_always_pure(this_def, arg, k) && - is_dim_always_pure(next_def, arg, k)) { - masked_dims.push_back(k); - } - } - } - // Expand the bounds required of all the producers found // (and we are checking until i, because stages are topologically sorted). for (size_t j = 0; j < i; j++) { @@ -927,7 +923,8 @@ class BoundsInference : public IRMutator { } if (producer.func.same_as(consumer.func) && - masked_dims.size() == (size_t)producer.func.dimensions()) { + consumer.pure_dims_shared_with_next_stage.size() == + (size_t)producer.func.dimensions()) { // This self-bounds relationship is completely masked by // another one, so just skip it. continue; @@ -952,13 +949,29 @@ class BoundsInference : public IRMutator { } // If producer = consumer, and a dim is pure, and there's - // another update def after this consumer where the dim is - // also pure, forget the dependence along this axis - it's - // redundant with the next stage. This avoids quadratic - // blow-up of bounds expressions for Funcs with lots of - // update stages where long runs of them share pure vars. + // another update def after this consumer where the dim is also + // pure, forget the dependence along this axis - it's redundant + // with the next stage. This avoids quadratic blow-up of bounds + // expressions for Funcs with lots of update stages where long + // runs of them share pure vars. if (producer.func.same_as(consumer.func)) { - for (int k : masked_dims) { + for (int k : consumer.pure_dims_shared_with_next_stage) { + b[k] = Interval::nothing(); + } + } + + // On the other hand, if the producer has another stage after it + // that shares the same pure vars, and that second stage refers + // to the earlier stage, we don't need to register our + // dependence on that consumer. It happens transitively via the + // next stage, and the dependence between update stages is + // constrained to be elementwise in the pure vars. + const Stage &next = stages[j + 1]; + if (!consumer.func.same_as(producer.func) && + j + 1 < i && + next.func.same_as(producer.func) && + producer.bounds.find({next.func.name(), next.stage}) != producer.bounds.end()) { + for (int k : producer.pure_dims_shared_with_next_stage) { b[k] = Interval::nothing(); } }