diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index c0f68be20fb3..507e7e0fd56d 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -60,6 +60,10 @@ void define_enums(py::module &m) { .value("GuardWithIf", PrefetchBoundStrategy::GuardWithIf) .value("NonFaulting", PrefetchBoundStrategy::NonFaulting); + py::enum_(m, "RFactorOptions") + .value("None", RFactorOptions::None) + .value("HoistInvariantFactor", RFactorOptions::HoistInvariantFactor); + py::enum_(m, "StmtOutputFormat") .value("Text", StmtOutputFormat::Text) .value("HTML", StmtOutputFormat::HTML); diff --git a/python_bindings/src/halide/halide_/PyStage.cpp b/python_bindings/src/halide/halide_/PyStage.cpp index 9a9f6394d7cd..d9fd72b3b5ab 100644 --- a/python_bindings/src/halide/halide_/PyStage.cpp +++ b/python_bindings/src/halide/halide_/PyStage.cpp @@ -14,10 +14,10 @@ void define_stage(py::module &m) { .def("dump_argument_list", &Stage::dump_argument_list) .def("name", &Stage::name) - .def("rfactor", static_cast> &)>(&Stage::rfactor), - py::arg("preserved")) - .def("rfactor", static_cast(&Stage::rfactor), - py::arg("r"), py::arg("v")) + .def("rfactor", static_cast> &, RFactorOptions)>(&Stage::rfactor), + py::arg("preserved"), py::arg("options") = RFactorOptions::None) + .def("rfactor", static_cast(&Stage::rfactor), + py::arg("r"), py::arg("v"), py::arg("options") = RFactorOptions::None) .def("split_vars", [](const Stage &stage) -> py::list { auto vars = stage.split_vars(); diff --git a/src/AssociativeOpsTable.cpp b/src/AssociativeOpsTable.cpp index 63b4add503e9..7ab139f06f81 100644 --- a/src/AssociativeOpsTable.cpp +++ b/src/AssociativeOpsTable.cpp @@ -109,6 +109,11 @@ struct TableKey { map> pattern_tables; +std::mutex &ops_table_lock() { + static std::mutex lock; + return lock; +} + #define declare_vars(t, index) \ Expr x##index = Variable::make((t), "x" + std::to_string(index)); \ Expr y##index = Variable::make((t), "y" + std::to_string(index)); \ @@ -362,8 +367,7 @@ const vector &get_ops_table(const vector &exprs) { const vector &table = [&]() -> decltype(auto) { // get_ops_table_helper() lazily initializes the table, so ensure // that multiple threads can't try to do so at the same time. - static std::mutex ops_table_lock; - std::scoped_lock lock_guard(ops_table_lock); + std::scoped_lock lock_guard(ops_table_lock()); return get_ops_table_helper(types, exprs[0].node_type(), exprs.size()); }(); @@ -376,5 +380,23 @@ const vector &get_ops_table(const vector &exprs) { return table; } +std::optional get_associative_identity(Type type, IRNodeType root) { + std::scoped_lock lock_guard(ops_table_lock()); + + const vector &table = get_ops_table_helper({type}, root, 1); + if (table.empty()) { + return std::nullopt; + } + + const Expr &identity = table.front().identities.front(); + for (const AssociativePattern &pattern : table) { + internal_assert(pattern.size() == 1); + if (!equal(pattern.identities.front(), identity)) { + return std::nullopt; + } + } + return identity; +} + } // namespace Internal } // namespace Halide diff --git a/src/AssociativeOpsTable.h b/src/AssociativeOpsTable.h index 9bbb12db7536..ab3568c951d3 100644 --- a/src/AssociativeOpsTable.h +++ b/src/AssociativeOpsTable.h @@ -8,6 +8,7 @@ #include "IREquality.h" #include "IROperator.h" +#include #include #include @@ -71,6 +72,10 @@ struct AssociativePattern { const std::vector &get_ops_table(const std::vector &exprs); +/** Return the identity for a single-output associative op, if the table has one + * and all matching patterns agree on it. */ +std::optional get_associative_identity(Type type, IRNodeType root); + } // namespace Internal } // namespace Halide diff --git a/src/Func.cpp b/src/Func.cpp index f4e4298fae07..f79c5d6ec855 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -21,6 +21,7 @@ #include "Function.h" #include "IR.h" #include "IREquality.h" +#include "IRMatch.h" #include "IRMutator.h" #include "IROperator.h" #include "IRPrinter.h" @@ -573,63 +574,37 @@ std::string Stage::dump_argument_list() const { return dump_dim_list(definition.schedule().dims()); } -namespace { - -class SubstituteSelfReference : public IRMutator { - using IRMutator::visit; - - const string func; - const Function substitute; - const vector new_args; - - Expr visit(const Call *c) override { - Expr expr = IRMutator::visit(c); - c = expr.as(); - internal_assert(c); - - if ((c->call_type == Call::Halide) && (func == c->name)) { - debug(4) << "...Replace call to Func \"" << c->name << "\" with " - << "\"" << substitute.name() << "\"\n"; - vector args; - args.insert(args.end(), c->args.begin(), c->args.end()); - args.insert(args.end(), new_args.begin(), new_args.end()); - expr = Call::make(substitute, args, c->value_index); - } - return expr; - } - -public: - SubstituteSelfReference(const string &func, const Function &substitute, - const vector &new_args) - : func(func), substitute(substitute), new_args(new_args) { - internal_assert(substitute.get_contents().defined()); - } -}; - -/** Substitute all self-reference calls to 'func' with 'substitute' which - * args (LHS) is the old args (LHS) plus 'new_args' in that order. - * Expect this method to be called on the value (RHS) of an update definition. */ -vector substitute_self_reference(const vector &values, const string &func, - const Function &substitute, const vector &new_args) { - SubstituteSelfReference subs(func, substitute, new_args); - vector result; - result.reserve(values.size()); - for (const auto &val : values) { - result.push_back(subs(val)); - } - return result; -} - -} // anonymous namespace - -Func Stage::rfactor(const RVar &r, const Var &v) { +Func Stage::rfactor(const RVar &r, const Var &v, RFactorOptions options) { definition.schedule().touched() = true; - return rfactor({{r, v}}); + return rfactor({{r, v}}, options); } // Helpers for rfactor implementation namespace { +// Replace self-references to `func_name` with calls to the intermediate `intm`, +// appending the preserved vars to the call's args. Expected to be called on the +// values (RHS) of an update definition. +vector substitute_self_reference(vector values, + const string &func_name, + const Function &intm, + const vector &preserved_vars) { + for (Expr &v : values) { + v = mutate_with(v, [&](auto *self, const Call *c) -> Expr { + Expr expr = self->visit_base(c); + c = expr.as(); + internal_assert(c); + if (c->call_type == Call::Halide && func_name == c->name) { + vector args(c->args); + args.insert(args.end(), preserved_vars.begin(), preserved_vars.end()); + expr = Call::make(intm, args, c->value_index); + } + return expr; + }); + } + return values; +} + optional find_dim(const vector &items, const VarOrRVar &v) { const auto has_v = std::find_if(items.begin(), items.end(), [&](auto &x) { return dim_match(x, v); @@ -682,17 +657,27 @@ string dequalify(string name) { return name; } -vector subst_dims(const SubstitutionMap &substitution_map, const vector &dims) { - auto new_dims = dims; - for (auto &dim : new_dims) { - if (const auto it = substitution_map.find(dim.var); it != substitution_map.end()) { - const Variable *new_var = it->second.as(); - internal_assert(new_var); - dim.var = new_var->name; +struct RFactorProjection { + const SubstitutionMap &rdom_promises; + const SubstitutionMap &vars; + + template + T operator()(const T &x) const { + return substitute(vars, substitute(rdom_promises, x)); + } + + vector operator()(const vector &dims) const { + auto new_dims = dims; + for (auto &dim : new_dims) { + if (const auto it = vars.find(dim.var); it != vars.end()) { + const Variable *new_var = it->second.as(); + internal_assert(new_var); + dim.var = new_var->name; + } } + return new_dims; } - return new_dims; -} +}; pair project_rdom(const vector &dims, const ReductionDomain &rdom, const vector &splits) { // The bounds projections maps expressions that reference the old RDom @@ -758,6 +743,232 @@ pair project_rdom(const vector &dims, con return {new_rdom, dim_projection}; } +// A semiring distributive law used by hoisted rfactor: `inner` distributes over +// the outer (reduction) op, so a loop-invariant operand of `inner` can be +// hoisted out of the reduction. +struct DistributiveLaw { + IRNodeType outer_op; + IRNodeType inner_op; + // The +/* ring auto-promotes an integer body to the factor's float type; + // when true, we strip that implicit promotion so the intermediate keeps + // the integer accumulation type. + bool strip_promotion_cast; +}; + +constexpr DistributiveLaw distributive_laws[] = { + {IRNodeType::Add, IRNodeType::Mul, true}, // sum_k(s * x_k) = s * sum_k(x_k) + {IRNodeType::Min, IRNodeType::Add, true}, // min_k(c + x_k) = c + min_k(x_k) + {IRNodeType::Max, IRNodeType::Add, true}, // max_k(c + x_k) = c + max_k(x_k) + {IRNodeType::Or, IRNodeType::And, false}, // or_k(p && x_k) = p && or_k(x_k) + {IRNodeType::And, IRNodeType::Or, false}, // and_k(p || x_k) = p || and_k(x_k) +}; + +bool distributive_law_valid_for_type(const DistributiveLaw &law, Type t) { + if (law.outer_op == IRNodeType::Min || law.outer_op == IRNodeType::Max) { + // Hoisting min/max over addition relies on addition being order-preserving. + // This is not true for unsigned or narrow signed integer wraparound. + return !t.can_overflow(); + } + return true; +} + +// nullopt if no law's outer_op matches, or if the matching law is invalid for `op`'s type. +optional distributive_law_for(const Expr &op) { + for (const DistributiveLaw &law : distributive_laws) { + if (law.outer_op == op.node_type()) { + return distributive_law_valid_for_type(law, op.type()) ? std::make_optional(law) : std::nullopt; + } + } + return std::nullopt; +} + +// If `e` is a binary op of node type `op` and exactly one of its two operands +// satisfies `is_selected`, returns {selected operand, other operand}. Returns +// nullopt if `e` isn't a binary `op`, or if neither/both operands match. +template +optional> select_binary_operand(const Expr &e, IRNodeType op, Predicate &&is_selected) { + if (e.node_type() != op) { + return std::nullopt; + } + // `op` is always a binary op, so this is guaranteed to have a value. + auto [a, b] = *as_binary_operands(e); + const bool a_sel = is_selected(a); + const bool b_sel = is_selected(b); + if (a_sel == b_sel) { + return std::nullopt; + } + return a_sel ? std::make_pair(a, b) : std::make_pair(b, a); +} + +// Collect the leaves of a chain of `op`-typed binary nodes. +// E.g., flatten (a*b)*(c*d) into [a, b, c, d] +void flatten_associative_chain(const Expr &e, IRNodeType op, vector &leaves) { + if (e.node_type() == op) { + auto [a, b] = *as_binary_operands(e); + flatten_associative_chain(a, op, leaves); + flatten_associative_chain(b, op, leaves); + } else { + leaves.push_back(e); + } +} + +struct HoistedFactor { + IRNodeType op; + Expr factor; // The loop-invariant distributable factor, expressed in + // the preserved update's coordinate system. + Expr inner_body; // The remaining body after removing the factor, with any + // implicit float-promotion cast stripped (for the +/* ring). +}; + +// 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. `intermediate_vars` is the set of RVar +// names the factor must NOT reference after projecting into the intermediate +// update (the rvars the intermediate reduces over). +optional extract_factor(const Expr &increment, + const DistributiveLaw &law, + const Scope<> &intermediate_vars, + const RFactorProjection &to_intermediate, + const RFactorProjection &to_preserved) { + auto is_intermediate_rvar_free = [&](const Expr &e) { + return !expr_uses_vars(to_intermediate(e), intermediate_vars); + }; + + // An invariant factor may be nested arbitrarily deep in an + // associative/commutative chain, so flatten the whole chain + // into leaves and partition by invariance. This is just + // commutative-ring algebra: l1*l2*...*lN can always be regrouped + // as (product of invariant leaves) * (product of dependent leaves), + // regardless of how the multiplication was parenthesized. + vector leaves; + flatten_associative_chain(increment, law.inner_op, leaves); + + vector invariant_leaves, dependent_leaves; + for (const Expr &leaf : leaves) { + if (is_intermediate_rvar_free(leaf)) { + invariant_leaves.push_back(to_preserved(leaf)); + } else { + dependent_leaves.push_back(leaf); + } + } + if (invariant_leaves.empty() || dependent_leaves.empty()) { + // Nothing to hoist, or the entire increment is invariant (a + // degenerate case not worth special-casing here). + return std::nullopt; + } + + Expr factor; + for (const Expr &leaf : invariant_leaves) { + factor = factor.defined() ? make_binary_op(law.inner_op, factor, leaf) : leaf; + } + Expr body; + for (const Expr &leaf : dependent_leaves) { + body = body.defined() ? make_binary_op(law.inner_op, body, leaf) : leaf; + } + + if (law.strip_promotion_cast) { + // Match some lossless patterns that expose a single cast: + Expr a_i8 = Variable::make(Int(8), "a"); + Expr b_i8 = Variable::make(Int(8), "b"); + Expr a_u8 = Variable::make(UInt(8), "a"); + Expr b_u8 = Variable::make(UInt(8), "b"); + + std::map matches; + static const std::pair patterns[] = { + {cast(Float(32), a_i8) * cast(Float(32), b_i8), cast(Float(32), widening_mul(a_i8, b_i8))}, + {cast(Float(32), a_i8) + cast(Float(32), b_i8), cast(Float(32), widening_add(a_i8, b_i8))}, + {cast(Float(32), a_i8) - cast(Float(32), b_i8), cast(Float(32), widening_sub(a_i8, b_i8))}, + {cast(Float(32), a_u8) * cast(Float(32), b_u8), cast(Float(32), widening_mul(a_u8, b_u8))}, + {cast(Float(32), a_u8) + cast(Float(32), b_u8), cast(Float(32), widening_add(a_u8, b_u8))}, + {cast(Float(32), a_u8) - cast(Float(32), b_u8), cast(Float(32), widening_sub(a_u8, b_u8))}, + }; + + for (const auto &[pattern, result] : patterns) { + if (expr_match(pattern, body, matches)) { + body = substitute(matches, result); + break; + } + } + + // Strip only the outer promotion cast; strict_cast inside the body is left alone. + if (const Cast *c = body.as()) { + int acc_bits = body.type().bits(); + if (body.type().is_float() && c->value.type().is_int_or_uint()) { + body = c->value; + // When strength reducing a float accumulator to an integer accumulator, + // avoid losing precision by picking a type that's at least as wide. + if (acc_bits > body.type().bits()) { + body = cast(body.type().with_bits(acc_bits), body); + } + } + } + } + + return HoistedFactor{law.inner_op, factor, body}; +} + +vector> extract_hoisted_factors(const vector &values, + const AssociativeOp &prover_result, + const string &func_name, + const Scope<> &intermediate_vars, + const RFactorProjection &to_intermediate, + const RFactorProjection &to_preserved) { + vector> result(values.size()); + + auto is_orig_self_ref = [&](const Expr &e) { + 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; + }; + + for (size_t i = 0; i < values.size(); ++i) { + if (optional law = distributive_law_for(prover_result.pattern.ops[i])) { + if (optional increment = extract_increment(values[i], *law)) { + result[i] = extract_factor(*increment, *law, intermediate_vars, + to_intermediate, to_preserved); + } + } + } + return result; +} + +// Build the identities for the intermediate's pure definition. A hoisted element +// accumulates its inner body at that body's (possibly narrower) type, so its +// identity is the outer reduction op's identity in that type; other elements +// keep the prover's identity. +Tuple build_intm_identities(const AssociativePattern &pattern, + const vector> &hoisted_factors) { + vector result = pattern.identities; + for (size_t i = 0; i < result.size(); ++i) { + if (hoisted_factors[i]) { + const HoistedFactor &factor = *hoisted_factors[i]; + const IRNodeType outer_op = pattern.ops[i].node_type(); + std::optional identity = get_associative_identity(factor.inner_body.type(), outer_op); + internal_assert(identity) + << "Could not find an unambiguous identity for " + << IRNodeType_string(outer_op) << " over " << factor.inner_body.type() << "\n"; + result[i] = *identity; + } + } + return Tuple(result); +} + +// The intermediate accumulated the inner body at its own (possibly narrower) +// type, so cast back to the outer op's type before recombining with the factor. +Expr apply_hoisted_factor(Expr r, const optional &factor, Type result_type) { + if (!factor) { + return r; + } + if (r.type() != result_type) { + r = cast(result_type, r); + } + return make_binary_op(factor->op, factor->factor, r); +} + } // namespace pair, vector> Stage::rfactor_validate_args(const std::vector> &preserved, const AssociativeOp &prover_result) { @@ -854,7 +1065,7 @@ pair, vector> Stage::rfactor_validate_args(const std::vecto return std::make_pair(std::move(var_splits), std::move(rvar_splits)); } -Func Stage::rfactor(const vector> &preserved) { +Func Stage::rfactor(const vector> &preserved, RFactorOptions options) { user_assert(!definition.is_init()) << "rfactor() must be called on an update definition\n"; definition.schedule().touched() = true; @@ -918,6 +1129,9 @@ Func Stage::rfactor(const vector> &preserved) { // Project the RDom into each side ReductionDomain intermediate_rdom, preserved_rdom; SubstitutionMap intermediate_map, preserved_map; + RFactorProjection to_intermediate{rdom_promises, intermediate_map}; + RFactorProjection to_preserved{rdom_promises, preserved_map}; + { // Intermediate std::tie(intermediate_rdom, intermediate_map) = project_rdom(intermediate_rdims, rdom, rvar_splits); @@ -926,9 +1140,7 @@ Func Stage::rfactor(const vector> &preserved) { } { - Expr pred = intermediate_rdom.predicate(); - pred = substitute(rdom_promises, pred); - pred = substitute(intermediate_map, pred); + Expr pred = to_intermediate(intermediate_rdom.predicate()); intermediate_rdom.set_predicate(simplify(pred)); } @@ -941,9 +1153,7 @@ Func Stage::rfactor(const vector> &preserved) { intm_rdom.push(var, Interval{min, min + extent - 1}); } { - Expr pred = preserved_rdom.predicate(); - pred = substitute(rdom_promises, pred); - pred = substitute(preserved_map, pred); + Expr pred = to_preserved(preserved_rdom.predicate()); pred = or_condition_over_domain(pred, intm_rdom); preserved_rdom.set_predicate(pred); } @@ -952,24 +1162,55 @@ Func Stage::rfactor(const vector> &preserved) { // Intermediate func Func intm(function.name() + "_intm"); + // Must happen before intm(args)=Tuple(...) so that output_types is set correctly + // from the start. substitute_self_reference uses intm.types() to determine the + // Call type, and changing output_types afterward has no effect. + const size_t n_vals = definition.values().size(); + vector> hoisted_factors(n_vals); + if (has_rfactor_option(options, RFactorOptions::HoistInvariantFactor)) { + // A factor is hoistable if, after projecting the original expression + // into the intermediate update, it does not depend on any RVar reduced + // by that intermediate update. It may depend on preserved RVars. + Scope<> intermediate_vars; + for (const auto &[var, min, extent] : intermediate_rdom.domain()) { + intermediate_vars.push(var); + } + hoisted_factors = extract_hoisted_factors(definition.values(), prover_result, + function.name(), intermediate_vars, + to_intermediate, to_preserved); + } + // Intermediate pure definition { vector args = dim_vars_exprs; args.insert(args.end(), preserved_vars.begin(), preserved_vars.end()); - intm(args) = Tuple(prover_result.pattern.identities); + intm(args) = build_intm_identities(prover_result.pattern, hoisted_factors); } // Intermediate update definition { vector args = definition.args(); args.insert(args.end(), preserved_vars.begin(), preserved_vars.end()); - args = substitute(rdom_promises, args); - args = substitute(intermediate_map, args); - + args = to_intermediate(args); + + // For a hoisted element, replace the value with the outer op applied to + // a self-reference and the inner body; the factor is re-applied later at + // write-back. The self-reference is a call to the original func typed at + // the intermediate's (possibly different) accumulation type, so the outer + // op does not re-promote the inner body; substitute_self_reference then + // retargets it to the intermediate. vector values = definition.values(); + for (size_t i = 0; i < n_vals; ++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(), preserved_vars); - values = substitute(rdom_promises, values); - values = substitute(intermediate_map, values); + values = to_intermediate(values); intm.function().define_update(args, values, intermediate_rdom); // Intermediate schedule @@ -1003,7 +1244,7 @@ Func Stage::rfactor(const vector> &preserved) { } intm.function().update(0).schedule() = definition.schedule().get_copy(); - intm.function().update(0).schedule().dims() = subst_dims(intermediate_map, intm_dims); + intm.function().update(0).schedule().dims() = to_intermediate(intm_dims); intm.function().update(0).schedule().rvars() = intermediate_rdom.domain(); intm.function().update(0).schedule().splits() = var_splits; } @@ -1019,11 +1260,12 @@ Func Stage::rfactor(const vector> &preserved) { 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(f_load_args)) : Expr(intm(f_load_args)[i]); + r = apply_hoisted_factor(r, hoisted_factors[i], prover_result.pattern.ops[i].type()); add_let(preserved_map, prover_result.ys[i].var, r); } if (!prover_result.xs[i].var.empty()) { - Expr prev_val = Call::make(intm.types()[i], function.name(), + Expr prev_val = Call::make(function.output_types()[i], function.name(), dim_vars_exprs, Call::CallType::Halide, FunctionPtr(), i); add_let(preserved_map, prover_result.xs[i].var, prev_val); @@ -1063,9 +1305,9 @@ Func Stage::rfactor(const vector> &preserved) { } definition.args() = dim_vars_exprs; - definition.values() = substitute(preserved_map, substitute(rdom_promises, prover_result.pattern.ops)); + definition.values() = to_preserved(prover_result.pattern.ops); definition.predicate() = preserved_rdom.predicate(); - definition.schedule().dims() = subst_dims(preserved_map, reducing_dims); + definition.schedule().dims() = to_preserved(reducing_dims); definition.schedule().rvars() = preserved_rdom.domain(); definition.schedule().splits() = var_splits; } diff --git a/src/Func.h b/src/Func.h index 0bfb591871c7..536de4fe0951 100644 --- a/src/Func.h +++ b/src/Func.h @@ -18,6 +18,7 @@ #include "Var.h" #include +#include #include namespace Halide { @@ -66,6 +67,26 @@ struct Split; struct StorageDim; } // namespace Internal +enum class RFactorOptions : uint32_t { + None = 0, + HoistInvariantFactor = 1 << 0, +}; + +inline RFactorOptions operator|(RFactorOptions a, RFactorOptions b) { + using T = std::underlying_type_t; + return static_cast(static_cast(a) | static_cast(b)); +} + +inline RFactorOptions operator&(RFactorOptions a, RFactorOptions b) { + using T = std::underlying_type_t; + return static_cast(static_cast(a) & static_cast(b)); +} + +inline bool has_rfactor_option(RFactorOptions options, RFactorOptions flag) { + using T = std::underlying_type_t; + return (static_cast(options) & static_cast(flag)) != 0; +} + /** A single definition of a Func. May be a pure or update definition. */ class Stage { /** Reference to the Function this stage (or definition) belongs to. */ @@ -143,52 +164,79 @@ class Stage { * with the new pure Vars added to the outermost. * * For example, f.update(0).rfactor({{r.y, u}}) would rewrite a pipeline like this: - \code - f(x, y) = 0; - f(x, y) += g(r.x, r.y); - \endcode + * \code + * f(x, y) = 0; + * f(x, y) += g(r.x, r.y); + * \endcode * into a pipeline like this: - \code - f_intm(x, y, u) = 0; - f_intm(x, y, u) += g(r.x, u); - - f(x, y) = 0; - f(x, y) += f_intm(x, y, r.y); - \endcode + * \code + * f_intm(x, y, u) = 0; + * f_intm(x, y, u) += g(r.x, u); + * + * f(x, y) = 0; + * f(x, y) += f_intm(x, y, r.y); + * \endcode * * This has a variety of uses. You can use it to split computation of an associative reduction: - \code - f(x, y) = 10; - RDom r(0, 96); - f(x, y) = max(f(x, y), g(x, y, r.x)); - f.update(0).split(r.x, rxo, rxi, 8).reorder(y, x).parallel(x); - f.update(0).rfactor({{rxo, u}}).compute_root().parallel(u).update(0).parallel(u); - \endcode + * \code + * f(x, y) = 10; + * RDom r(0, 96); + * f(x, y) = max(f(x, y), g(x, y, r.x)); + * f.update(0).split(r.x, rxo, rxi, 8).reorder(y, x).parallel(x); + * f.update(0).rfactor({{rxo, u}}).compute_root().parallel(u).update(0).parallel(u); + * \endcode * *, which is equivalent to: - \code - parallel for u = 0 to 11: - for y: - for x: - f_intm(x, y, u) = -inf - parallel for x: - for y: - parallel for u = 0 to 11: - for rxi = 0 to 7: - f_intm(x, y, u) = max(f_intm(x, y, u), g(8*u + rxi)) - for y: - for x: - f(x, y) = 10 - parallel for x: - for y: - for rxo = 0 to 11: - f(x, y) = max(f(x, y), f_intm(x, y, rxo)) - \endcode - * + * \code + * parallel for u = 0 to 11: + * for y: + * for x: + * f_intm(x, y, u) = -inf + * parallel for x: + * for y: + * parallel for u = 0 to 11: + * for rxi = 0 to 7: + * f_intm(x, y, u) = max(f_intm(x, y, u), g(8*u + rxi)) + * for y: + * for x: + * f(x, y) = 10 + * parallel for x: + * for y: + * for rxo = 0 to 11: + * f(x, y) = max(f(x, y), f_intm(x, y, rxo)) + * \endcode + * + * If `options` contains `RFactorOptions::HoistInvariantFactor`, rfactor applies + * the distributive law of a semiring to hoist a loop-invariant factor from the + * intermediate accumulation to the write-back step. For a factor that does not + * depend on any RVar being reduced in the intermediate, the valid hoistings are: + * + * Outer op Inner combine Law + * --------- ------------- --- + * + (sum) * sum_k(s * x_k) = s * sum_k(x_k) + * min + min_k(c + x_k) = c + min_k(x_k) + * max + max_k(c + x_k) = c + max_k(x_k) + * || (bool) && or_k(p && x_k) = p && or_k(x_k) + * && (bool) || and_k(p || x_k) = p || and_k(x_k) + * + * For the arithmetic laws (every row above except the boolean ||/&&), rfactor + * also strips the implicit float-promotion cast that Halide inserts when a + * float factor or offset combines with an integer body, allowing the + * intermediate to accumulate at the integer type (e.g. Int(32) instead of + * Float(32)). It also recognizes exact products such as + * cast(int8_a) * cast(int8_b) as + * cast(widening_mul(int8_a, int8_b)) before choosing the + * intermediate type, which can expose dot-product-friendly integer + * accumulations. Strict casts are preserved; use strict_float(cast(...)) + * when this conversion must not be deferred. + * + * This reduces the number of factor applications from |K_inner| × |K_pres| to + * |K_pres|. If no distributable factor is found, hoisted rfactor falls back + * silently to normal rfactor behavior. */ // @{ - Func rfactor(const std::vector> &preserved); - Func rfactor(const RVar &r, const Var &v); + Func rfactor(const std::vector> &preserved, RFactorOptions options = RFactorOptions::None); + Func rfactor(const RVar &r, const Var &v, RFactorOptions options = RFactorOptions::None); // @} /** Schedule the iteration over this stage to be fused with another diff --git a/src/IROperator.cpp b/src/IROperator.cpp index c109f1d2ccff..d7951b9c2870 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -246,6 +246,76 @@ std::optional is_const_power_of_two_integer(int64_t val) { return val < 0 ? std::nullopt : is_const_power_of_two_integer((uint64_t)val); } +std::optional> as_binary_operands(const Expr &e) { + // We switch on the actual node type, so we can downcast e.get() directly + // rather than going through Expr::as<>(), which would redundantly re-check + // the node type the switch case has already established. + switch (e.node_type()) { +#define HANDLE_BINARY_OP(NodeType) \ + case IRNodeType::NodeType: { \ + const NodeType *op = static_cast(e.get()); \ + return std::pair{op->a, op->b}; \ + } + HANDLE_BINARY_OP(Add) + HANDLE_BINARY_OP(Sub) + HANDLE_BINARY_OP(Mul) + HANDLE_BINARY_OP(Div) + HANDLE_BINARY_OP(Mod) + HANDLE_BINARY_OP(Min) + HANDLE_BINARY_OP(Max) + HANDLE_BINARY_OP(EQ) + HANDLE_BINARY_OP(NE) + HANDLE_BINARY_OP(LT) + HANDLE_BINARY_OP(LE) + HANDLE_BINARY_OP(GT) + HANDLE_BINARY_OP(GE) + HANDLE_BINARY_OP(And) + HANDLE_BINARY_OP(Or) +#undef HANDLE_BINARY_OP + default: + return std::nullopt; + } +} + +Expr make_binary_op(IRNodeType t, const Expr &a, const Expr &b) { + switch (t) { + case IRNodeType::Add: + return a + b; + case IRNodeType::Sub: + return a - b; + case IRNodeType::Mul: + return a * b; + case IRNodeType::Div: + return a / b; + case IRNodeType::Mod: + return a % b; + case IRNodeType::Min: + return min(a, b); + case IRNodeType::Max: + return max(a, b); + case IRNodeType::EQ: + return a == b; + case IRNodeType::NE: + return a != b; + case IRNodeType::LT: + return a < b; + case IRNodeType::LE: + return a <= b; + case IRNodeType::GT: + return a > b; + case IRNodeType::GE: + return a >= b; + case IRNodeType::And: + return a && b; + case IRNodeType::Or: + return a || b; + default: + internal_error << "make_binary_op: " << IRNodeType_string(t) + << " is not a binary operator\n"; + return Expr(); + } +} + bool is_positive_const(const Expr &e) { if (const IntImm *i = e.as()) { return i->value > 0; diff --git a/src/IROperator.h b/src/IROperator.h index 797f12870f5d..ccceeabcc341 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -10,6 +10,7 @@ #include #include #include +#include #include "ConstantInterval.h" #include "Expr.h" @@ -50,6 +51,14 @@ std::optional is_const_power_of_two_integer(uint64_t); std::optional is_const_power_of_two_integer(int64_t); // @} +/** If `e` is a binary operator, return its two operands; otherwise return std::nullopt. */ +std::optional> as_binary_operands(const Expr &e); + +/** Build a binary expression of node type `t` from operands `a` and `b`, using + * the corresponding operator overload (so the usual type matching and constant + * folding apply). `t` must be a binary operator; it is an internal error otherwise. */ +Expr make_binary_op(IRNodeType t, const Expr &a, const Expr &b); + /** Is the expression a const (as defined by is_const), and also * strictly greater than zero (in all lanes, if a vector expression) */ bool is_positive_const(const Expr &e); diff --git a/src/StrictifyFloat.cpp b/src/StrictifyFloat.cpp index 9deb86679808..40dc753655e7 100644 --- a/src/StrictifyFloat.cpp +++ b/src/StrictifyFloat.cpp @@ -85,8 +85,7 @@ class Strictify : public IRMutator { } Expr visit(const Cast *op) override { - if (op->value.type().is_float() && - op->type.is_float()) { + if (op->value.type().is_float() || op->type.is_float()) { return Call::make(op->type, Call::strict_cast, {mutate(op->value)}, Call::PureIntrinsic); } else { diff --git a/test/correctness/rfactor.cpp b/test/correctness/rfactor.cpp index b3d598117168..af9b38c8bdb8 100644 --- a/test/correctness/rfactor.cpp +++ b/test/correctness/rfactor.cpp @@ -1298,6 +1298,503 @@ int isnan_max_rfactor_test() { return 0; } +int hoisted_rfactor_test() { + ImageParam A{Int(8), 2, "A"}; + ImageParam B{Int(8), 2, "B"}; + ImageParam As{Float(16), 1, "As"}; + ImageParam Bs{Float(16), 1, "Bs"}; + + Var i{"i"}, j{"j"}, u{"u"}; + RDom k({{0, A.dim(1).extent() / 4 * 4}}, "k"); + RVar ko{"ko"}, ki{"ki"}; + + Func C{"C"}; + C(i, j) += widening_mul(As(i), Bs(j)) * cast(Int(32), widening_mul(A(i, k), B(j, k))); + C.bound(i, 0, A.dim(0).extent()); + C.bound(j, 0, B.dim(0).extent()); + + // Split k into (ko, ki) and rfactor ko to u, so the intermediate accumulates + // over ki. widening_mul(As(i), Bs(j)) is invariant in k, so by distributivity: + // + // sum_k(scale(i,j) * inner(i,j,k)) = scale(i,j) * sum_k(inner(i,j,k)) + // + // Hoisting should move the loop-invariant scale factor out of the + // intermediate reduction, so the intermediate accumulates only the inner + // integer product and the scale is applied during write-back. + Func C_intm = C.update().split(k, ko, ki, 4).rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + + internal_assert(C_intm.types()[0] == Int(32)) + << "hoisted rfactor: expected C_intm to be Int(32), got " << C_intm.types()[0] << "\n"; + + // Numerical correctness: result must match a reference that applies the full + // non-hoisted reduction. + const int M = 8, N = 8, K = 16; + Buffer a_buf(M, K), b_buf(N, K); + Buffer as_buf(M), bs_buf(N); + for (int m = 0; m < M; m++) { + for (int n_k = 0; n_k < K; n_k++) { + a_buf(m, n_k) = (int8_t)((m + n_k) % 7 - 3); + } + as_buf(m) = float16_t((float)(m + 1) * 0.5f); + } + for (int n = 0; n < N; n++) { + for (int n_k = 0; n_k < K; n_k++) { + b_buf(n, n_k) = (int8_t)((n + n_k + 1) % 5 - 2); + } + bs_buf(n) = float16_t((float)(n + 1) * 0.25f); + } + A.set(a_buf); + B.set(b_buf); + As.set(as_buf); + Bs.set(bs_buf); + + Buffer result = C.realize({M, N}); + + // Reference: plain reduction without rfactor. + Buffer ref(M, N); + ref.fill(0.f); + for (int m = 0; m < M; m++) { + for (int n = 0; n < N; n++) { + for (int kk = 0; kk < K; kk++) { + ref(m, n) += (float)as_buf(m) * (float)bs_buf(n) * + (float)((int32_t)(int16_t)((int16_t)a_buf(m, kk) * (int16_t)b_buf(n, kk))); + } + } + } + + for (int m = 0; m < M; m++) { + for (int n = 0; n < N; n++) { + internal_assert(std::abs(result(m, n) - ref(m, n)) < 1e-6f) + << "hoisted rfactor mismatch at (" << m << ", " << n << "): " + << result(m, n) << " vs ref " << ref(m, n) << "\n"; + } + } + + return 0; +} + +// When HoistInvariantFactor strips a float-promotion cast to expose a +// narrower int type underneath (e.g. widening_mul(int8, int8), which is +// only Int(16)), the intermediate must not accumulate at that raw, +// possibly-too-narrow width: the original (float) accumulation had no +// overflow risk, and the intermediate shouldn't introduce one just because +// a narrow type happened to be sitting under the promotion cast. This +// reproduces exactly that scenario with a reduction long enough (64 terms +// of magnitude up to 127*127) to overflow Int(16) if the fix regresses. +int hoisted_rfactor_widens_narrow_int_test() { + const int K = 64; + ImageParam A{Int(8), 1, "A"}; + ImageParam B{Int(8), 1, "B"}; + ImageParam Scale{Float(32), 1, "Scale"}; + + Var i{"i"}; + RDom r(0, K, "r"); + + Func Acc{"Acc"}; + Acc(i) = 0.0f; + // No explicit cast() here -- unlike hoisted_rfactor_test above, + // this is deliberately the "naive" shape a decode()-then-multiply + // composition produces (see test/correctness/approximation_composition.cpp), + // relying on HoistInvariantFactor to pick a safe accumulation width on + // its own rather than the caller having to know to widen it by hand. + Acc(i) += Scale(i) * cast(widening_mul(A(r), B(r))); + + Func Acc_intm = Acc.update().rfactor({}, RFactorOptions::HoistInvariantFactor); + internal_assert(Acc_intm.types()[0] == Int(32)) + << "hoisted rfactor: expected the narrow Int(16) product to be widened to " + << "Int(32) for safe accumulation, got " << Acc_intm.types()[0] << "\n"; + Acc_intm.compute_root(); + + Buffer a_buf(K), b_buf(K); + for (int k = 0; k < K; k++) { + // Same sign every term (max magnitude, same direction) so the + // partial sums only ever grow -- the worst case for narrow-int + // overflow, and guaranteed to exceed Int(16)'s range well before k + // reaches K. + a_buf(k) = 127; + b_buf(k) = 127; + } + Buffer scale_buf(1); + scale_buf(0) = 1.0f; + A.set(a_buf); + B.set(b_buf); + Scale.set(scale_buf); + + Buffer result = Acc.realize({1}); + const float expected = (float)K * 127.0f * 127.0f; + internal_assert(result(0) == expected) + << "hoisted rfactor narrow-int widening: got " << result(0) << ", expected " << expected << "\n"; + + return 0; +} + +// A single top-level split can only find an invariant factor that sits as +// one of a Mul's two *immediate* operands. This reproduces a shape where +// two independent invariant factors are each nested inside their own +// sub-product instead -- (scaleA(i) * castA) * (scaleB(i) * castB) -- the +// way two independently-quantized operands, each decode()d on its own, +// naturally compose (see test/correctness/approximation_composition.cpp). +// Finding either factor requires flattening the whole multiplicative chain +// into leaves and partitioning every leaf, not just checking a binary +// node's two immediate children. Once both scales are hoisted out, the +// remaining body is exactly cast(a) * cast(b) for int8 a, b +// -- which HoistInvariantFactor should also fold to a widening_mul, +// reaching the same Int(32) accumulation as +// hoisted_rfactor_test above, despite neither being spelled out explicitly +// anywhere in this reduction's source. +int hoisted_rfactor_scattered_factors_test() { + const int K = 64; + ImageParam A{Int(8), 1, "A"}; + ImageParam B{Int(8), 1, "B"}; + ImageParam ScaleA{Float(32), 1, "ScaleA"}; + ImageParam ScaleB{Float(32), 1, "ScaleB"}; + + Var i{"i"}; + RDom r(0, K, "r"); + + Func Acc{"Acc"}; + Acc(i) = 0.0f; + Acc(i) += (cast(A(r)) * ScaleA(i)) * (cast(B(r)) * ScaleB(i)); + + Func Acc_intm = Acc.update().rfactor({}, RFactorOptions::HoistInvariantFactor); + internal_assert(Acc_intm.types()[0] == Int(32)) + << "hoisted rfactor: expected both scattered scale factors to be found and hoisted, " + << "and the remaining int8*int8 product widened to Int(32), got " << Acc_intm.types()[0] << "\n"; + Acc_intm.compute_root(); + + Buffer a_buf(K), b_buf(K); + Buffer scale_a_buf(1), scale_b_buf(1); + for (int k = 0; k < K; k++) { + a_buf(k) = 127; + b_buf(k) = 127; + } + scale_a_buf(0) = 2.0f; + scale_b_buf(0) = 3.0f; + A.set(a_buf); + B.set(b_buf); + ScaleA.set(scale_a_buf); + ScaleB.set(scale_b_buf); + + Buffer result = Acc.realize({1}); + const float expected = 2.0f * 3.0f * (float)K * 127.0f * 127.0f; + internal_assert(result(0) == expected) + << "hoisted rfactor scattered factors: got " << result(0) << ", expected " << expected << "\n"; + + return 0; +} + +// Same shape as hoisted_rfactor_scattered_factors_test, but with UInt(8) +// operands instead of Int(8) -- the widening-mul pattern table in +// extract_factor() has separate entries for signed and unsigned sources, so +// this exercises the unsigned ones. +int hoisted_rfactor_scattered_factors_unsigned_test() { + const int K = 64; + ImageParam A{UInt(8), 1, "A"}; + ImageParam B{UInt(8), 1, "B"}; + ImageParam ScaleA{Float(32), 1, "ScaleA"}; + ImageParam ScaleB{Float(32), 1, "ScaleB"}; + + Var i{"i"}; + RDom r(0, K, "r"); + + Func Acc{"Acc"}; + Acc(i) = 0.0f; + Acc(i) += (cast(A(r)) * ScaleA(i)) * (cast(B(r)) * ScaleB(i)); + + Func Acc_intm = Acc.update().rfactor({}, RFactorOptions::HoistInvariantFactor); + internal_assert(Acc_intm.types()[0] == UInt(32)) + << "hoisted rfactor: expected both scattered scale factors to be found and hoisted, " + << "and the remaining uint8*uint8 product widened to UInt(32), got " << Acc_intm.types()[0] << "\n"; + Acc_intm.compute_root(); + + Buffer a_buf(K), b_buf(K); + Buffer scale_a_buf(1), scale_b_buf(1); + for (int k = 0; k < K; k++) { + a_buf(k) = 255; + b_buf(k) = 255; + } + scale_a_buf(0) = 2.0f; + scale_b_buf(0) = 3.0f; + A.set(a_buf); + B.set(b_buf); + ScaleA.set(scale_a_buf); + ScaleB.set(scale_b_buf); + + Buffer result = Acc.realize({1}); + const float expected = 2.0f * 3.0f * (float)K * 255.0f * 255.0f; + internal_assert(result(0) == expected) + << "hoisted rfactor scattered factors (unsigned): got " << result(0) << ", expected " << expected << "\n"; + + return 0; +} + +// Hoisted rfactor 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. +int hoisted_rfactor_min_test() { + ImageParam offset_p{Float(32), 1, "offset_p"}; + ImageParam data_p{Float(32), 2, "data_p"}; + + Var i{"i"}, u{"u"}; + const int K = 16; + RDom k(0, K, "k"); + + Func C{"C"}; + C(i) = Float(32).max(); + C(i) = min(C(i), offset_p(i) + data_p(i, k)); + + RVar ko{"ko"}, ki{"ki"}; + C.update().split(k, ko, ki, 4); + Func C_intm = C.update().rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + C_intm.compute_root(); + + const int M = 8; + Buffer off(M), dat(M, K); + for (int m = 0; m < M; m++) { + off(m) = (float)(m + 1); + for (int kk = 0; kk < K; kk++) { + dat(m, kk) = (float)(((m * K + kk) % 7) - 3); + } + } + offset_p.set(off); + data_p.set(dat); + + Buffer result = C.realize({M}); + + Buffer ref(M); + for (int m = 0; m < M; m++) { + float v = std::numeric_limits::max(); + for (int kk = 0; kk < K; kk++) { + v = std::min(v, off(m) + dat(m, kk)); + } + ref(m) = v; + } + + for (int m = 0; m < M; m++) { + internal_assert(result(m) == ref(m)) + << "hoisted rfactor min mismatch at " << m << ": " + << result(m) << " vs ref " << ref(m) << "\n"; + } + return 0; +} + +// Hoisted rfactor with outer Or (bool) and And factor: +// or_k(mask(i) && check(i, k)) = mask(i) && or_k(check(i, k)) +// The intermediate accumulates or without the mask; write-back applies it once. +int hoisted_rfactor_or_test() { + ImageParam mask_p{Bool(), 1, "mask_p"}; + ImageParam check_p{Bool(), 2, "check_p"}; + + Var i{"i"}, u{"u"}; + const int K = 16; + RDom k(0, K, "k"); + + Func valid{"valid"}; + valid(i) = cast(false); + valid(i) = valid(i) || (mask_p(i) && check_p(i, k)); + + RVar ko{"ko"}, ki{"ki"}; + valid.update().split(k, ko, ki, 4); + Func valid_intm = valid.update().rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + valid_intm.compute_root(); + + const int M = 8; + Buffer mask(M), chk(M, K); + for (int m = 0; m < M; m++) { + mask(m) = (m % 2 == 0); + for (int kk = 0; kk < K; kk++) { + chk(m, kk) = ((m + kk) % 3 == 0); + } + } + mask_p.set(mask); + check_p.set(chk); + + Buffer result = valid.realize({M}); + + for (int m = 0; m < M; m++) { + bool ref = false; + for (int kk = 0; kk < K; kk++) { + ref = ref || (mask(m) && chk(m, kk)); + } + internal_assert(result(m) == ref) + << "hoisted rfactor or mismatch at " << m << ": " + << (int)result(m) << " vs ref " << (int)ref << "\n"; + } + return 0; +} + +// Hoisted rfactor with outer Min and an additive float offset over an integer +// body: +// min_k(offset(i) + body(i, k)) = offset(i) + min_k(body(i, k)) +// The float offset implicitly promotes the integer body to float. Hoisting +// strips that promotion so the intermediate accumulates min over Int(32); the +// offset is added once at write-back. This exercises the min/max narrowing path: +// the intermediate's identity must be Int(32).max(), not a cast of the original +// Float(32) identity. +int hoisted_rfactor_min_narrow_test() { + ImageParam offset_p{Float(32), 1, "offset_p"}; + ImageParam data_p{Int(32), 2, "data_p"}; + + Var i{"i"}, u{"u"}; + const int K = 16; + RDom k(0, K, "k"); + + Func C{"C"}; + C(i) = Float(32).max(); + C(i) = min(C(i), offset_p(i) + data_p(i, k)); + + RVar ko{"ko"}, ki{"ki"}; + C.update().split(k, ko, ki, 4); + Func C_intm = C.update().rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + C_intm.compute_root(); + + internal_assert(C_intm.types()[0] == Int(32)) + << "hoisted rfactor min: expected C_intm to be Int(32), got " << C_intm.types()[0] << "\n"; + + const int M = 8; + Buffer off(M); + Buffer dat(M, K); + for (int m = 0; m < M; m++) { + off(m) = (float)(m + 1) * 0.5f; + for (int kk = 0; kk < K; kk++) { + dat(m, kk) = ((m * K + kk) % 7) - 3; + } + } + offset_p.set(off); + data_p.set(dat); + + Buffer result = C.realize({M}); + + Buffer ref(M); + for (int m = 0; m < M; m++) { + float v = std::numeric_limits::max(); + for (int kk = 0; kk < K; kk++) { + v = std::min(v, off(m) + (float)dat(m, kk)); + } + ref(m) = v; + } + + for (int m = 0; m < M; m++) { + internal_assert(result(m) == ref(m)) + << "hoisted rfactor min narrow mismatch at " << m << ": " + << result(m) << " vs ref " << ref(m) << "\n"; + } + return 0; +} + +// Hoisted rfactor strips implicit int-to-float Cast nodes so the intermediate +// can accumulate at integer type. It must not strip strict_cast, since that +// explicitly requests that the conversion happen before the reduction. +int hoisted_rfactor_strict_cast_test() { + Buffer data(2); + data(0) = 16777217; + data(1) = -16777216; + + Var u{"u"}; + RDom k(0, 2, "k"); + + Func f{"f"}; + f() = 0.0f; + f() += 1.5f * strict_float(cast(data(k))); + + RVar ko{"ko"}, ki{"ki"}; + f.update().split(k, ko, ki, 2); + Func intm = f.update().rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + intm.compute_root(); + + internal_assert(intm.types()[0] == Float(32)) + << "hoisted rfactor strict_cast: expected intm to remain Float(32), got " + << intm.types()[0] << "\n"; + + Buffer result = f.realize(); + internal_assert(result() == 0.0f) + << "hoisted rfactor strict_cast mismatch: " << result() + << " vs ref 0\n"; + + return 0; +} + +// A factor may depend on an RVar that rfactor preserves. In the intermediate +// update below, r.y is replaced by the pure Var u, so scale(u) is invariant +// over the intermediate's reduction over r.x and can be hoisted to writeback. +int hoisted_rfactor_preserved_rvar_factor_test() { + ImageParam scale_p{Float(32), 1, "scale_p"}; + ImageParam data_p{Int(32), 2, "data_p"}; + + Var u{"u"}; + const int X = 8, Y = 4; + RDom r(0, X, 0, Y, "r"); + + Func f{"f"}; + f() = 0.0f; + f() += scale_p(r.y) * data_p(r.x, r.y); + + Func intm = f.update().rfactor({{r.y, u}}, RFactorOptions::HoistInvariantFactor); + intm.compute_root(); + + internal_assert(intm.types()[0] == Int(32)) + << "hoisted rfactor preserved-rvar factor: expected intm to be Int(32), got " + << intm.types()[0] << "\n"; + + Buffer scale(Y); + Buffer data(X, Y); + for (int y = 0; y < Y; y++) { + scale(y) = (float)(y + 1); + for (int x = 0; x < X; x++) { + data(x, y) = (x + 2 * y) % 7 - 3; + } + } + scale_p.set(scale); + data_p.set(data); + + Buffer result = f.realize(); + + float ref = 0.0f; + for (int y = 0; y < Y; y++) { + int partial = 0; + for (int x = 0; x < X; x++) { + partial += data(x, y); + } + ref += scale(y) * (float)partial; + } + + internal_assert(result() == ref) + << "hoisted rfactor preserved-rvar factor mismatch: " + << result() << " vs ref " << ref << "\n"; + + return 0; +} + +// The min/max + add hoisting law is only valid for integer types where addition +// has no defined wraparound behavior. For UInt(8), hoisting the invariant 250 +// would incorrectly turn min_k((250 + x_k) mod 256) into +// (250 + min_k(x_k)) mod 256. +int hoisted_rfactor_unsigned_min_fallback_test() { + Buffer data(2); + data(0) = 1; + data(1) = 10; + + Var u{"u"}; + RDom k(0, 2, "k"); + + Func f{"f"}; + f() = UInt(8).max(); + f() = min(f(), cast(250) + data(k)); + + RVar ko{"ko"}, ki{"ki"}; + f.update().split(k, ko, ki, 2); + Func intm = f.update().rfactor({{ko, u}}, RFactorOptions::HoistInvariantFactor); + intm.compute_root(); + + Buffer result = f.realize(); + const uint8_t expected = 4; + internal_assert(result() == expected) + << "hoisted rfactor unsigned min fallback mismatch: " + << (int)result() << " vs ref " << (int)expected << "\n"; + + return 0; +} + } // namespace int main(int argc, char **argv) { @@ -1347,6 +1844,16 @@ int main(int argc, char **argv) { {"rfactor bounds tests", rfactor_precise_bounds_test}, {"isnan max rfactor test (bitwise or)", isnan_max_rfactor_test}, {"isnan max rfactor test (logical or)", isnan_max_rfactor_test}, + {"hoisted rfactor test (add/mul)", hoisted_rfactor_test}, + {"hoisted rfactor test (add/mul, narrow int widened)", hoisted_rfactor_widens_narrow_int_test}, + {"hoisted rfactor test (add/mul, scattered factors)", hoisted_rfactor_scattered_factors_test}, + {"hoisted rfactor test (add/mul, scattered factors, unsigned)", hoisted_rfactor_scattered_factors_unsigned_test}, + {"hoisted rfactor test (min/add)", hoisted_rfactor_min_test}, + {"hoisted rfactor test (min/add, narrowing)", hoisted_rfactor_min_narrow_test}, + {"hoisted rfactor test (strict_cast preserved)", hoisted_rfactor_strict_cast_test}, + {"hoisted rfactor test (preserved-rvar factor)", hoisted_rfactor_preserved_rvar_factor_test}, + {"hoisted rfactor test (min/add, unsigned fallback)", hoisted_rfactor_unsigned_min_fallback_test}, + {"hoisted rfactor test (or/and)", hoisted_rfactor_or_test}, }; using Sharder = Halide::Internal::Test::Sharder; diff --git a/test/performance/CMakeLists.txt b/test/performance/CMakeLists.txt index cee54a47c132..35af5ddd5c05 100644 --- a/test/performance/CMakeLists.txt +++ b/test/performance/CMakeLists.txt @@ -28,6 +28,7 @@ tests( realize_overhead.cpp rgb_interleaved.cpp tiled_matmul.cpp + tiled_matmul_arm_neon.cpp vectorize.cpp wrap.cpp # keep-sorted end diff --git a/test/performance/tiled_matmul_arm_neon.cpp b/test/performance/tiled_matmul_arm_neon.cpp new file mode 100644 index 000000000000..544cb97e175f --- /dev/null +++ b/test/performance/tiled_matmul_arm_neon.cpp @@ -0,0 +1,197 @@ +#include "Halide.h" +#include "halide_benchmark.h" + +using namespace Halide; + +// Quantized (int8 x int8 -> int32) mat-vec with a per-row weight scale and a +// per-vector scale, targeting ARM's SDOT instruction, scheduled via rfactor's +// distributive-law hoisting mode. +// +// WtPacked(k, i, blk) = Wt(blk*reduce+k, i): the weight matrix, repacked so +// the dimension being vectorized across (i) sits next to the reduction chunk +// (k), not behind the block index. That makes each sdot operand load a +// single contiguous 16-byte read instead of a gather across K-byte strides. +// +// Two variants are compared: Hoisted factors WtScale(i) * VecScale (the +// invariant scale product) out of the reduction before splitting it by +// block, so the per-block partial sums are a pure Int(32) dot product +// eligible for SDOT. PlainRfactor uses the same split/atomic/vectorize +// schedule but without hoisting, so the scale multiply stays inline in every +// term of the per-block reduction: the partial-sum intermediate is +// Float(32)-typed, and CodeGen_ARM's i32(widening_mul(i8, i8)) pattern can't +// match it, so no sdot is generated even though the reduction is still +// turned into a horizontal vector reduce. + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::ARMDotProd)) { + printf("[SKIP] This test requires ARM's SDOT instruction.\n"); + return 0; + } + + const int M = 1024; // output rows + const int K = 1024; // reduction extent + const int reduce = 4; // SDOT contracts 4 int8 lanes per output lane + + enum { Hoisted, + PlainRfactor, + NumVariants }; + double times[NumVariants]; + Buffer outs[NumVariants]; + + for (int variant = 0; variant < NumVariants; variant++) { + Var i{"i"}, u{"u"}; + RDom r(0, K, "r"); + + ImageParam WtPacked(Int(8), 3, "WtPacked"); // WtPacked(k, i, blk) = Wt(blk*reduce+k, i) + ImageParam WtScale(Float(32), 1, "WtScale"); // WtScale(i): per-row weight scale + ImageParam Vec(Int(8), 1, "Vec"); // Vec(k): quantized vector + Param VecScale("VecScale"); // single scale for the vector + + // Without this, WtPacked's row stride is an opaque runtime value (an + // ImageParam doesn't otherwise promise anything about it), so the + // simplifier can't prove that 4 consecutive rows' data is *also* + // contiguous with the reduction dimension -- it's forced to build each + // sdot operand via a scalar gather-and-insert instead of a single dense + // vld1q_s8. Declaring the true packed strides lets it prove the combined + // (k, i) index is one flat dense ramp. + WtPacked.dim(0).set_stride(1); + WtPacked.dim(1).set_stride(reduce); + + Func Acc("Acc"); + Acc(i) = 0.0f; + Acc(i) += WtScale(i) * VecScale * cast(widening_mul(WtPacked(r % reduce, i, r / reduce), Vec(r))); + + // ro: which block of `reduce` elements a term belongs to; ri: its + // position within that block. + RVar ro{"ro"}, ri{"ri"}; + Acc.update().split(r, ro, ri, reduce); + + const int panel = target.natural_vector_size() * 4; + + Func Result("Result"); + Result(i) = Acc(i); + + Var io, ii; + Result.bound(i, 0, M); + Result.split(i, io, ii, panel).vectorize(ii, panel); + + if (variant == Hoisted) { + // Factor the reduction in two steps. First, preserving nothing + // hoists WtScale(i) * VecScale -- invariant across all of r -- + // out of the reduction entirely: Acc_wb becomes a pure Int(32) + // dot product over all of K, eligible for SDOT, and Acc's own + // update collapses to a single multiply per row: + // Acc(i) += WtScale(i) * VecScale * Acc_wb(i). + Func Acc_wb = Acc.update().rfactor({}, RFactorOptions::HoistInvariantFactor); + + // Second, factor Acc_wb's own (now scale-free) reduction by + // block: no hoisting needed here, just combining per-block + // partial sums. Preserving ro turns it into a new dimension u of + // Acc_dot, so Acc_dot(u, i) holds one block's partial dot + // product (reduced over ri only), while Acc_wb sums + // Acc_dot(ro, i) across blocks. ro keeps its identity across + // both rfactor calls, so it's still the same RVar driving + // Acc_wb's reduction loop -- usable directly below as the + // compute_at level for Acc_dot. + Func Acc_dot = Acc_wb.update().rfactor(ro, u); + + // Compute a whole panel of rows through every stage (Acc, + // Acc_wb, Acc_dot) before moving to the next panel, rather than + // fully materializing each stage before the next reads it: Acc + // and Acc_wb are computed_at the panel loop, and Acc_dot is + // fused one level deeper still, into Acc_wb's own per-block + // reduction loop, so it only ever needs panel-many elements -- + // small enough to live on the stack/in registers instead of + // being heap-allocated. + Acc_wb.compute_at(Result, io) + .vectorize(i, panel) + .update() + .vectorize(i, panel); + + // Acc_dot's row dimension (i) is deliberately left unvectorized: + // fused this deep inside Acc_wb's own per-block reduction loop, + // vectorizing it makes the vectorizer conflate the two loops and + // build a wildly oversized (and out-of-bounds) vector + // expression. The sdot reduction itself is still + // atomic-vectorized over ri, so this costs nothing. + Acc_dot.compute_at(Acc_wb, ro) + .update() + .reorder(ri, i, u) + .atomic() + .vectorize(ri, reduce) + .unroll(ri); + } else { + // No invariant factor to hoist out this time, so a single + // rfactor preserving ro is enough: Acc_dot(u, i) holds one + // block's partial sum of WtScale(i) * VecScale * term(ri, i) + // (Float(32), since the scale multiply is still inline), and + // Acc sums Acc_dot(ro, i) across blocks directly -- there's no + // separate write-back stage, since Acc's own combine was never + // factored out. + Func Acc_dot = Acc.update().rfactor(ro, u, RFactorOptions::None); + + Acc_dot.compute_at(Acc, ro) + .update() + .reorder(ri, i, u) + .atomic() + .vectorize(ri, reduce) + .unroll(ri); + } + + Acc.compute_at(Result, io) + .vectorize(i, panel) + .update() + .vectorize(i, panel); + + Buffer wt_packed(reduce, M, K / reduce); + for (int blk = 0; blk < K / reduce; blk++) { + for (int y = 0; y < M; y++) { + for (int k = 0; k < reduce; k++) { + wt_packed(k, y, blk) = (int8_t)((((reduce * blk + k) + y) % 15) - 7); + } + } + } + + Buffer wt_scale_buf(M); + for (int y = 0; y < M; y++) { + wt_scale_buf(y) = 0.01f * ((y % 7) + 1); + } + + Buffer vec_buf(K); + for (int x = 0; x < K; x++) { + vec_buf(x) = (int8_t)((x % 13) - 6); + } + + WtPacked.set(wt_packed); + WtScale.set(wt_scale_buf); + Vec.set(vec_buf); + VecScale.set(0.5f); + + Buffer out(M); + Result.realize(out, target); + outs[variant] = out; + + times[variant] = Tools::benchmark([&] { + Result.realize(out, target); + }); + } + + for (int y = 0; y < M; y++) { + float ref = outs[PlainRfactor](y); + if (std::abs(ref - outs[Hoisted](y)) > 1e-3f * std::abs(ref)) { + printf("Quantized mat-vec mismatch at %d: %f (plain) vs %f (hoisted)\n", y, ref, outs[Hoisted](y)); + return 1; + } + } + + printf("Quantized mat-vec (int8 x int8 -> f32, SDOT)\n" + "Time with non-hoisting rfactor: %0.4f ms\n" + "Time with hoisted rfactor: %0.4f ms (%0.2fx vs non-hoisting)\n", + times[PlainRfactor] * 1000, + times[Hoisted] * 1000, + times[PlainRfactor] / times[Hoisted]); + + printf("Success!\n"); + return 0; +}