Skip to content

Fix quadratic bounds blow-up at the source #9448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
abadams merged 4 commits into from
Sep 16, 2026
+133 −94
Merged

Fix quadratic bounds blow-up at the source #9448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ae8efcd
Fix quadratic bounds blow-up at the source
abadams Sep 14, 2026
4fe5f0c
uncomment debug
abadams Sep 14, 2026
3cb716c
Move comment
abadams Sep 14, 2026
5896dd8
Skip more kinds of relationships
abadams Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/Bounds.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
220 changes: 126 additions & 94 deletions src/BoundsInference.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,49 @@ class BoundsInference : public IRMutator {
string stage_prefix;
size_t fused_group_index;
Inliner *inliner;
vector<int> 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<Variable>();
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
Expand Down Expand Up @@ -360,23 +403,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<Variable>();
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,
Expand Down Expand Up @@ -413,36 +439,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<bool> 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
Expand Down Expand Up @@ -837,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);
}
}

Expand All @@ -864,12 +849,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++) {
Expand Down Expand Up @@ -926,41 +913,86 @@ class BoundsInference : public IRMutator {
// (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()];
auto it = boxes.find(producer.func.name());
if (it == boxes.end() ||
it->second.empty()) {
// No dependence
continue;
}

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();
if (producer.func.same_as(consumer.func) &&
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;
}

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();
}
Comment on lines +937 to 948

@alexreinking alexreinking Sep 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

user_assert(b[k].is_bounded()) << [&] {
    std::ostringstream err;
    // ...
    return 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 : consumer.pure_dims_shared_with_next_stage) {
b[k] = Interval::nothing();
}
debug(0) << "\n";
*/
}

producer.bounds[{consumer.name, consumer.stage}] = b;
producer.consumers.push_back((int)i);
// 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();
}
}

// Dump out the region required of each stage for debugging.
debug(4) << [&] {
std::ostringstream s;
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);
}
}

Expand Down
Loading