diff --git a/src/BoundsInference.cpp b/src/BoundsInference.cpp index 73554dbaa1e2..563480de7d47 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,18 +230,45 @@ 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; 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; + + // 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 @@ -377,6 +421,54 @@ 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; + } + } + } + + 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 Stmt define_bounds(Stmt s, const Function &producing_func, @@ -398,37 +490,51 @@ 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); } } 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) + "."; @@ -771,6 +877,12 @@ class BoundsInference : public IRMutator { // different reduction variables as well. void populate_scope(Scope &result) { 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"), @@ -843,6 +955,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; @@ -923,6 +1036,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++) { @@ -959,7 +1088,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); } } @@ -995,7 +1128,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; } } }