From e291904aac20d42c8f63a22c3ea70222aa88c950 Mon Sep 17 00:00:00 2001 From: Alex Reinking Date: Tue, 25 Aug 2026 14:12:49 -0400 Subject: [PATCH] Split multi-term reductions in hoist_invariants --- .../halide/src/halide_/PyStage.cpp | 4 +- src/Func.cpp | 129 +++++++++++------- src/Func.h | 26 ++-- test/correctness/hoist_invariants.cpp | 113 ++++++++++++++- 4 files changed, 209 insertions(+), 63 deletions(-) diff --git a/python_bindings/halide/src/halide_/PyStage.cpp b/python_bindings/halide/src/halide_/PyStage.cpp index 492e53fdd6db..e81a5d28afe6 100644 --- a/python_bindings/halide/src/halide_/PyStage.cpp +++ b/python_bindings/halide/src/halide_/PyStage.cpp @@ -23,7 +23,9 @@ void define_stage(py::module &m) { py::arg("preserved")) .def("rfactor", static_cast(&Stage::rfactor), py::arg("r"), py::arg("v")) - .def("hoist_invariants", &Stage::hoist_invariants) + .def("hoist_invariants", [](Stage &stage) { + return std::vector(stage.hoist_invariants()); + }) .def("eager_inline", (Stage & (Stage::*)(const std::vector &)) & Stage::eager_inline, py::arg("fs")) .def("eager_inline", [](Stage &stage, const py::args &args) -> Stage & { diff --git a/src/Func.cpp b/src/Func.cpp index 7ecd7ed98ee1..5fdf7a5e7f98 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -842,12 +842,16 @@ struct HoistedFactor { // job of the separate change_type() directive. }; +struct HoistedTerm { + optional factor; + Expr body; + size_t intermediate_index; +}; + // Given the non-self-reference increment from an update body and the // distributive law of the outer associative op, extract a loop-invariant factor // that distributes over the outer op. `reduction_vars` is the set of RVar names // the factor must not reference. -// TODO: if we flatten by the outer op here we can make tuple-valued reductions -// for things like: f(r) += a * g(r) + b * h(r) optional extract_factor(const Expr &increment, const DistributiveLaw &law, const Scope<> &reduction_vars) { @@ -890,20 +894,16 @@ optional extract_factor(const Expr &increment, return HoistedFactor{law.inner_op, factor, body}; } -vector> extract_hoisted_factors(const vector &values, - const AssociativeOp &prover_result, - const string &func_name, - const Scope<> &reduction_vars) { - vector> result(values.size()); +vector> extract_hoisted_terms(const vector &values, + const AssociativeOp &prover_result, + const string &func_name, + const Scope<> &reduction_vars) { + vector> result(values.size()); - auto is_orig_self_ref = [&](const Expr &e) { + auto is_orig_self_ref = [&](const Expr &e, size_t value_index) { const Call *c = e.as(); - return c && c->name == func_name && c->call_type == Call::Halide; - }; - - auto extract_increment = [&](const Expr &val, const DistributiveLaw &law) -> optional { - optional> split = select_binary_operand(val, law.outer_op, is_orig_self_ref); - return split ? std::make_optional(split->second) : std::nullopt; + return c && c->name == func_name && c->call_type == Call::Halide && + c->value_index == (int)value_index; }; for (size_t i = 0; i < values.size(); ++i) { @@ -912,8 +912,29 @@ vector> extract_hoisted_factors(const vector &valu // update introduced promise_clamped bindings for a preserved RVar). // Inline them so the outer op is visible to the pattern match. Expr value = substitute_in_all_lets(values[i]); - if (optional increment = extract_increment(value, *law)) { - result[i] = extract_factor(*increment, *law, reduction_vars); + vector outer_leaves; + flatten_associative_chain(value, law->outer_op, outer_leaves); + const auto self = std::find_if(outer_leaves.begin(), outer_leaves.end(), + [&](const Expr &e) { return is_orig_self_ref(e, i); }); + if (self != outer_leaves.end() && + std::find_if(std::next(self), outer_leaves.end(), + [&](const Expr &e) { return is_orig_self_ref(e, i); }) == outer_leaves.end()) { + for (const Expr &term : outer_leaves) { + if (!is_orig_self_ref(term, i)) { + optional factor = extract_factor(term, *law, reduction_vars); + result[i].push_back({factor, factor ? factor->inner_body : term, 0}); + } + } + } + } else { + // A tuple component without a distributive law must still be carried + // through an intermediate if another component is being hoisted. + Expr value = substitute_in_all_lets(values[i]); + const IRNodeType outer_op = prover_result.pattern.ops[i].node_type(); + optional> split = select_binary_operand( + value, outer_op, [&](const Expr &e) { return is_orig_self_ref(e, i); }); + if (split) { + result[i].push_back({std::nullopt, split->second, 0}); } } } @@ -1242,7 +1263,7 @@ Func Stage::rfactor(const vector> &preserved) { return intm; } -Func Stage::hoist_invariants() { +FuncVec Stage::hoist_invariants() { user_assert(!definition.is_init()) << "hoist_invariants() must be called on an update definition\n"; definition.schedule().touched() = true; @@ -1260,35 +1281,46 @@ Func Stage::hoist_invariants() { reduction_vars.push(var); reduction_bounds.push(var, Interval{min, min + extent - 1}); } - vector> hoisted_factors = - extract_hoisted_factors(definition.values(), prover_result, - function.name(), reduction_vars); - const bool any_hoisted = std::any_of(hoisted_factors.begin(), hoisted_factors.end(), - [](const auto &f) { return f.has_value(); }); - user_assert(any_hoisted) - << "hoist_invariants() could not find a distributable loop-invariant " - << "factor in the update definition of " << function.name() << ".\n"; + vector> hoisted_terms = + extract_hoisted_terms(definition.values(), prover_result, + function.name(), reduction_vars); + const bool any_hoisted = std::any_of(hoisted_terms.begin(), hoisted_terms.end(), + [](const auto &terms) { + return std::any_of(terms.begin(), terms.end(), + [](const HoistedTerm &term) { + return term.factor.has_value(); + }); + }); + const bool any_split = std::any_of(hoisted_terms.begin(), hoisted_terms.end(), + [](const auto &terms) { return terms.size() > 1; }); + user_assert(any_hoisted || any_split) + << "hoist_invariants() could not find multiple reduction terms or a " + << "distributable loop-invariant factor in the update definition of " + << function.name() << ".\n"; + + size_t intermediate_count = 0; + for (auto &terms : hoisted_terms) { + for (HoistedTerm &term : terms) { + term.intermediate_index = intermediate_count++; + } + } + FuncVec intms(function.name() + "_intm", intermediate_count); - Func intm(function.name() + "_intm"); - intm(dim_vars_exprs) = Tuple(prover_result.pattern.identities); + // Define one factor-free scalar intermediate reduction per outer term. + for (size_t i = 0; i < hoisted_terms.size(); ++i) { + for (const HoistedTerm &term : hoisted_terms[i]) { + Func &intm = intms[term.intermediate_index]; + intm(dim_vars_exprs) = prover_result.pattern.identities[i]; - // Define the factor-free intermediate reduction. - { - vector values = definition.values(); - for (size_t i = 0; i < values.size(); ++i) { - if (hoisted_factors[i]) { - Expr self_ref = Call::make(hoisted_factors[i]->inner_body.type(), function.name(), - dim_vars_exprs, Call::Halide, FunctionPtr(), (int)i); - values[i] = make_binary_op(prover_result.pattern.ops[i].node_type(), - self_ref, hoisted_factors[i]->inner_body); - } - } - values = substitute_self_reference(values, function.name(), intm.function(), {}); + Expr self_ref = Call::make(term.body.type(), intm.name(), dim_vars_exprs, + Call::Halide, FunctionPtr()); + Expr value = make_binary_op(prover_result.pattern.ops[i].node_type(), self_ref, term.body); - // The args and values still refer to the original RDom, so define_update() - // discovers and reuses it. The entire update schedule transfers unchanged. - intm.function().define_update(definition.args(), values); - intm.function().update(0).schedule() = definition.schedule().get_copy(); + // The args and value still refer to the original RDom, so define_update() + // discovers and reuses it. The entire update schedule transfers unchanged. + intm.function().define_update(definition.args(), {value}); + intm.function().update(0).schedule() = definition.schedule().get_copy(); + } } // Replace the original reduction with a factor-applying write-back update. @@ -1296,8 +1328,13 @@ Func Stage::hoist_invariants() { SubstitutionMap writeback_map; for (size_t i = 0; i < definition.values().size(); ++i) { if (!prover_result.ys[i].var.empty()) { - Expr r = (definition.values().size() == 1) ? Expr(intm(dim_vars_exprs)) : Expr(intm(dim_vars_exprs)[i]); - r = apply_hoisted_factor(r, hoisted_factors[i]); + Expr r; + for (const HoistedTerm &term : hoisted_terms[i]) { + Expr term_result = intms[term.intermediate_index](dim_vars_exprs); + term_result = apply_hoisted_factor(term_result, term.factor); + r = r.defined() ? make_binary_op(prover_result.pattern.ops[i].node_type(), r, term_result) : term_result; + } + internal_assert(r.defined()); add_let(writeback_map, prover_result.ys[i].var, r); } @@ -1337,7 +1374,7 @@ Func Stage::hoist_invariants() { definition.schedule().splits() = var_splits; } - return intm; + return intms; } void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, bool exact, TailStrategy tail) { diff --git a/src/Func.h b/src/Func.h index 583efbce14f1..3cd265113e62 100644 --- a/src/Func.h +++ b/src/Func.h @@ -59,6 +59,7 @@ struct VarOrRVar { }; class ImageParam; +class FuncVec; namespace Internal { struct AssociativeOp; @@ -214,12 +215,13 @@ class Stage { eager_inline(const Func &first, Args &&...args); // @} - /** Hoist a loop-invariant factor out of an associative reduction by applying + /** Hoist loop-invariant factors out of an associative reduction by applying * the distributive law of a semiring. Like rfactor(), this must be called on * an update definition; it splits the update into an intermediate that - * accumulates the factor-free reduction over all of the update's RVars and a - * write-back that applies the hoisted factor once. The intermediate Func is - * returned. + * accumulates one factor-free outer term over all of the update's RVars and + * a write-back that applies the hoisted factors once. The intermediate Funcs + * are returned in a FuncVec. For tuple-valued reductions, they are ordered + * by tuple output index, then by outer-term order. * * A factor is hoistable if it does not depend on any RVar being reduced. It * may be nested at any depth of an associative/commutative chain. The valid @@ -236,21 +238,25 @@ class Stage { * For example, hoist_invariants() rewrites a pipeline like this: * \code * f(x) = 0; - * f(x) += s(x) * g(x, r); + * f(x) += a(x) * g(x, r) + b(x) * h(x, r); * \endcode * into a pipeline like this: * \code - * f_intm(x) = 0; - * f_intm(x) += g(x, r); + * f_intm0(x) = 0; + * f_intm0(x) += g(x, r); + * f_intm1(x) = 0; + * f_intm1(x) += h(x, r); * * f(x) = 0; - * f(x) += s(x) * f_intm(x); + * f(x) += a(x) * f_intm0(x) + b(x) * f_intm1(x); * \endcode * * This reduces the number of factor applications from |R| to one per pure - * point. It is an error if no distributable invariant factor is found. + * point. Terms without a hoistable factor are still split into separate + * intermediates. It is an error if there is neither a distributable invariant + * factor nor more than one outer term to split. */ - Func hoist_invariants(); + FuncVec hoist_invariants(); /** Schedule the iteration over this stage to be fused with another * stage 's' from outermost loop to a given LoopLevel. 'this' stage will diff --git a/test/correctness/hoist_invariants.cpp b/test/correctness/hoist_invariants.cpp index 73c4e5e38667..7f61562ca580 100644 --- a/test/correctness/hoist_invariants.cpp +++ b/test/correctness/hoist_invariants.cpp @@ -182,6 +182,104 @@ int hoist_invariants_scattered_factors_unsigned_test() { return 0; } +// A sum reduction can contain independently scaled terms. Each outer addend +// gets its own scalar intermediate, and the write-back applies the scale while +// merging those intermediates into the original accumulator. +int hoist_invariants_multiple_terms_test() { + const int K = 16; + RDom r(0, K, "r"); + + Func f{"multiple_terms"}; + f() = 0; + f() += 2 * (r + 1) + 3 * (2 * r + 1); + + FuncVec intms = f.update().hoist_invariants(); + internal_assert(intms.size() == 2) + << "hoist_invariants multiple terms: expected two intermediates, got " + << intms.size() << "\n"; + internal_assert(intms[0].name() == f.name() + "_intm0" && + intms[1].name() == f.name() + "_intm1") + << "hoist_invariants multiple terms: unexpected intermediate names\n"; + for (Func &intm : intms) { + intm.compute_root(); + } + + Buffer result = f.realize(); + int expected = 0; + for (int k = 0; k < K; ++k) { + expected += 2 * (k + 1) + 3 * (2 * k + 1); + } + internal_assert(result() == expected) + << "hoist_invariants multiple terms: got " << result() + << ", expected " << expected << "\n"; + return 0; +} + +// Splitting outer terms is useful even when none has an invariant factor. +// Each term is reduced independently and merged at the original accumulator. +int hoist_invariants_unscaled_terms_test() { + const int K = 16; + Var x{"x"}; + RDom r(0, K, "r"); + + Func g{"unscaled_g"}, h{"unscaled_h"}, f{"unscaled_terms"}; + g(x) = x + 1; + h(x) = x * x + 3; + f() = 0; + f() += g(r) + h(r); + + FuncVec intms = f.update().hoist_invariants(); + internal_assert(intms.size() == 2) + << "hoist_invariants unscaled terms: expected two intermediates, got " + << intms.size() << "\n"; + for (Func &intm : intms) { + intm.compute_root(); + } + + Buffer result = f.realize(); + int expected = 0; + for (int k = 0; k < K; ++k) { + expected += (k + 1) + (k * k + 3); + } + internal_assert(result() == expected) + << "hoist_invariants unscaled terms: got " << result() + << ", expected " << expected << "\n"; + return 0; +} + +// Tuple outputs are flattened independently. The FuncVec is ordered by tuple +// output index, then by the order of that output's flattened outer terms. +int hoist_invariants_tuple_terms_test() { + const int K = 8; + RDom r(0, K, "r"); + + Func f{"tuple_terms"}; + f() = Tuple(0, 0); + f() = Tuple(f()[0] + 2 * (r + 1) + 3 * (r + 2), + f()[1] + 4 * (r + 3)); + + FuncVec intms = f.update().hoist_invariants(); + internal_assert(intms.size() == 3) + << "hoist_invariants tuple terms: expected three intermediates, got " + << intms.size() << "\n"; + for (size_t i = 0; i < intms.size(); ++i) { + internal_assert(intms[i].name() == f.name() + "_intm" + std::to_string(i)) + << "hoist_invariants tuple terms: unexpected intermediate ordering\n"; + intms[i].compute_root(); + } + + Realization result = f.realize(); + int expected0 = 0, expected1 = 0; + for (int k = 0; k < K; ++k) { + expected0 += 2 * (k + 1) + 3 * (k + 2); + expected1 += 4 * (k + 3); + } + internal_assert(result[0].as()() == expected0 && + result[1].as()() == expected1) + << "hoist_invariants tuple terms: incorrect result\n"; + return 0; +} + // hoist_invariants() with outer Min and additive factor: // min_k(offset(i) + body(i, k)) = offset(i) + min_k(body(i, k)) // The intermediate accumulates min without the offset; write-back adds it once. @@ -439,8 +537,8 @@ int hoist_invariants_invalid_law_rejected_test() { } catch (const Halide::CompileError &e) { error = true; const string expected = - "hoist_invariants() could not find a distributable loop-invariant " - "factor in the update definition of " + + "hoist_invariants() could not find multiple reduction terms or a " + "distributable loop-invariant factor in the update definition of " + f.name() + "."; if (string(e.what()).find(expected) == string::npos) { printf("Unexpected error for unsigned min hoisting:\n%s\n", e.what()); @@ -454,8 +552,8 @@ int hoist_invariants_invalid_law_rejected_test() { return 0; } -// hoist_invariants() errors when there is no distributable invariant factor to -// hoist, rather than silently behaving like a plain rfactor(). +// hoist_invariants() errors when there is neither an invariant factor to hoist +// nor multiple outer terms to split, rather than behaving like plain rfactor(). int hoist_invariants_nothing_to_hoist_rejected_test() { if (!Halide::exceptions_enabled()) { return 0; @@ -475,8 +573,8 @@ int hoist_invariants_nothing_to_hoist_rejected_test() { } catch (const Halide::CompileError &e) { error = true; const string expected = - "hoist_invariants() could not find a distributable loop-invariant " - "factor in the update definition of " + + "hoist_invariants() could not find multiple reduction terms or a " + "distributable loop-invariant factor in the update definition of " + f.name() + "."; if (string(e.what()).find(expected) == string::npos) { printf("Unexpected error when no invariant is hoistable:\n%s\n", e.what()); @@ -502,6 +600,9 @@ int main(int argc, char **argv) { {"hoist_invariants test (add/mul)", hoist_invariants_test}, {"hoist_invariants test (add/mul, scattered factors)", hoist_invariants_scattered_factors_test}, {"hoist_invariants test (add/mul, scattered factors, unsigned)", hoist_invariants_scattered_factors_unsigned_test}, + {"hoist_invariants test (multiple terms)", hoist_invariants_multiple_terms_test}, + {"hoist_invariants test (unscaled terms)", hoist_invariants_unscaled_terms_test}, + {"hoist_invariants test (tuple terms)", hoist_invariants_tuple_terms_test}, {"hoist_invariants test (min/add)", hoist_invariants_min_test}, {"hoist_invariants test (or/and)", hoist_invariants_or_test}, {"hoist_invariants test (strict_float preserved)", hoist_invariants_strict_float_test},