From 089a07469a11900154e797f7759d22016b1c8f0f Mon Sep 17 00:00:00 2001 From: Andrew Adams Date: Fri, 11 Sep 2026 11:25:23 -0700 Subject: [PATCH 1/2] Speed up bounds inference for funcs with many pure update dims For dimensions that are pure in every update stage of a Func, there is a single required region shared by all stages rather than one per stage. Previously populate_scope still created a distinct per-stage bound variable (f.sK.x.min/max) for these dimensions, so when a many-staged Func was a consumer, merge_boxes had to combine one structurally distinct term per stage, building large min/max expressions and calling simplify on each. define_bounds then discarded those pure-dimension results and aliased them to the last stage anyway. Instead, key the required region of an always-pure dimension off the last stage's bound variables directly in populate_scope. The per-stage boxes are then structurally identical in that dimension and collapse on merge instead of growing one term per stage. On the bgu app (whose solve is expressed as a Func with 100 update stages, pure in x/y/z), this drops computation bounds inference from ~381ms to ~23ms and total lowering from ~578ms to ~210ms, with no change in generated code. Co-Authored-By: Claude Opus 4.8 --- src/BoundsInference.cpp | 60 +++++++++++++++++++++++++++++------------ 1 file changed, 43 insertions(+), 17 deletions(-) diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index ba8266883e4b..eb10fde2c3ea 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -225,6 +225,12 @@ class BoundsInference : public IRMutator { size_t fused_group_index; Inliner *inliner; + // Which dimensions of 'func' are pure (i.e. equal to the pure + // Var/RVar in that position) across the pure definition and every + // update definition. Shared by every Stage of the same Func, and + // computed once by compute_always_pure_dims() below. + vector always_pure_dims; + // Computed expressions on the left and right-hand sides. // Note that a function definition might have different LHS or reduction domain // (if it's an update def) or RHS per specialization. All specializations @@ -377,6 +383,24 @@ class BoundsInference : public IRMutator { return true; } + // Populate 'always_pure_dims': for each dimension of 'func', whether + // it's pure (i.e. equal to the pure Var/RVar in that position) across + // the pure definition and every update definition. Only depends on + // 'func', so it only needs to be computed once and shared across all + // Stages of the same Func. + void compute_always_pure_dims() { + const vector func_args = func.args(); + always_pure_dims.assign(func_args.size(), true); + for (const Definition &def : func.updates()) { + for (size_t j = 0; j < always_pure_dims.size(); j++) { + if (always_pure_dims[j] && + !is_dim_always_pure(def, func_args[j], (int)j)) { + always_pure_dims[j] = false; + } + } + } + } + // Wrap a statement in let stmts defining the box Stmt define_bounds(Stmt s, const Function &producing_func, @@ -414,21 +438,13 @@ 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; - } - } - } - + // Optimization: If a dimension is pure in every update step + // of a func, then there's a single bound for that dimension, + // instead of one bound per stage. All stages but the last use + // the bounds for the last stage. populate_scope already keys + // the required region of pure dims off the last stage's bound + // variables, so here we just define the current stage's bound + // variables to alias the last stage's. if (stage < func.updates().size()) { size_t stages = func.updates().size(); string last_stage = func.name() + ".s" + std::to_string(stages) + "."; @@ -769,8 +785,17 @@ class BoundsInference : public IRMutator { // We need to take into account specializations which may refer to // different reduction variables as well. void populate_scope(Scope &result) { - for (const string &farg : func.args()) { - string arg = name + ".s" + std::to_string(stage) + "." + farg; + const vector func_args = func.args(); + for (size_t i = 0; i < func_args.size(); i++) { + const string &farg = func_args[i]; + // If a dimension is pure in every stage, all stages share a + // single bound (that of the last stage), so key the required + // region off the last stage's bound variables. This keeps the + // per-stage boxes structurally identical in this dimension, so + // they collapse when merged rather than growing one term per + // stage. + size_t bound_stage = always_pure_dims[i] ? func.updates().size() : stage; + string arg = name + ".s" + std::to_string(bound_stage) + "." + farg; result.push(farg, Interval(Variable::make(Int(32), arg + ".min"), Variable::make(Int(32), arg + ".max"))); @@ -842,6 +867,7 @@ class BoundsInference : public IRMutator { s.stage = 0; s.name = s.func.name(); s.fused_group_index = find_fused_group_index(s.func, fused_groups); + s.compute_always_pure_dims(); s.compute_exprs(); s.stage_prefix = s.name + ".s0."; s.inliner = &inliner; From eae5a5d5e6dde812774f3f8507c118fceb5d7054 Mon Sep 17 00:00:00 2001 From: Martijn Courteaux Date: Sat, 12 Sep 2026 15:25:33 +0200 Subject: [PATCH 2/2] Collapse pure-dim bounds at the loop level that knows the bounds Keying the required region of an always-pure dimension off the last stage's bound variables in populate_scope is wrong: those variables are only narrowed to the current iteration for the stage that owns the loop nest we are in. A producer computed inside stage k's nest therefore saw the last stage's unnarrowed bound for such a dimension. gpu_mixed_dimensionality hits this directly. It tiles the pure definition of out over x, y and z but its update over x and y only, so out.s1 has no z loop to narrow out.s1.z.min/max. Inside out.s0's nest the region required of h in z widened from the current 4-tile to the full extent, and cascaded through h, g and f: the kernel ended up with 64x64x4 thread extents and a 2MB per-thread local depot, and the test no longer finished. Do the collapse in define_bounds instead, which knows what loop level it is at. Each entry of a stage's bounds now carries the region both in terms of the consumer stage's own bound variables and with the always-pure dimensions phrased via the last stage's. A consumer stage that owns or is fused with the current loop nest uses the former; a consumer produced further in has every stage's bounds defined here, so the two alias each other and the latter is equivalent and smaller. The last stage's bound variables are built once per Func and shared by all of its stages, so merge_boxes' same_as fast path fires. On bgu, computation bounds inference goes from ~746ms to ~25ms. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016Gn5nH7e378aGqdHmo7Mdm --- src/BoundsInference.cpp | 149 ++++++++++++++++++++++++++++++++++------ 1 file changed, 128 insertions(+), 21 deletions(-) diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index eb10fde2c3ea..1ffed45fe6dd 100644 --- a/src/BoundsInference.cpp +++ b/src/BoundsInference.cpp @@ -11,6 +11,7 @@ #include "Qualify.h" #include "Scope.h" #include "Simplify.h" +#include "Substitute.h" #include #include @@ -144,6 +145,22 @@ size_t find_fused_group_index(const Function &producing_func, return iter - fused_groups.begin(); } +Box substitute_in_box(const map &replacements, const Box &b) { + Box result = b; + for (size_t i = 0; i < result.size(); i++) { + if (result[i].min.defined()) { + result[i].min = substitute(replacements, result[i].min); + } + if (result[i].max.defined()) { + result[i].max = substitute(replacements, result[i].max); + } + } + if (result.used.defined()) { + result.used = substitute(replacements, result.used); + } + return result; +} + // Determine if the current producing stage is fused with other // stage (i.e. the consumer stage) at dimension 'var'. bool is_fused_with_others(const vector> &fused_groups, @@ -213,12 +230,27 @@ class BoundsInference : public IRMutator { } }; + // The region required of one Func by one stage of one of its consumers, + // in two equivalent-but-not-interchangeable forms. + struct RequiredRegion { + // Phrased in terms of the consumer stage's own bound variables + // (c.sK.x.min/.max). Correct at any loop level. + Box per_stage; + + // The same region with the consumer's always-pure dimensions phrased + // in terms of its *last* stage's bound variables instead. Only + // correct at loop levels where every stage of the consumer has its + // bounds defined, because that's what makes the two alias each other + // (see define_bounds). Empty when the rewrite would change nothing. + Box canonical; + }; + struct Stage { Function func; size_t stage; // 0 is the pure definition, 1 is the first update string name; vector consumers; - map, Box> bounds; + map, RequiredRegion> bounds; vector exprs; set rvars; string stage_prefix; @@ -231,6 +263,12 @@ class BoundsInference : public IRMutator { // computed once by compute_always_pure_dims() below. vector always_pure_dims; + // The last stage's bound variables for each always-pure dimension. + // These Exprs are built once and shared by every Stage of the same + // Func, so that when they end up in several stages' boxes, + // merge_boxes' same_as fast path fires and the merge is free. + vector last_stage_min, last_stage_max; + // Computed expressions on the left and right-hand sides. // Note that a function definition might have different LHS or reduction domain // (if it's an update def) or RHS per specialization. All specializations @@ -399,6 +437,36 @@ class BoundsInference : public IRMutator { } } } + + const string last_stage = + func.name() + ".s" + std::to_string(func.updates().size()) + "."; + last_stage_min.resize(func_args.size()); + last_stage_max.resize(func_args.size()); + for (size_t j = 0; j < func_args.size(); j++) { + if (always_pure_dims[j]) { + last_stage_min[j] = Variable::make(Int(32), last_stage + func_args[j] + ".min"); + last_stage_max[j] = Variable::make(Int(32), last_stage + func_args[j] + ".max"); + } + } + } + + // The substitution that rewrites this stage's bound variables for the + // always-pure dimensions into those of the Func's last stage. Empty if + // this *is* the last stage, or if no dimension is always pure. + map pure_dim_aliases() const { + map aliases; + if (stage >= func.updates().size()) { + return aliases; + } + const vector func_args = func.args(); + const string prefix = name + ".s" + std::to_string(stage) + "."; + for (size_t j = 0; j < func_args.size(); j++) { + if (always_pure_dims[j]) { + aliases[prefix + func_args[j] + ".min"] = last_stage_min[j]; + aliases[prefix + func_args[j] + ".max"] = last_stage_max[j]; + } + } + return aliases; } // Wrap a statement in let stmts defining the box @@ -422,16 +490,38 @@ class BoundsInference : public IRMutator { size_t last_dot = loop_level.rfind('.'); string var = loop_level.substr(last_dot + 1); - for (const pair, Box> &i : bounds) { + // Whether the stage we're producing has any fused partners at all. + // If it doesn't, no consumer can be fused with it, and we can skip + // the per-entry test below. + const bool maybe_fused = + producing_func.get_contents().defined() && + !producing_func.has_extern_definition() && + !fused_pairs_in_groups[find_fused_group_index(producing_func, fused_groups)].empty(); + + for (const pair, RequiredRegion> &i : bounds) { string func_name = i.first.first; int func_stage_index = i.first.second; string stage_name = func_name + ".s" + std::to_string(func_stage_index); - if (stage_name == producing_stage_index || - inner_productions.count(func_name) || - is_fused_with_others(fused_groups, fused_pairs_in_groups, - producing_func, producing_stage_index_index, - func_name, func_stage_index, var)) { - merge_boxes(b, i.second); + + // A consumer stage that owns the loop nest we're in, or that + // shares it by fusion, has had only its *own* bound variables + // narrowed to the current iteration, so we have to use its + // per-stage box. Any other relevant consumer is produced + // further in, which means every one of its stages has its + // bounds defined at this loop level, and the always-pure + // dimensions of all of them alias the last stage's. There the + // canonical box says the same thing in fewer terms. + const bool owns_loop_nest = + stage_name == producing_stage_index || + (maybe_fused && + is_fused_with_others(fused_groups, fused_pairs_in_groups, + producing_func, producing_stage_index_index, + func_name, func_stage_index, var)); + + if (owns_loop_nest) { + merge_boxes(b, i.second.per_stage); + } else if (inner_productions.count(func_name)) { + merge_boxes(b, i.second.canonical.empty() ? i.second.per_stage : i.second.canonical); } } @@ -785,17 +875,14 @@ class BoundsInference : public IRMutator { // We need to take into account specializations which may refer to // different reduction variables as well. void populate_scope(Scope &result) { - const vector func_args = func.args(); - for (size_t i = 0; i < func_args.size(); i++) { - const string &farg = func_args[i]; - // If a dimension is pure in every stage, all stages share a - // single bound (that of the last stage), so key the required - // region off the last stage's bound variables. This keeps the - // per-stage boxes structurally identical in this dimension, so - // they collapse when merged rather than growing one term per - // stage. - size_t bound_stage = always_pure_dims[i] ? func.updates().size() : stage; - string arg = name + ".s" + std::to_string(bound_stage) + "." + farg; + for (const string &farg : func.args()) { + // Must be this stage's own bound variables: they're the ones + // narrowed to the current iteration as we descend this stage's + // loop nest. The bounds of a dimension that is pure in every + // stage still collapse across stages on merge, but that + // happens in define_bounds, which knows what loop level it is + // at. See RequiredRegion. + string arg = name + ".s" + std::to_string(stage) + "." + farg; result.push(farg, Interval(Variable::make(Int(32), arg + ".min"), Variable::make(Int(32), arg + ".max"))); @@ -948,6 +1035,22 @@ class BoundsInference : public IRMutator { } } + // A dimension that is pure in every stage of this consumer has a + // single bound shared by all of them, so phrase those dimensions + // in terms of the last stage's bound variables as well. Doing it + // once here, per consumer stage, means every stage's box is + // structurally identical in those dimensions (and shares the very + // same Exprs), so they collapse when merged instead of growing one + // term per stage. Only define_bounds knows whether this form is + // usable at a given loop level, so keep both. + const map pure_aliases = consumer.pure_dim_aliases(); + map canonical_boxes; + if (!pure_aliases.empty()) { + for (const auto &p : boxes) { + canonical_boxes.emplace(p.first, substitute_in_box(pure_aliases, p.second)); + } + } + // 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++) { @@ -984,7 +1087,11 @@ class BoundsInference : public IRMutator { debug(0) << "\n"; */ - producer.bounds[{consumer.name, consumer.stage}] = b; + RequiredRegion &required = producer.bounds[{consumer.name, consumer.stage}]; + required.per_stage = b; + if (!canonical_boxes.empty()) { + required.canonical = canonical_boxes[producer.func.name()]; + } producer.consumers.push_back((int)i); } } @@ -1020,7 +1127,7 @@ class BoundsInference : public IRMutator { if (!s.func.same_as(output)) { continue; } - s.bounds[{s.name, s.stage}] = output_box; + s.bounds[{s.name, s.stage}].per_stage = output_box; } } }