diff --git a/python_bindings/src/halide/halide_/PyFunc.cpp b/python_bindings/src/halide/halide_/PyFunc.cpp index 9daa3fae3e6f..b6a6b3ce48b5 100644 --- a/python_bindings/src/halide/halide_/PyFunc.cpp +++ b/python_bindings/src/halide/halide_/PyFunc.cpp @@ -219,6 +219,8 @@ void define_func(py::module &m) { .def("bound_storage", &Func::bound_storage) .def("memoize", &Func::memoize, py::arg("eviction_key") = EvictionKey()) .def("compute_inline", &Func::compute_inline) + .def("eager_inline", &Func::eager_inline, py::arg("fs")) + .def("change_type", &Func::change_type, py::arg("type"), py::arg("unsafe") = false) .def("compute_root", &Func::compute_root) .def("store_root", &Func::store_root) diff --git a/python_bindings/src/halide/halide_/PyStage.cpp b/python_bindings/src/halide/halide_/PyStage.cpp index 93629beb90a2..47071ed0af57 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("split_vars", [](const Stage &stage) -> py::list { auto vars = stage.split_vars(); diff --git a/src/AddTypeChangeChecks.cpp b/src/AddTypeChangeChecks.cpp new file mode 100644 index 000000000000..a9c3d72609fc --- /dev/null +++ b/src/AddTypeChangeChecks.cpp @@ -0,0 +1,35 @@ +#include "AddTypeChangeChecks.h" +#include "Function.h" +#include "IR.h" +#include "IROperator.h" +#include "Schedule.h" +#include "Simplify.h" + +namespace Halide { +namespace Internal { + +Stmt add_type_change_checks(const Stmt &s, const std::map &env) { + std::vector stmts; + + for (const auto &p : env) { + const Function &f = p.second; + for (const auto &[condition, message] : f.schedule().type_change_checks()) { + if (!condition.defined()) { + continue; + } + Expr proven = simplify(condition); + if (is_const_one(proven)) { + // Statically proven; no runtime check needed. + continue; + } + Expr error = requirement_failed_error(condition, {Expr(message)}); + stmts.push_back(AssertStmt::make(condition, error)); + } + } + + stmts.push_back(s); + return Block::make(stmts); +} + +} // namespace Internal +} // namespace Halide diff --git a/src/AddTypeChangeChecks.h b/src/AddTypeChangeChecks.h new file mode 100644 index 000000000000..194830a7f6b7 --- /dev/null +++ b/src/AddTypeChangeChecks.h @@ -0,0 +1,29 @@ +#ifndef HALIDE_ADD_TYPE_CHANGE_CHECKS_H +#define HALIDE_ADD_TYPE_CHANGE_CHECKS_H + +/** \file + * Defines the lowering pass that injects the overflow-safety preconditions + * recorded by Func::change_type() into the pipeline's assertion block. + */ + +#include +#include + +#include "Expr.h" + +namespace Halide { +namespace Internal { + +class Function; + +/** Prepend assertions for any static preconditions that Func::change_type() + * recorded on the funcs in `env` (that it could not discharge at schedule time, + * e.g. because a reduction extent was symbolic). Statically-true conditions are + * dropped. Like the other check passes, the resulting asserts are removed later + * when the no_asserts target feature is set. */ +Stmt add_type_change_checks(const Stmt &s, const std::map &env); + +} // namespace Internal +} // namespace Halide + +#endif 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/CMakeLists.txt b/src/CMakeLists.txt index 1e2229eca2c7..d5862b51d3f4 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -57,6 +57,7 @@ target_sources( AddImageChecks.h AddParameterChecks.h AddSplitFactorChecks.h + AddTypeChangeChecks.h AlignLoads.h AllocationBoundsInference.h ApplySplit.h @@ -239,6 +240,7 @@ target_sources( AddImageChecks.cpp AddParameterChecks.cpp AddSplitFactorChecks.cpp + AddTypeChangeChecks.cpp AlignLoads.cpp AllocationBoundsInference.cpp ApplySplit.cpp diff --git a/src/Func.cpp b/src/Func.cpp index 659efac305a4..680acaa100d2 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -14,6 +15,7 @@ #include "Associativity.h" #include "Callable.h" #include "CodeGen_LLVM.h" +#include "ConstantBounds.h" #include "Debug.h" #include "ExprUsesVar.h" #include "FindCalls.h" @@ -21,10 +23,13 @@ #include "Function.h" #include "IR.h" #include "IREquality.h" +#include "IRMatch.h" #include "IRMutator.h" #include "IROperator.h" #include "IRPrinter.h" +#include "IRVisitor.h" #include "ImageParam.h" +#include "Inline.h" #include "LLVM_Output.h" #include "Lower.h" #include "Param.h" @@ -573,55 +578,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}}); @@ -630,6 +586,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); @@ -682,17 +661,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 +747,172 @@ 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. +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. 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. `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; + } + + 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])) { + // 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, intermediate_vars, + to_intermediate, to_preserved); + } + } + } + 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) { @@ -854,7 +1009,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_impl(const vector> &preserved, bool hoist_invariant_factor) { user_assert(!definition.is_init()) << "rfactor() must be called on an update definition\n"; definition.schedule().touched() = true; @@ -918,6 +1073,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 +1084,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 +1097,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,6 +1106,29 @@ 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 (hoist_invariant_factor) { + // 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); + 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"; + } + // Intermediate pure definition { vector args = dim_vars_exprs; @@ -963,13 +1140,24 @@ 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); + // 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. substitute_self_reference then retargets the self-reference + // 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 +1191,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 +1207,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]); 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 +1252,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; } @@ -1073,6 +1262,123 @@ Func Stage::rfactor(const vector> &preserved) { return intm; } +Func Stage::rfactor(const vector> &preserved) { + return rfactor_impl(preserved, /*hoist_invariant_factor=*/false); +} + +Func Stage::hoist_invariants() { + definition.schedule().touched() = true; + return rfactor_impl({}, /*hoist_invariant_factor=*/true); +} + +// Helpers for change_type implementation +namespace { + +// Does `e` contain a direct Halide call to `fname` (a self-reference)? +bool contains_self_reference(const Expr &e, const string &fname) { + class Finder : public IRGraphVisitor { + using IRGraphVisitor::visit; + const string &fname; + void visit(const Call *c) override { + if (c->call_type == Call::Halide && c->name == fname) { + found = true; + } + IRGraphVisitor::visit(c); + } + + public: + bool found = false; + explicit Finder(const string &f) + : fname(f) { + } + } finder(fname); + e.accept(&finder); + return finder.found; +} + +// Rewrite a factor-free leaf expression `e` to type `t`, preferring integer +// forms over float round-trips so patterns like widening_mul survive to +// instruction selection. This is the type-narrowing logic previously wired into +// hoisted rfactor, now driven by an explicit target type. +Expr retype_leaf(const Expr &e, Type t) { + if (e.type() == t) { + return e; + } + + // Expose a single promotion cast for products/sums/diffs of two same-width + // integer operands, e.g. cast(a) * cast(b) == cast(widening_mul(a, b)). + Expr folded = e; + { + Expr a8i = Variable::make(Int(8), "a"), b8i = Variable::make(Int(8), "b"); + Expr a8u = Variable::make(UInt(8), "a"), b8u = Variable::make(UInt(8), "b"); + Expr a16i = Variable::make(Int(16), "a"), b16i = Variable::make(Int(16), "b"); + Expr a16u = Variable::make(UInt(16), "a"), b16u = Variable::make(UInt(16), "b"); + std::map matches; + const std::pair patterns[] = { + {cast(Float(32), a8i) * cast(Float(32), b8i), cast(Float(32), widening_mul(a8i, b8i))}, + {cast(Float(32), a8u) * cast(Float(32), b8u), cast(Float(32), widening_mul(a8u, b8u))}, + {cast(Float(32), a16i) * cast(Float(32), b16i), cast(Float(32), widening_mul(a16i, b16i))}, + {cast(Float(32), a16u) * cast(Float(32), b16u), cast(Float(32), widening_mul(a16u, b16u))}, + {cast(Float(32), a8i) + cast(Float(32), b8i), cast(Float(32), widening_add(a8i, b8i))}, + {cast(Float(32), a8u) + cast(Float(32), b8u), cast(Float(32), widening_add(a8u, b8u))}, + {cast(Float(32), a8i) - cast(Float(32), b8i), cast(Float(32), widening_sub(a8i, b8i))}, + {cast(Float(32), a8u) - cast(Float(32), b8u), cast(Float(32), widening_sub(a8u, b8u))}, + }; + for (const auto &[pattern, result] : patterns) { + if (expr_match(pattern, folded, matches)) { + folded = substitute(matches, result); + break; + } + } + } + + // Strip an outer int->float promotion cast, widening the integer to the + // target's width if needed (a strict_cast is a Call, not a Cast, and is left + // alone). This keeps the accumulation at an integer type. + if (const Cast *c = folded.as()) { + if (folded.type().is_float() && c->value.type().is_int_or_uint() && + t.is_int_or_uint() && t.bits() >= c->value.type().bits()) { + return (c->value.type() == t) ? c->value : cast(t, c->value); + } + } + + return cast(t, folded); +} + +// Retype a whole definition value to type `t`, retargeting self-references from +// `fname` to `dst` and pushing casts down to the increment leaves. Only reduction +// updates shaped as a tree of binary combiners over a self-reference and an +// increment are supported. +Expr retype_value(const Expr &e, const string &fname, const Function &dst, Type t) { + if (const Call *c = e.as()) { + if (c->call_type == Call::Halide && c->name == fname) { + return Call::make(dst, c->args, c->value_index); + } + } + if (contains_self_reference(e, fname)) { + optional> operands = as_binary_operands(e); + user_assert(operands) + << "change_type() only supports update definitions built from binary " + << "operators over the accumulator; " << fname << " has an unsupported shape.\n"; + return make_binary_op(e.node_type(), + retype_value(operands->first, fname, dst, t), + retype_value(operands->second, fname, dst, t)); + } + return retype_leaf(e, t); +} + +// The top-level associative combiner of a (let-stripped) reduction update value, +// i.e. the node type of the binary op whose operands are the self-reference and +// the increment. Returns nullopt if `val` isn't such a shape. +optional reduction_op(const Expr &val, const string &fname) { + optional> split = select_binary_operand(val, val.node_type(), [&](const Expr &e) { + return contains_self_reference(e, fname); + }); + return split ? std::make_optional(val.node_type()) : std::nullopt; +} + +} // namespace + 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"; @@ -3144,6 +3450,202 @@ Func &Func::compute_inline() { return compute_at(LoopLevel::inlined()); } +Func &Func::eager_inline(const std::vector &fs) { + invalidate_cache(); + for (const Func &f : fs) { + user_assert(f.defined()) + << "eager_inline() was passed an undefined Func.\n"; + user_assert(f.function().can_be_inlined()) + << "eager_inline() cannot inline " << f.name() + << ": it must be a pure Func with no update or extern definition and " + << "no specializations.\n"; + // Rewrites this Func's pure and update definitions in place, replacing + // every direct call to f with f's body. Processing fs left to right means + // a body spliced in by an earlier inline exposes its own direct calls to + // later fs, which the next iteration then inlines. + Internal::inline_function(func, f.function()); + } + return *this; +} + +namespace { + +// Prove that computing `typed`'s reduction at type `t` cannot overflow. Returns +// true if it is safe; if safety can only be guaranteed under a runtime +// precondition, that condition is returned in *condition. Returns false with an +// explanation in *why if it cannot be proven. +bool change_type_prove_safe(const Func &typed, Type t, Expr *condition, string *why) { + *condition = Expr(); + const Function fn = typed.function(); + const ConstantInterval limit = ConstantInterval::bounds_of_type(t); + + auto fits = [&](const ConstantInterval &ci) { + return ci.min_defined && ci.max_defined && limit.min <= ci.min && ci.max <= limit.max; + }; + + // Pure / identity values must be representable at the new type. + for (const Expr &v : fn.values()) { + if (!fits(constant_integer_bounds(v))) { + *why = "the initial value may not be representable in the target type"; + return false; + } + } + + for (const Definition &def : fn.updates()) { + Expr val = substitute_in_all_lets(def.values()[0]); + optional op = reduction_op(val, fn.name()); + + // The increment is the non-self-reference operand of the combiner. + Expr increment = val; + if (op) { + optional> split = select_binary_operand(val, *op, [&](const Expr &e) { + return contains_self_reference(e, fn.name()); + }); + if (split) { + increment = split->second; + } + } + const ConstantInterval term = constant_integer_bounds(increment); + + if (op && *op == IRNodeType::Add) { + // Bound the accumulator by (number of terms) x (per-term range). + int64_t n_max = 1; + bool symbolic = false; + Expr n_terms = make_const(Int(64), 1); + for (const auto &rv : def.schedule().rvars()) { + n_terms = simplify(n_terms * cast(Int(64), rv.extent)); + // Only a literal extent is known at compile time; a symbolic + // extent (e.g. an ImageParam dimension) gets only type-based + // bounds, which we must not treat as a static bound. + if (optional ext = as_const_int(simplify(rv.extent)); ext && *ext >= 0) { + n_max *= *ext; + } else { + symbolic = true; + } + } + if (!symbolic) { + if (fits(term * ConstantInterval(0, n_max))) { + continue; + } + *why = "the accumulated sum may exceed the target type's range"; + return false; + } + // Symbolic term count: emit a runtime precondition instead. + if (!term.min_defined || !term.max_defined) { + *why = "the per-term magnitude is unbounded"; + return false; + } + Expr cond = (make_const(Int(64), term.max) * n_terms <= make_const(Int(64), limit.max)) && + (make_const(Int(64), term.min) * n_terms >= make_const(Int(64), limit.min)); + *condition = condition->defined() ? (*condition && cond) : cond; + continue; + } + + // min / max / and / or (and unrecognized shapes): the result stays within + // a single term's range, so no accumulation overflow is possible. + if (fits(term)) { + continue; + } + *why = "a term may not be representable in the target type"; + return false; + } + return true; +} + +} // namespace + +Func Func::change_type(Type t, bool unsafe) { + user_assert(defined()) << "change_type() called on undefined Func.\n"; + user_assert(!func.has_extern_definition()) + << "change_type() cannot be applied to the extern Func " << name() << ".\n"; + user_assert(outputs() == 1) + << "change_type() currently supports only single-output Funcs, but " + << name() << " has " << outputs() << " outputs.\n"; + + invalidate_cache(); + + const Type old_t = func.output_types()[0]; + if (old_t == t) { + return *this; + } + + const string fname = func.name(); + const vector pure_vars = args(); + const vector pure_arg_exprs(pure_vars.begin(), pure_vars.end()); + + // Determine the reduction op (if any), so min/max accumulations get the + // right identity at the new type rather than a lossy cast of e.g. +inf. + optional op; + if (func.has_update_definition()) { + op = reduction_op(substitute_in_all_lets(func.update(0).values()[0]), fname); + } + const bool is_min_max = op && (*op == IRNodeType::Min || *op == IRNodeType::Max); + + // Build the retyped clone. + Func typed(fname + "_typed"); + + // Pure definition. + { + vector retyped; + for (const Expr &v : func.values()) { + if (is_min_max) { + optional id = get_associative_identity(t, *op); + user_assert(id) << "change_type() could not find an identity for " + << IRNodeType_string(*op) << " at type " << t << ".\n"; + retyped.push_back(*id); + } else { + retyped.push_back(retype_leaf(v, t)); + } + } + // Single-output only (asserted above), so there is exactly one value. + typed(pure_vars) = retyped[0]; + } + + // Update definitions. The retyped values still reference the original + // reduction domain, so pass a default domain and let define_update discover + // it from the values (passing a freshly-built one would trip its identity + // check). + for (size_t u = 0; u < func.updates().size(); u++) { + const Definition &def = func.update(u); + vector vals; + vals.reserve(def.values().size()); + for (const Expr &v : def.values()) { + vals.push_back(retype_value(substitute_in_all_lets(v), fname, typed.function(), t)); + } + typed.function().define_update(def.args(), vals, ReductionDomain{}); + typed.function().update(u).schedule() = def.schedule().get_copy(); + } + + // Safety check. + if (!unsafe && t.is_int_or_uint()) { + Expr condition; + string why; + const bool ok = change_type_prove_safe(typed, t, &condition, &why); + user_assert(ok) + << "change_type(" << t << ") on " << fname << " may overflow: " << why << ".\n" + << "Pass unsafe=true to change_type() to bypass this check.\n"; + if (condition.defined()) { + std::ostringstream msg; + msg << "change_type(" << t << ") on " << fname + << " requires the reduction extent to be small enough not to overflow"; + typed.function().schedule().type_change_checks().emplace_back(condition, msg.str()); + } + } + + // Rewrite this Func into an inline cast-back wrapper of the retyped clone, so + // that every existing consumer keeps seeing the original type. + const Expr wrapped = cast(old_t, Call::make(typed.function(), pure_arg_exprs, 0)); + vector arg_names; + arg_names.reserve(pure_vars.size()); + for (const Var &v : pure_vars) { + arg_names.push_back(v.name()); + } + func.clear_definition(); + func.define(arg_names, {wrapped}); + + return typed; +} + Func &Func::trace_loads() { invalidate_cache(); func.trace_loads(); diff --git a/src/Func.h b/src/Func.h index 000b44561efa..295a4ca84ed7 100644 --- a/src/Func.h +++ b/src/Func.h @@ -18,6 +18,7 @@ #include "Var.h" #include +#include #include namespace Halide { @@ -92,6 +93,11 @@ class Stage { std::pair, std::vector> rfactor_validate_args(const std::vector> &preserved, const Internal::AssociativeOp &prover_result); + /** Shared implementation of rfactor() and hoist_invariants(). When + * `hoist_invariant_factor` is set, a distributive loop-invariant factor is + * hoisted from the intermediate reduction to the write-back step. */ + Func rfactor_impl(const std::vector> &preserved, bool hoist_invariant_factor); + public: Stage(Internal::Function f, Internal::Definition d, size_t stage_index) : function(std::move(f)), definition(std::move(d)), stage_index(stage_index) { @@ -143,54 +149,93 @@ 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 */ // @{ Func rfactor(const std::vector> &preserved); Func rfactor(const RVar &r, const Var &v); // @} + /** 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. Unlike rfactor(), hoist_invariants() does not take a list of + * preserved RVars (call rfactor() on the returned intermediate to split its + * reduction further) and it does not change any types (see change_type()). + * + * 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 @@ -2613,6 +2658,52 @@ class Func { */ Func &compute_inline(); + /** Change the type at which this Func computes and stores its values, + * subject to a reduction-aware safety check. + * + * This splits the Func in two: a new intermediate Func (returned) that + * copies this Func's definitions but accumulates at the requested type `t` + * (with the appropriate casts inserted, preferring integer forms such as + * widening_mul over float round-trips), and this Func, which is rewritten in + * place into an inline wrapper that casts the intermediate's result back to + * the original type. Every existing consumer therefore keeps seeing the + * original type, while the returned intermediate can be scheduled to exploit + * the new type (e.g. an Int(32) accumulator eligible for dot-product + * instructions). Schedule the returned Func to control the retyped + * computation. + * + * The change is validated with the bounds machinery: for an integer target, + * change_type() proves the accumulation cannot overflow by combining the + * per-term value range with the reduction extent. If it can only be + * guaranteed under a runtime precondition (e.g. the RDom extent isn't too + * wide), that precondition is injected into the pipeline's assertion block + * (and removed by the no_asserts target feature). If safety cannot be + * established, change_type() errors unless `unsafe` is true, which bypasses + * the check entirely. + * + * Currently supports single-output Funcs whose update definitions are built + * from binary operators over the accumulator. */ + Func change_type(Type t, bool unsafe = false); + + /** Immediately inline direct calls to each of the given Funcs into this + * Func's definitions, processed left to right so that inlining an earlier + * Func can expose direct calls to a later one (e.g. when an earlier Func's + * body itself calls a later one). + * + * Unlike compute_inline(), which merely marks a Func to be inlined during + * lowering, eager_inline() performs the substitution now, at schedule time, + * rewriting this Func's pure and update definitions in place. This is useful + * to surface structure that other schedule-time directives need to see -- + * for example, exposing a loop-invariant factor buried inside a call, so that + * hoist_invariants() can then hoist it out of a reduction like + * h(x) += f(x) * g(x). + * + * Each inlined Func must be inlinable: a pure Func (no update or extern + * definition) with no specializations, and with a schedule compatible with + * inlining (as for compute_inline()). The inlined Funcs are otherwise + * unchanged; only this Func's calls to them are replaced. */ + Func &eager_inline(const std::vector &fs); + /** Get a handle on an update step for the purposes of scheduling * it. */ Stage update(int idx = 0); diff --git a/src/Function.cpp b/src/Function.cpp index 36dac5d1fda6..5849771e4713 100644 --- a/src/Function.cpp +++ b/src/Function.cpp @@ -669,6 +669,21 @@ void Function::define(const vector &args, vector values) { } } +void Function::clear_definition() { + contents->output_types.clear(); + contents->args.clear(); + contents->func_schedule = FuncSchedule(); + contents->init_def = Definition(); + contents->updates.clear(); + contents->output_buffers.clear(); + contents->extern_arguments.clear(); + contents->extern_function_name.clear(); + contents->extern_mangling = NameMangling::Default; + contents->extern_function_device_api = DeviceAPI::Host; + contents->extern_proxy_expr = Expr(); + contents->frozen = false; +} + void Function::create_output_buffers(const std::vector &types, int dims) const { internal_assert(contents->output_buffers.empty()); internal_assert(!types.empty() && dims != AnyDims); diff --git a/src/Function.h b/src/Function.h index 7a9fce63163d..0006cb1d9b78 100644 --- a/src/Function.h +++ b/src/Function.h @@ -119,6 +119,13 @@ class Function { * reduction domain */ void define(const std::vector &args, std::vector values); + /** Reset this Function to an undefined state in place (clearing all pure, + * update, and extern definitions, output types/buffers, and schedule) while + * preserving the Function's object identity, so existing references to it + * remain valid and it can be given a fresh definition with define(). Used by + * Func::change_type() to turn the original Func into an inline wrapper. */ + void clear_definition(); + /** Add an update definition to this function. It must already have a pure * definition but not an update definition, and the length of args must * match the length of args used in the pure definition. 'value' may depend diff --git a/src/IROperator.cpp b/src/IROperator.cpp index ffd33122d741..1a00a994d8e5 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/Inline.cpp b/src/Inline.cpp index 70e936f07898..02624ce2fe28 100644 --- a/src/Inline.cpp +++ b/src/Inline.cpp @@ -189,7 +189,10 @@ class Inliner : public IRMutator { Inliner(const Function &f) : func(f) { internal_assert(f.can_be_inlined()) << "Illegal to inline " << f.name() << "\n"; - validate_schedule_inlined_function(f); + // Note: the inlined function's schedule is validated at lowering time in + // ScheduleFunctions (schedule_functions), where loop levels are locked. + // Validating here too is redundant and breaks schedule-time inlining + // (e.g. Func::eager_inline), where loop levels are not yet locked. } }; diff --git a/src/Lower.cpp b/src/Lower.cpp index cd7ccc9a03f4..cf89410dbfb3 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -11,6 +11,7 @@ #include "AddImageChecks.h" #include "AddParameterChecks.h" #include "AddSplitFactorChecks.h" +#include "AddTypeChangeChecks.h" #include "AllocationBoundsInference.h" #include "AsyncProducers.h" #include "BoundConstantExtentLoops.h" @@ -215,6 +216,10 @@ void lower_impl(const vector &output_funcs, s = add_split_factor_checks(s, env); log("Lowering after asserting that all split factors are positive:", s); + debug(1) << "Asserting change_type() accumulations cannot overflow...\n"; + s = add_type_change_checks(s, env); + log("Lowering after asserting change_type() accumulations cannot overflow:", s); + debug(1) << "Removing extern loops...\n"; s = remove_extern_loops(s); log("Lowering after removing extern loops:", s); diff --git a/src/Schedule.cpp b/src/Schedule.cpp index e3e5576e051d..2ac05c200714 100644 --- a/src/Schedule.cpp +++ b/src/Schedule.cpp @@ -244,6 +244,11 @@ struct FuncScheduleContents { // This is an extent of the ring buffer and expected to be a positive integer. Expr ring_buffer; Expr memoize_eviction_key; + // Static preconditions injected by change_type() that must hold for the + // retyped accumulation not to overflow. Each is a (condition, message) pair; + // a lowering pass turns them into assertions in the pipeline's initial + // assertion block (removed by the no_asserts target feature). + std::vector> type_change_checks; FuncScheduleContents() : store_level(LoopLevel::inlined()), compute_level(LoopLevel::inlined()), hoist_storage_level(LoopLevel::inlined()) { @@ -279,6 +284,11 @@ struct FuncScheduleContents { b.remainder = mutator(b.remainder); } } + for (auto &check : type_change_checks) { + if (check.first.defined()) { + check.first = mutator(check.first); + } + } } }; @@ -369,6 +379,7 @@ FuncSchedule FuncSchedule::deep_copy( copy.contents->memoize_eviction_key = contents->memoize_eviction_key; copy.contents->async = contents->async; copy.contents->ring_buffer = contents->ring_buffer; + copy.contents->type_change_checks = contents->type_change_checks; // Deep-copy wrapper functions. for (const auto &iter : contents->wrappers) { @@ -420,6 +431,14 @@ Expr &FuncSchedule::ring_buffer() const { return contents->ring_buffer; } +const std::vector> &FuncSchedule::type_change_checks() const { + return contents->type_change_checks; +} + +std::vector> &FuncSchedule::type_change_checks() { + return contents->type_change_checks; +} + std::vector &FuncSchedule::storage_dims() { return contents->storage_dims; } diff --git a/src/Schedule.h b/src/Schedule.h index 8ce415df9497..d50240bad0d7 100644 --- a/src/Schedule.h +++ b/src/Schedule.h @@ -617,6 +617,16 @@ class FuncSchedule { Expr &ring_buffer(); Expr &ring_buffer() const; + /** Static preconditions injected by Func::change_type() that guarantee the + * retyped accumulation cannot overflow. Each entry is a (condition, message) + * pair; a lowering pass (add_type_change_checks) asserts them in the + * pipeline's initial assertion block, and they are removed by the no_asserts + * target feature. */ + // @{ + const std::vector> &type_change_checks() const; + std::vector> &type_change_checks(); + // @} + /** The list and order of dimensions used to store this * function. The first dimension in the vector corresponds to the * innermost dimension for storage (i.e. which dimension is 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/CMakeLists.txt b/test/correctness/CMakeLists.txt index 3a05d2e9bece..c43166c155ee 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -42,6 +42,7 @@ tests( cascaded_filters.cpp cast.cpp cast_handle.cpp + change_type.cpp chunk.cpp chunk_sharing.cpp circular_reference_leak.cpp @@ -94,6 +95,7 @@ tests( downsampling_reduce.cpp dynamic_allocation_in_gpu_kernel.cpp dynamic_reduction_bounds.cpp + eager_inline.cpp early_out.cpp embed_bitcode.cpp erf.cpp diff --git a/test/correctness/change_type.cpp b/test/correctness/change_type.cpp new file mode 100644 index 000000000000..f041acf92476 --- /dev/null +++ b/test/correctness/change_type.cpp @@ -0,0 +1,301 @@ +#include "Halide.h" +#include + +using namespace Halide; + +namespace { + +// change_type() retypes a Func's accumulation. Applied to a float dot-product +// intermediate produced by hoist_invariants(), it exposes an Int(32) +// accumulation (widening_mul-friendly) while every consumer keeps seeing the +// original Float(32) type via the inline cast-back wrapper. +int change_type_basic_test() { + const int K = 64; + ImageParam A{Int(8), 1, "A"}, 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; + Acc(i) += Scale(i) * cast(widening_mul(A(r), B(r))); + + Func Acc_wb = Acc.update().hoist_invariants(); + internal_assert(Acc_wb.types()[0] == Float(32)) + << "expected the hoisted intermediate to stay Float(32)\n"; + + Func Acc_i32 = Acc_wb.change_type(Int(32)); + internal_assert(Acc_i32.types()[0] == Int(32)) + << "change_type: expected the retyped intermediate to be Int(32), got " + << Acc_i32.types()[0] << "\n"; + // The original intermediate remains Float(32) for its consumers. + internal_assert(Acc_wb.types()[0] == Float(32)) + << "change_type: the original Func must keep its type for consumers, got " + << Acc_wb.types()[0] << "\n"; + Acc_i32.compute_root(); + + const int M = 4; + Buffer a(K), b(K); + for (int k = 0; k < K; k++) { + a(k) = (int8_t)((k % 15) - 7); + b(k) = (int8_t)((k % 11) - 5); + } + Buffer s(M); + for (int m = 0; m < M; m++) { + s(m) = 0.5f * (m + 1); + } + A.set(a); + B.set(b); + Scale.set(s); + + Buffer result = Acc.realize({M}); + for (int m = 0; m < M; m++) { + int32_t dot = 0; + for (int k = 0; k < K; k++) { + dot += (int32_t)a(k) * (int32_t)b(k); + } + float expected = s(m) * (float)dot; + internal_assert(result(m) == expected) + << "change_type basic mismatch at " << m << ": " << result(m) + << " vs " << expected << "\n"; + } + return 0; +} + +// change_type() on a min reduction must use the reduction identity at the new +// type (Int(16).max()), not a lossy cast of the original Float(32) +inf. +int change_type_min_test() { + ImageParam offset_p{Float(32), 1, "offset_p"}; + ImageParam data_p{Int(16), 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_wb = C.update().hoist_invariants(); + Func C_i16 = C_wb.change_type(Int(16)); + internal_assert(C_i16.types()[0] == Int(16)) + << "change_type min: expected Int(16), got " << C_i16.types()[0] << "\n"; + C_i16.compute_root(); + + 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) = (int16_t)(((m * K + kk) % 21) - 10); + } + } + offset_p.set(off); + data_p.set(dat); + + Buffer result = C.realize({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)); + } + internal_assert(result(m) == v) + << "change_type min mismatch at " << m << ": " << result(m) << " vs " << v << "\n"; + } + return 0; +} + +// change_type() must not fold a strict_float() through into a widening integer +// product: the reduction below rounds each term to float before summing, so the +// result must observe that rounding (giving exactly 0). +int change_type_strict_float_test() { + Buffer data(2); + data(0) = 16777217; // not representable exactly as float32 + 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(); + // Retype to Int(32): the strict_float call must survive (it is not an + // implicit promotion cast), so each term is still rounded to float first. + Func intm_i = intm.change_type(Int(32), /*unsafe=*/true); + intm_i.compute_root(); + + Buffer result = f.realize(); + internal_assert(result() == 0.0f) + << "change_type strict_float mismatch: " << result() << " vs 0\n"; + return 0; +} + +// A symbolic reduction extent can't be bounded at schedule time, so change_type +// injects a runtime precondition. With a valid (small) extent it passes and the +// result is correct. +int change_type_symbolic_extent_test() { + ImageParam A{Int(8), 1, "A"}, B{Int(8), 1, "B"}; + + Var i{"i"}; + // The extent is a runtime value (an ImageParam dimension), so it can't be + // bounded at schedule time. + RDom r(0, A.dim(0).extent(), "r"); + + Func Acc{"Acc"}; + Acc(i) = 0.0f; + Acc(i) += cast(widening_mul(A(r), B(r))); + + // int8*int8 accumulated over a symbolic number of terms: change_type injects + // a runtime precondition guaranteeing the sum fits in Int(32). + Func Acc_i32 = Acc.change_type(Int(32)); + internal_assert(Acc_i32.types()[0] == Int(32)) + << "change_type symbolic: expected Int(32), got " << Acc_i32.types()[0] << "\n"; + Acc_i32.compute_root(); + + const int K = 100; + Buffer a(K), b(K); + for (int k = 0; k < K; k++) { + a(k) = (int8_t)((k % 9) - 4); + b(k) = (int8_t)((k % 7) - 3); + } + A.set(a); + B.set(b); + + Buffer result = Acc.realize({4}); + int32_t dot = 0; + for (int k = 0; k < K; k++) { + dot += (int32_t)a(k) * (int32_t)b(k); + } + for (int m = 0; m < 4; m++) { + float expected = (float)dot; + internal_assert(result(m) == expected) + << "change_type symbolic-extent mismatch at " << m << ": " << result(m) + << " vs " << expected << "\n"; + } + return 0; +} + +// change_type() can be applied more than once, retyping the intermediate +// returned by a previous change_type(). Each step must remain safe and correct. +int change_type_twice_test() { + const int K = 32; + ImageParam A{Int(8), 1, "A"}; + + Var i{"i"}; + RDom r(0, K, "r"); + + Func Acc{"Acc"}; + Acc(i) = 0.0f; + // Sum of K int8 values: |sum| <= 32 * 127 = 4064, which fits Int(16), so both + // retypes (Float(32) -> Int(32) -> Int(16)) are statically safe. + Acc(i) += cast(A(r)); + + Func Acc_i32 = Acc.change_type(Int(32)); + Func Acc_i16 = Acc_i32.change_type(Int(16)); + internal_assert(Acc.types()[0] == Float(32) && + Acc_i32.types()[0] == Int(32) && + Acc_i16.types()[0] == Int(16)) + << "change_type twice: unexpected types " + << Acc.types()[0] << " / " << Acc_i32.types()[0] << " / " << Acc_i16.types()[0] << "\n"; + Acc_i16.compute_root(); + Acc_i32.compute_root(); + + Buffer a(K); + for (int k = 0; k < K; k++) { + a(k) = (int8_t)((k % 15) - 7); + } + A.set(a); + + Buffer result = Acc.realize({2}); + int32_t sum = 0; + for (int k = 0; k < K; k++) { + sum += (int32_t)a(k); + } + for (int m = 0; m < 2; m++) { + internal_assert(result(m) == (float)sum) + << "change_type twice mismatch at " << m << ": " << result(m) << " vs " << sum << "\n"; + } + return 0; +} + +#if HALIDE_WITH_EXCEPTIONS +// change_type() into a type too narrow to hold the accumulated sum must be +// rejected unless the unsafe flag is set. +int change_type_overflow_rejected_test() { + if (!Halide::exceptions_enabled()) { + return 0; + } + + const int K = 256; + ImageParam A{Int(8), 1, "A"}, B{Int(8), 1, "B"}; + Var i{"i"}; + RDom r(0, K, "r"); + + auto build = [&]() { + Func Acc{"Acc"}; + Acc(i) = 0.0f; + Acc(i) += cast(i) * cast(widening_mul(A(r), B(r))); + return Acc; + }; + + // 256 terms up to 127*127 ~= 4.1M overflows Int(16). + bool error = false; + try { + Func Acc = build(); + Func wb = Acc.update().hoist_invariants(); + wb.change_type(Int(16)); + } catch (const Halide::CompileError &e) { + error = true; + printf("Expected overflow error:\n%s\n", e.what()); + } + if (!error) { + printf("change_type should have rejected the overflowing Int(16) accumulation!\n"); + return 1; + } + + // With unsafe=true, the same call is allowed (the user takes responsibility). + Func Acc = build(); + Func wb = Acc.update().hoist_invariants(); + Func typed = wb.change_type(Int(16), /*unsafe=*/true); + internal_assert(typed.types()[0] == Int(16)) + << "change_type unsafe: expected Int(16), got " << typed.types()[0] << "\n"; + return 0; +} +#endif + +} // namespace + +int main(int argc, char **argv) { + printf("Running change_type_basic_test\n"); + if (change_type_basic_test()) { + return 1; + } + printf("Running change_type_min_test\n"); + if (change_type_min_test()) { + return 1; + } + printf("Running change_type_strict_float_test\n"); + if (change_type_strict_float_test()) { + return 1; + } + printf("Running change_type_symbolic_extent_test\n"); + if (change_type_symbolic_extent_test()) { + return 1; + } + printf("Running change_type_twice_test\n"); + if (change_type_twice_test()) { + return 1; + } +#if HALIDE_WITH_EXCEPTIONS + printf("Running change_type_overflow_rejected_test\n"); + if (change_type_overflow_rejected_test()) { + return 1; + } +#endif + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/eager_inline.cpp b/test/correctness/eager_inline.cpp new file mode 100644 index 000000000000..5af4ef3d0f76 --- /dev/null +++ b/test/correctness/eager_inline.cpp @@ -0,0 +1,130 @@ +#include "Halide.h" +#include + +using namespace Halide; + +namespace { + +// eager_inline() inlines direct calls at schedule time (unlike compute_inline(), +// which defers to lowering), which lets later schedule-time directives see the +// exposed structure. Here an invariant factor S(x) is buried inside a call +// f(x, r); inlining f (and the g it calls) exposes it so hoist_invariants() can +// lift it out of the reduction. +int eager_inline_exposes_factor_test() { + const int K = 16, M = 4; + ImageParam S{Float(32), 1, "S"}; + ImageParam data{Float(32), 2, "data"}; + + Var x{"x"}, y{"y"}; + RDom r(0, K, "r"); + + Func g{"g"}, f{"f"}, h{"h"}; + g(x) = S(x); + f(x, y) = g(x) * data(x, y); + h(x) += f(x, r); + + // Inline f then g into h, left to right: inlining f exposes the call to g. + h.eager_inline({f, g}); + // Now h(x) += S(x) * data(x, r), so the invariant S(x) can be hoisted. + Func h_wb = h.update().hoist_invariants(); + h_wb.compute_root(); + + Buffer sbuf(M), dbuf(M, K); + for (int i = 0; i < M; i++) { + sbuf(i) = 0.5f * (i + 1); + for (int k = 0; k < K; k++) { + dbuf(i, k) = (float)((i + k) % 5 - 2); + } + } + S.set(sbuf); + data.set(dbuf); + + Buffer out = h.realize({M}); + for (int i = 0; i < M; i++) { + float ref = 0.f; + for (int k = 0; k < K; k++) { + ref += sbuf(i) * dbuf(i, k); + } + if (std::abs(out(i) - ref) > 1e-4f) { + printf("eager_inline exposes factor mismatch at %d: %f vs %f\n", i, out(i), ref); + return 1; + } + } + return 0; +} + +// eager_inline() performs the substitution immediately, so the caller's +// definition no longer references the inlined Funcs (they are inlined by value). +// Verify the numerics of a simple chained inline match a plain inlined pipeline. +int eager_inline_chain_test() { + Var x{"x"}; + Func a{"a"}, b{"b"}, c{"c"}; + a(x) = x + 1; + b(x) = a(x) * 2; // calls a + c(x) = b(x) + a(x); // calls b (which calls a) and a directly + + // Inline b then a into c. Inlining b splices in its call to a, which the + // subsequent inline of a then also folds. + c.eager_inline({b, a}); + + Buffer out = c.realize({8}); + for (int i = 0; i < 8; i++) { + int ref = (i + 1) * 2 + (i + 1); + if (out(i) != ref) { + printf("eager_inline chain mismatch at %d: %d vs %d\n", i, out(i), ref); + return 1; + } + } + return 0; +} + +#if HALIDE_WITH_EXCEPTIONS +// A Func with an update definition is not inlinable, so eager_inline() rejects it. +int eager_inline_non_pure_rejected_test() { + if (!Halide::exceptions_enabled()) { + return 0; + } + + Var x{"x"}; + RDom r(0, 4); + Func reduced{"reduced"}, consumer{"consumer"}; + reduced(x) = 0; + reduced(x) += r; // update definition -> not pure + consumer(x) = reduced(x); + + bool error = false; + try { + consumer.eager_inline({reduced}); + } catch (const Halide::CompileError &e) { + error = true; + printf("Expected error (cannot inline impure Func):\n%s\n", e.what()); + } + if (!error) { + printf("eager_inline should have rejected a Func with an update definition!\n"); + return 1; + } + return 0; +} +#endif + +} // namespace + +int main(int argc, char **argv) { + printf("Running eager_inline_exposes_factor_test\n"); + if (eager_inline_exposes_factor_test()) { + return 1; + } + printf("Running eager_inline_chain_test\n"); + if (eager_inline_chain_test()) { + return 1; + } +#if HALIDE_WITH_EXCEPTIONS + printf("Running eager_inline_non_pure_rejected_test\n"); + if (eager_inline_non_pure_rejected_test()) { + return 1; + } +#endif + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/rfactor.cpp b/test/correctness/rfactor.cpp index b3d598117168..d7018403a129 100644 --- a/test/correctness/rfactor.cpp +++ b/test/correctness/rfactor.cpp @@ -1298,6 +1298,408 @@ 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"); + + 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()); + + // 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; +} + +// 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; +} + +#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; + printf("Expected error (unsigned min not hoistable):\n%s\n", e.what()); + } + 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; + printf("Expected error (nothing to hoist):\n%s\n", e.what()); + } + 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 +1749,17 @@ 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 (after rfactor)", hoist_invariants_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..ebe1112a5410 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 diff --git a/test/performance/tiled_matmul_arm_neon.cpp b/test/performance/tiled_matmul_arm_neon.cpp new file mode 100644 index 000000000000..0a92f4e0ea02 --- /dev/null +++ b/test/performance/tiled_matmul_arm_neon.cpp @@ -0,0 +1,200 @@ +#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 +// three directives: 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) hoist_invariants() lifts +// the invariant scale product WtScale(i) * VecScale out of the reduction, +// (b) rfactor() splits the scale-free reduction by block, and (c) +// 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"}, 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) { + // Compose the three directives. First, hoist_invariants() lifts + // WtScale(i) * VecScale -- invariant across all of r -- 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().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; +}