diff --git a/python_bindings/src/halide/halide_/PyStage.cpp b/python_bindings/src/halide/halide_/PyStage.cpp index 934c58bd9439..492e53fdd6db 100644 --- a/python_bindings/src/halide/halide_/PyStage.cpp +++ b/python_bindings/src/halide/halide_/PyStage.cpp @@ -23,6 +23,7 @@ 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("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 b356665cbb99..80f82b05b689 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -581,55 +581,6 @@ 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) { definition.schedule().touched() = true; return rfactor({{r, v}}); @@ -638,6 +589,29 @@ Func Stage::rfactor(const RVar &r, const Var &v) { // 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); @@ -690,17 +664,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 @@ -766,6 +750,41 @@ pair project_rdom(const vector &dims, con return {new_rdom, dim_projection}; } +// A semiring distributive law used by hoist_invariants(): `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; +}; + +constexpr DistributiveLaw distributive_laws[] = { + {IRNodeType::Add, IRNodeType::Mul}, // sum_k(s * x_k) = s * sum_k(x_k) + {IRNodeType::Min, IRNodeType::Add}, // min_k(c + x_k) = c + min_k(x_k) + {IRNodeType::Max, IRNodeType::Add}, // max_k(c + x_k) = c + max_k(x_k) + {IRNodeType::Or, IRNodeType::And}, // or_k(p && x_k) = p && or_k(x_k) + {IRNodeType::And, IRNodeType::Or}, // 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. @@ -784,6 +803,115 @@ optional> select_binary_operand(const Expr &e, IRNodeType op, P 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. It keeps + // its natural type; changing the accumulation type is the + // job of the separate change_type() directive. +}; + +// 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) { + auto is_rvar_free = [&](const Expr &e) { + return !expr_uses_vars(e, reduction_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_rvar_free(leaf)) { + invariant_leaves.push_back(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; + } + + 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()); + + 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])) { + // The value may be wrapped in Let nodes (e.g. an rfactor of this same + // 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); + } + } + } + return result; +} + +// Re-apply the hoisted factor to the intermediate's result at the write-back +// step. The intermediate accumulates the inner body at its natural type (the +// same type as the outer op), so no cast is needed. +Expr apply_hoisted_factor(const Expr &r, const optional &factor) { + if (!factor) { + return 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) { @@ -944,6 +1072,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); @@ -952,9 +1083,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)); } @@ -967,9 +1096,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); } @@ -989,13 +1116,11 @@ Func Stage::rfactor(const vector> &preserved) { { 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); vector values = definition.values(); 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 @@ -1029,7 +1154,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; } @@ -1089,9 +1214,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; } @@ -1099,6 +1224,104 @@ Func Stage::rfactor(const vector> &preserved) { return intm; } +Func Stage::hoist_invariants() { + user_assert(!definition.is_init()) << "hoist_invariants() must be called on an update definition\n"; + + definition.schedule().touched() = true; + + // Check whether the operator is associative and determine the operator and + // its identity for each value in the definition if it is a Tuple. + const auto &prover_result = prove_associativity(function.name(), definition.args(), definition.values()); + const auto &[var_splits, _] = rfactor_validate_args({}, prover_result); + + const vector dim_vars_exprs(dim_vars.begin(), dim_vars.end()); + + Scope<> reduction_vars; + Scope reduction_bounds; + for (const auto &[var, min, extent] : definition.schedule().rvars()) { + 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"; + + Func intm(function.name() + "_intm"); + intm(dim_vars_exprs) = Tuple(prover_result.pattern.identities); + + // 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(), {}); + + // 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(); + } + + // Replace the original reduction with a factor-applying write-back update. + { + 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]); + add_let(writeback_map, prover_result.ys[i].var, r); + } + + if (!prover_result.xs[i].var.empty()) { + Expr prev_val = Call::make(function.output_types()[i], function.name(), + dim_vars_exprs, Call::Halide, + FunctionPtr(), (int)i); + add_let(writeback_map, prover_result.xs[i].var, prev_val); + } else { + user_warning << "Update definition of " << name() << " at index " << i + << " doesn't depend on the previous value. This isn't a" + << " reduction operation\n"; + } + } + + vector writeback_dims; + for (const Dim &dim : definition.schedule().dims()) { + if (!dim.is_rvar()) { + writeback_dims.push_back(dim); + } + } + // Add pure vars not referenced by the original update just before + // outermost, as rfactor() does for histogram-style updates. + for (size_t i = 0; i < dim_vars.size(); i++) { + if (!expr_uses_var(definition.args()[i], dim_vars[i].name())) { + Dim d = {dim_vars[i].name(), ForType::Serial, DeviceAPI::None, + DimType::PureVar, Partition::Auto}; + writeback_dims.insert(writeback_dims.end() - 1, d); + } + } + + definition.args() = dim_vars_exprs; + definition.values() = substitute(writeback_map, prover_result.pattern.ops); + definition.predicate() = or_condition_over_domain(definition.predicate(), reduction_bounds); + definition.schedule().dims() = std::move(writeback_dims); + definition.schedule().rvars().clear(); + definition.schedule().splits() = var_splits; + } + + return intm; +} + void Stage::split(const string &old, const string &outer, const string &inner, const Expr &factor_arg, bool exact, TailStrategy tail) { debug(4) << "In schedule for " << name() << ", split " << old << " into " << outer << " and " << inner << " with factor of " << factor_arg << "\n"; diff --git a/src/Func.h b/src/Func.h index 4a940eb67969..67da971e49ce 100644 --- a/src/Func.h +++ b/src/Func.h @@ -18,6 +18,7 @@ #include "Var.h" #include +#include #include namespace Halide { @@ -216,6 +217,44 @@ class Stage { } // @} + /** Hoist a loop-invariant factor 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. + * + * 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 + * 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 example, hoist_invariants() rewrites a pipeline like this: + * \code + * f(x) = 0; + * f(x) += s(x) * g(x, r); + * \endcode + * into a pipeline like this: + * \code + * f_intm(x) = 0; + * f_intm(x) += g(x, r); + * + * f(x) = 0; + * f(x) += s(x) * f_intm(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. + */ + Func 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 * be computed AFTER 's' in the innermost fused dimension. There should not diff --git a/test/correctness/rfactor.cpp b/test/correctness/rfactor.cpp index b3d598117168..eaa5a038f373 100644 --- a/test/correctness/rfactor.cpp +++ b/test/correctness/rfactor.cpp @@ -1298,6 +1298,545 @@ int isnan_max_rfactor_test() { return 0; } +// hoist_invariants() lifts a loop-invariant factor out of a sum reduction: +// sum_k(scale(i,j) * inner(i,j,k)) = scale(i,j) * sum_k(inner(i,j,k)) +// It does not change any types: the returned intermediate accumulates the +// factor-free body at its natural type. (Retyping the accumulation for a +// dot-product-friendly integer type is the job of change_type().) +int hoist_invariants_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"}; + 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()); + C.update().split(k, ko, ki, 4); + + // widening_mul(As(i), Bs(j)) is invariant in k, so hoisting moves it out of + // the reduction: the intermediate accumulates only the inner product and the + // scale is applied once during write-back. The body's type (Float(32)) is + // unchanged. + Func C_intm = C.update().hoist_invariants(); + + internal_assert(C_intm.types()[0] == Float(32)) + << "hoist_invariants: expected C_intm to keep its natural Float(32) type, 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 hoisting. + 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) + << "hoist_invariants mismatch at (" << m << ", " << n << "): " + << result(m, n) << " vs ref " << ref(m, n) << "\n"; + } + } + + return 0; +} + +// An invariant factor may be nested inside its own sub-product rather than +// sitting as one of a Mul's two immediate operands: +// (scaleA(i) * castA) * (scaleB(i) * castB) +// -- the way two independently-scaled operands naturally compose. Finding both +// scales requires flattening the whole multiplicative chain into leaves and +// partitioning each one by invariance, not just checking a binary node's two +// immediate children. +int hoist_invariants_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().hoist_invariants(); + internal_assert(Acc_intm.types()[0] == Float(32)) + << "hoist_invariants: expected the intermediate to keep Float(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) + << "hoist_invariants scattered factors: got " << result(0) << ", expected " << expected << "\n"; + + return 0; +} + +// Same shape as hoist_invariants_scattered_factors_test, but with UInt(8) +// operands, to exercise the unsigned path. +int hoist_invariants_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().hoist_invariants(); + internal_assert(Acc_intm.types()[0] == Float(32)) + << "hoist_invariants: expected the intermediate to keep Float(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) + << "hoist_invariants scattered factors (unsigned): got " << result(0) << ", expected " << expected << "\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. +int hoist_invariants_min_test() { + ImageParam offset_p{Float(32), 1, "offset_p"}; + ImageParam data_p{Float(32), 2, "data_p"}; + + Var i{"i"}; + 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)); + + Func C_intm = C.update().hoist_invariants(); + 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)) + << "hoist_invariants min mismatch at " << m << ": " + << result(m) << " vs ref " << ref(m) << "\n"; + } + return 0; +} + +// hoist_invariants() 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 hoist_invariants_or_test() { + ImageParam mask_p{Bool(), 1, "mask_p"}; + ImageParam check_p{Bool(), 2, "check_p"}; + + Var i{"i"}; + 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)); + + Func valid_intm = valid.update().hoist_invariants(); + 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) + << "hoist_invariants or mismatch at " << m << ": " + << (int)result(m) << " vs ref " << (int)ref << "\n"; + } + return 0; +} + +// hoist_invariants() must not disturb strict_float: the reduction below rounds +// each term to float before summing, and both the reference and the hoisted +// intermediate must observe that same rounding (giving exactly 0). +int hoist_invariants_strict_float_test() { + Buffer data(2); + data(0) = 16777217; + data(1) = -16777216; + + RDom k(0, 2, "k"); + + Func f{"f"}; + f() = 0.0f; + f() += 1.5f * strict_float(cast(data(k))); + + Func intm = f.update().hoist_invariants(); + intm.compute_root(); + + internal_assert(intm.types()[0] == Float(32)) + << "hoist_invariants strict_float: expected intm to remain Float(32), got " + << intm.types()[0] << "\n"; + + Buffer result = f.realize(); + internal_assert(result() == 0.0f) + << "hoist_invariants strict_float mismatch: " << result() + << " vs ref 0\n"; + + return 0; +} + +// A loop-invariant RDom predicate must be preserved on the write-back update. +// When enabled is false, the original reduction executes no iterations, so the +// hoisted factor (and its failing require) must not be evaluated. +int hoist_invariants_predicated_rdom_test() { + Param enabled{"enabled"}; + RDom r(0, 4, "r"); + r.where(enabled); + + Func f{"f"}; + f() = 0; + f() += require(enabled, 2, "hoisted factor evaluated outside RDom predicate") * (r + 1); + + Func intm = f.update().hoist_invariants(); + intm.compute_root(); + + Module module = f.compile_to_module(f.infer_arguments(), "hoist_invariants_predicated_rdom"); + bool inside_predicate = false; + bool found_writeback = false; + for (const LoweredFunc &lowered_func : module.functions()) { + visit_with( + lowered_func.body, + [&](auto *self, const IfThenElse *op) { + const bool old_inside_predicate = inside_predicate; + inside_predicate |= expr_uses_var(op->condition, enabled.name()); + (*self)(op->then_case); + inside_predicate = old_inside_predicate; + + if (op->else_case.defined()) { + (*self)(op->else_case); + } + }, + [&](auto *self, const Store *op) { + const bool predicated = + inside_predicate || expr_uses_var(op->predicate, enabled.name()); + if (op->name == f.name() && predicated) { + visit_with(op->value, [&](auto *, const Load *load) { + found_writeback |= load->name == intm.name(); + }); + } + self->visit_base(op); + }); + } + internal_assert(found_writeback) + << "hoist_invariants predicated RDom: lowered pipeline had no write-back " + << "guarded by " << enabled.name() << " that reads " << intm.name() << "\n"; + + enabled.set(false); + Buffer disabled_result = f.realize(); + internal_assert(disabled_result() == 0) + << "hoist_invariants predicated RDom: disabled result was " + << disabled_result() << ", expected 0\n"; + + enabled.set(true); + Buffer enabled_result = f.realize(); + internal_assert(enabled_result() == 20) + << "hoist_invariants predicated RDom: enabled result was " + << enabled_result() << ", expected 20\n"; + + return 0; +} + +// hoist_invariants() composes with rfactor(): rfactor first preserves r.y as a +// pure Var u, so scale(u) becomes invariant over the intermediate's remaining +// reduction over r.x and can then be hoisted from that intermediate. +int hoist_invariants_after_rfactor_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); + + // Preserve r.y as u; the intermediate now reduces only over r.x, and + // scale_p(u) is invariant across that reduction. + Func intm = f.update().rfactor({{r.y, u}}); + Func intm2 = intm.update().hoist_invariants(); + intm.compute_root(); + intm2.compute_root(); + + 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) + << "hoist_invariants after rfactor mismatch: " + << result() << " vs ref " << ref << "\n"; + + return 0; +} + +// A direct 2D Gaussian blur can be made separable by preserving r.y with +// rfactor(), then hoisting kernel(r.y) out of the intermediate's remaining +// r.x reduction. The factor-free intermediate performs the horizontal blur; +// its write-back applies the vertical kernel weight, and the original Func +// combines those weighted rows. +int hoist_invariants_separable_gaussian_after_rfactor_test() { + constexpr int radius = 2; + constexpr int diameter = 2 * radius + 1; + constexpr int width = 17; + constexpr int height = 13; + + Buffer kernel(diameter); + kernel.set_min(-radius); + const int binomial_weights[diameter] = {1, 4, 6, 4, 1}; + for (int i = 0; i < diameter; i++) { + kernel(i - radius) = (float)binomial_weights[i] / 16.0f; + } + + Buffer input(width + 2 * radius, height + 2 * radius); + input.set_min(-radius, -radius); + for (int y = input.dim(1).min(); y <= input.dim(1).max(); y++) { + for (int x = input.dim(0).min(); x <= input.dim(0).max(); x++) { + input(x, y) = (float)((7 * x + 13 * y + 101) % 29 - 14); + } + } + + Var x{"x"}, y{"y"}, dy{"dy"}; + + auto make_pipeline = [&](const std::string &func_name, const std::string &rdom_name) { + RDom r(-radius, diameter, -radius, diameter, rdom_name); + Func f{func_name}; + f(x, y) = 0.0f; + f(x, y) += kernel(r.x) * kernel(r.y) * input(x + r.x, y + r.y); + return std::make_pair(f, r); + }; + + auto [reference, ref_r] = make_pipeline("gaussian_reference", "ref_r"); + auto [blur, r] = make_pipeline("gaussian_blur", "r"); + + Func vertical_partials = blur.update().rfactor(r.y, dy); + Func horizontal = vertical_partials.update().hoist_invariants(); + + vertical_partials.compute_root(); + horizontal.compute_root(); + + Buffer expected = reference.realize({width, height}); + Buffer actual = blur.realize({width, height}); + for (int yy = 0; yy < height; yy++) { + for (int xx = 0; xx < width; xx++) { + const float error = std::abs(actual(xx, yy) - expected(xx, yy)); + internal_assert(error < 1e-5f) + << "separable Gaussian after rfactor/hoist_invariants mismatch at (" + << xx << ", " << yy << "): " << actual(xx, yy) + << " vs ref " << expected(xx, yy) << "\n"; + } + } + + return 0; +} + +#if HALIDE_WITH_EXCEPTIONS +// 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, so hoist_invariants() must refuse rather than silently miscompile. +int hoist_invariants_invalid_law_rejected_test() { + if (!Halide::exceptions_enabled()) { + return 0; + } + + Buffer data(2); + data(0) = 1; + data(1) = 10; + + RDom k(0, 2, "k"); + + Func f{"f"}; + f() = UInt(8).max(); + f() = min(f(), cast(250) + data(k)); + + bool error = false; + try { + f.update().hoist_invariants(); + } 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 " + + f.name() + "."; + if (string(e.what()).find(expected) == string::npos) { + printf("Unexpected error for unsigned min hoisting:\n%s\n", e.what()); + return 1; + } + } + if (!error) { + printf("hoist_invariants should have rejected the unsigned min law!\n"); + return 1; + } + return 0; +} + +// hoist_invariants() errors when there is no distributable invariant factor to +// hoist, rather than silently behaving like a plain rfactor(). +int hoist_invariants_nothing_to_hoist_rejected_test() { + if (!Halide::exceptions_enabled()) { + return 0; + } + + ImageParam data_p{Int(32), 2, "data_p"}; + Var i{"i"}; + RDom k(0, 8, "k"); + + Func f{"f"}; + f(i) = 0; + f(i) += data_p(i, k); + + bool error = false; + try { + f.update().hoist_invariants(); + } 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 " + + f.name() + "."; + if (string(e.what()).find(expected) == string::npos) { + printf("Unexpected error when no invariant is hoistable:\n%s\n", e.what()); + return 1; + } + } + if (!error) { + printf("hoist_invariants should have errored when there is nothing to hoist!\n"); + return 1; + } + return 0; +} +#endif + } // namespace int main(int argc, char **argv) { @@ -1347,6 +1886,19 @@ 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}, + {"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 (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}, + {"hoist_invariants test (predicated RDom)", hoist_invariants_predicated_rdom_test}, + {"hoist_invariants test (after rfactor)", hoist_invariants_after_rfactor_test}, + {"hoist_invariants test (separable Gaussian after rfactor)", hoist_invariants_separable_gaussian_after_rfactor_test}, +#if HALIDE_WITH_EXCEPTIONS + {"hoist_invariants test (invalid law rejected)", hoist_invariants_invalid_law_rejected_test}, + {"hoist_invariants test (nothing to hoist rejected)", hoist_invariants_nothing_to_hoist_rejected_test}, +#endif }; using Sharder = Halide::Internal::Test::Sharder; diff --git a/test/performance/CMakeLists.txt b/test/performance/CMakeLists.txt index 30492c2a8514..ac90ddabbad9 100644 --- a/test/performance/CMakeLists.txt +++ b/test/performance/CMakeLists.txt @@ -29,6 +29,7 @@ tests( realize_overhead.cpp rgb_interleaved.cpp tiled_matmul.cpp + tiled_matmul_arm_neon.cpp vectorize.cpp wrap.cpp # keep-sorted end @@ -46,6 +47,7 @@ tests( parallel_scenarios.cpp profiler.cpp rfactor.cpp + separable_downsample.cpp sort.cpp stack_vs_heap.cpp thread_safe_jit_callable.cpp diff --git a/test/performance/separable_downsample.cpp b/test/performance/separable_downsample.cpp new file mode 100644 index 000000000000..6d4114a2c48e --- /dev/null +++ b/test/performance/separable_downsample.cpp @@ -0,0 +1,146 @@ +#include "Halide.h" +#include "halide_benchmark.h" + +#include +#include +#include +#include + +using namespace Halide; +using namespace Halide::Tools; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (target.arch == Target::WebAssembly) { + printf("[SKIP] Performance tests are meaningless and/or misleading under WebAssembly interpreter.\n"); + return 0; + } + + constexpr int stride = 2; + constexpr int radius = 7; + constexpr int diameter = 2 * radius + 1; + constexpr int output_width = 1024; + constexpr int output_height = 768; + + Buffer input(stride * output_width + 2 * radius, + stride * output_height + 2 * radius); + input.set_min(-radius, -radius); + for (int y = input.dim(1).min(); y <= input.dim(1).max(); y++) { + for (int x = input.dim(0).min(); x <= input.dim(0).max(); x++) { + input(x, y) = static_cast((7 * x + 13 * y + 1001) % 61 - 30); + } + } + + Buffer kernel(diameter); + kernel.set_min(-radius); + constexpr std::array binomial_weights = { + 1, 14, 91, 364, 1001, 2002, 3003, 3432, + 3003, 2002, 1001, 364, 91, 14, 1}; + for (int i = 0; i < diameter; i++) { + kernel(i - radius) = static_cast(binomial_weights[i]) / 16384.0f; + } + + Var x{"x"}, y{"y"}; + const int vec = target.natural_vector_size(); + + // The one-step implementation filters and downsamples both dimensions in a + // single 2D reduction. + RDom direct_r(-radius, diameter, -radius, diameter, "direct_r"); + Func direct{"direct_downsample"}; + direct(x, y) = 0.0f; + direct(x, y) += + kernel(direct_r.x) * kernel(direct_r.y) * + input(stride * x + direct_r.x, stride * y + direct_r.y); + + direct.compute_root() + .parallel(y) + .vectorize(x, vec); + direct.update() + .reorder(x, direct_r.x, direct_r.y, y) + .vectorize(x, vec) + .unroll(direct_r.x) + .unroll(direct_r.y); + + // Start with the same one-step definition. Preserving r.x leaves r.y as the + // reduction dimension of vertical_weighted. kernel(r.x) is then invariant + // over that reduction, so hoist_invariants() produces: + // + // vertical(x, y, dx) = + // sum_ry kernel(ry) * input(2*x + dx, 2*y + ry) + // separated(x, y) = + // sum_dx kernel(dx) * vertical(x, y, dx) + // + // Thus the first stage filters and downsamples in y, and the second stage + // filters and downsamples in x. + RDom r(-radius, diameter, -radius, diameter, "r"); + Func separated{"separable_downsample"}; + separated(x, y) = 0.0f; + separated(x, y) += + kernel(r.x) * kernel(r.y) * + input(stride * x + r.x, stride * y + r.y); + + Var dx{"dx"}; + Func vertical_weighted = separated.update().rfactor(r.x, dx); + Func vertical = vertical_weighted.update().hoist_invariants(); + + separated.compute_root() + .parallel(y) + .vectorize(x, vec); + separated.update() + .reorder(x, r.x, y) + .vectorize(x, vec) + .unroll(r.x); + + // Compute one vertical pass for the current x-kernel offset, consume it + // across the output row, then move to the next offset. + vertical_weighted.compute_at(separated, r.x) + .vectorize(x, vec); + vertical_weighted.update() + .reorder(x, dx, y) + .vectorize(x, vec); + + vertical.compute_at(vertical_weighted, x) + .vectorize(x, vec); + vertical.update() + .reorder(x, r.y, dx, y) + .vectorize(x, vec) + .unroll(r.y); + + Buffer direct_output(output_width, output_height); + Buffer separated_output(output_width, output_height); + + // Warm up and JIT both variants before timing. + direct.realize(direct_output, target); + separated.realize(separated_output, target); + + for (int yy = 0; yy < output_height; yy++) { + for (int xx = 0; xx < output_width; xx++) { + const float ref = direct_output(xx, yy); + const float actual = separated_output(xx, yy); + const float tolerance = 1e-5f * std::max(1.0f, std::abs(ref)); + if (std::abs(actual - ref) > tolerance) { + printf("Separable downsample mismatch at (%d, %d): %f vs direct %f\n", + xx, yy, actual, ref); + return 1; + } + } + } + + const double direct_time = benchmark([&] { + direct.realize(direct_output, target); + }); + const double separated_time = benchmark([&] { + separated.realize(separated_output, target); + }); + + printf("2x downsample with %dx%d separable kernel\n" + "Direct 2D reduction: %0.4f ms\n" + "rfactor + hoist_invariants: %0.4f ms (%0.2fx vs direct)\n", + diameter, diameter, + direct_time * 1000, + separated_time * 1000, + direct_time / separated_time); + + printf("Success!\n"); + return 0; +} diff --git a/test/performance/tiled_matmul_arm_neon.cpp b/test/performance/tiled_matmul_arm_neon.cpp new file mode 100644 index 000000000000..807049e1e7b3 --- /dev/null +++ b/test/performance/tiled_matmul_arm_neon.cpp @@ -0,0 +1,222 @@ +#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 by composing +// four directives: eager_inline(), hoist_invariants(), rfactor(), and +// change_type(). +// +// 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. The Hoisted variant (a) eager_inline()s Wt and +// VecDq into Acc's update so the scale factors buried inside them -- WtScale(i) +// and VecScale -- surface as explicit leaves of the update expression, which +// (b) hoist_invariants() then lifts out of the reduction as the invariant +// product WtScale(i) * VecScale, (c) rfactor() splits the now scale-free +// reduction by block, and (d) change_type(Int(32)) retypes the innermost +// per-block dot product to Int(32), yielding a pure i32(widening_mul(i8, i8)) +// accumulation that CodeGen_ARM matches to SDOT. PlainRfactor uses the same +// split/atomic/vectorize schedule +// but without hoisting or retyping, 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"}, j{"j"}, 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 Wt("Wt"); + Wt(i, j) = WtScale(i) * WtPacked(j % reduce, i, j / reduce); + + Func VecDq("VecDq"); + VecDq(i) = VecScale * Vec(i); + + Func Acc("Acc"); + Acc(i) = 0.0f; + Acc(i) += Wt(i, r) * VecDq(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) { + // Compose the directives. hoist_invariants() only analyzes Acc's + // update expression; it can't see inside the Funcs that expression + // calls. As written, Acc(i) += Wt(i, r) * VecDq(r) is a product of two + // opaque Calls that both depend on r, so the scale factors -- which + // are invariant across r -- are invisible, hidden in the definitions + // of Wt and VecDq. + // + // eager_inline(Wt, VecDq), called on Acc's update stage, substitutes + // those definitions into that update now, at schedule time, so it + // becomes + // Acc(i) += (WtScale(i) * WtPacked(...)) * (VecScale * Vec(r)). + // WtScale(i) (depends only on i) and VecScale (a Param) are now + // explicit leaves of the update expression, so hoist_invariants() can + // recognize them as loop-invariant and lift their product out of the + // reduction: Acc_wb accumulates the scale-free product over all of K + // and Acc's own update collapses to one multiply per row, + // Acc(i) = WtScale(i) * VecScale * Acc_wb(i). Acc_wb still accumulates + // at Float(32) at this point. + Func Acc_wb = Acc.update().eager_inline(Wt, VecDq).hoist_invariants(); + + // Second, factor Acc_wb's own (now scale-free) reduction by block. + // 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. + Func Acc_dot = Acc_wb.update().rfactor(ro, u); + + // Third, retype the innermost per-block dot product to Int(32). + // change_type() proves K int8 x int8 terms can't overflow Int(32), + // rewrites Acc_dot into an Int(32) accumulation (a pure + // i32(widening_mul(i8, i8)) dot product that CodeGen_ARM matches to + // SDOT), and leaves the original Acc_dot as an inline cast back to + // Float(32) so Acc_wb is unaffected. Schedule the returned Int(32) + // Func -- it holds the real reduction now. + Func Acc_dot_i32 = Acc_dot.change_type(Int(32)); + + // Compute a whole panel of rows through every stage before moving to + // the next panel: Acc and Acc_wb are computed_at the panel loop, and + // Acc_dot_i32 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 in registers. + Acc_wb.compute_at(Result, io) + .vectorize(i, panel) + .update() + .vectorize(i, panel); + + // Acc_dot_i32'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_i32.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); + + 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; +}