diff --git a/python_bindings/src/halide/halide_/PyFunc.cpp b/python_bindings/src/halide/halide_/PyFunc.cpp index 9daa3fae3e6f..cccd87bc1789 100644 --- a/python_bindings/src/halide/halide_/PyFunc.cpp +++ b/python_bindings/src/halide/halide_/PyFunc.cpp @@ -219,6 +219,10 @@ 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 & (Func::*)(const std::vector &)) & Func::eager_inline, py::arg("fs")) + .def("eager_inline", [](Func &func, const py::args &args) -> Func & { + return func.eager_inline(args_to_vector(args)); + }) .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..934c58bd9439 100644 --- a/python_bindings/src/halide/halide_/PyStage.cpp +++ b/python_bindings/src/halide/halide_/PyStage.cpp @@ -24,6 +24,11 @@ void define_stage(py::module &m) { .def("rfactor", static_cast(&Stage::rfactor), py::arg("r"), py::arg("v")) + .def("eager_inline", (Stage & (Stage::*)(const std::vector &)) & Stage::eager_inline, py::arg("fs")) + .def("eager_inline", [](Stage &stage, const py::args &args) -> Stage & { + return stage.eager_inline(args_to_vector(args)); + }) + .def("split_vars", [](const Stage &stage) -> py::list { auto vars = stage.split_vars(); py::list result; diff --git a/src/Func.cpp b/src/Func.cpp index db18fd7a590b..90beb3326a56 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -25,10 +25,12 @@ #include "IROperator.h" #include "IRPrinter.h" #include "ImageParam.h" +#include "Inline.h" #include "LLVM_Output.h" #include "Lower.h" #include "Param.h" #include "PrintLoopNest.h" +#include "RealizationOrder.h" #include "Simplify.h" #include "Solve.h" #include "Substitute.h" @@ -3232,6 +3234,46 @@ Func &Func::compute_inline() { return compute_at(LoopLevel::inlined()); } +Stage &Stage::eager_inline(const std::vector &fs) { + vector funcs; + map by_name; + 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"; + funcs.push_back(f.function()); + by_name.emplace(f.name(), f); + } + + // Inlining rewrites this stage's definition in place, replacing every direct + // call to f with f's body. A call to g inside f's body only becomes visible + // to us after f is inlined, so a caller must be inlined before its callees. + // topological_order() gives realization order (callees before callers); + // reverse it to inline callers first, so passing fs in any order works. + map env = Internal::build_environment(funcs); + vector order = Internal::topological_order(funcs, env); + std::reverse(order.begin(), order.end()); + + for (const string &name : order) { + auto it = by_name.find(name); + if (it != by_name.end()) { + Internal::inline_function(definition, it->second.function()); + } + } + return *this; +} + +Func &Func::eager_inline(const std::vector &fs) { + invalidate_cache(); + // Target the initial (pure) definition, mirroring other Func-level scheduling + // shorthands; use f.update(n).eager_inline(...) to inline into an update. + Stage(func, func.definition(), 0).eager_inline(fs); + return *this; +} + Func &Func::trace_loads() { invalidate_cache(); func.trace_loads(); diff --git a/src/Func.h b/src/Func.h index 396d2d436a0c..57265e950001 100644 --- a/src/Func.h +++ b/src/Func.h @@ -191,6 +191,31 @@ class Stage { Func rfactor(const RVar &r, const Var &v); // @} + /** Immediately inline direct calls to each of the given Funcs into this + * stage's definition. The Funcs are inlined in dependency order regardless + * of the order they are passed, so if one inlined Func's body calls another, + * both are fully folded in. + * + * Unlike compute_inline(), which merely marks a Func to be inlined during + * lowering, eager_inline() performs the substitution now, at schedule time, + * rewriting only this stage's definition in place. This is useful to surface + * structure that other schedule-time directives (e.g. rfactor()) need to see. + * + * 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 stage's calls to them are replaced. */ + // @{ + Stage &eager_inline(const std::vector &fs); + + template + HALIDE_NO_USER_CODE_INLINE std::enable_if_t::value, Stage &> + eager_inline(const Func &first, Args &&...args) { + std::vector collected_args{first, std::forward(args)...}; + return eager_inline(collected_args); + } + // @} + /** 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 @@ -2631,6 +2656,33 @@ class Func { */ Func &compute_inline(); + /** Immediately inline direct calls to each of the given Funcs into this + * Func's initial (pure) definition. The Funcs are inlined in dependency + * order regardless of the order they are passed, so if one inlined Func's + * body calls another, both are fully folded in. This is shorthand for + * update(0)-style scheduling: to inline into an update definition, call + * eager_inline() on that stage, e.g. f.update(n).eager_inline(...). + * + * Unlike compute_inline(), which merely marks a Func to be inlined during + * lowering, eager_inline() performs the substitution now, at schedule time, + * rewriting the definition in place. This is useful to surface structure that + * other schedule-time directives need to see. + * + * 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 definition's calls to them are replaced. */ + // @{ + Func &eager_inline(const std::vector &fs); + + template + HALIDE_NO_USER_CODE_INLINE std::enable_if_t::value, Func &> + eager_inline(const Func &first, Args &&...args) { + std::vector collected_args{first, std::forward(args)...}; + return eager_inline(collected_args); + } + // @} + /** Get a handle on an update step for the purposes of scheduling * it. */ Stage update(int idx = 0); diff --git a/src/Inline.cpp b/src/Inline.cpp index 2d1976c9e461..f5b5a260ec8a 100644 --- a/src/Inline.cpp +++ b/src/Inline.cpp @@ -1,6 +1,7 @@ #include "Inline.h" #include "CSE.h" #include "Debug.h" +#include "Definition.h" #include "ExternFuncArgument.h" #include "IRMutator.h" #include "IROperator.h" @@ -149,7 +150,6 @@ void validate_schedule_inlined_function(Function f) { Inliner::Inliner(const Function &f) { internal_assert(f.can_be_inlined()) << "Illegal to inline " << f.name() << "\n"; - validate_schedule_inlined_function(f); add(f); } @@ -356,5 +356,10 @@ void inline_function(Function caller, const Function &f) { } } +void inline_function(Definition &def, const Function &f) { + Inliner i(f); + def.mutate(&i); +} + } // namespace Internal } // namespace Halide diff --git a/src/Inline.h b/src/Inline.h index 40be123150ae..a2b7457b4acd 100644 --- a/src/Inline.h +++ b/src/Inline.h @@ -81,11 +81,13 @@ class Inliner : public IRMutator { /** Inline a single named function, which must be pure. For a pure function to * be inlined, it must not have any specializations (i.e. it can only have one - * values definition). */ + * values definition). The Definition overload rewrites just one stage's + * definition in place; the Function overload rewrites all of a Func's stages. */ // @{ Stmt inline_function(const Stmt &s, const Function &f); Expr inline_function(const Expr &e, const Function &f); void inline_function(Function caller, const Function &f); +void inline_function(Definition &def, const Function &f); // @} /** Inline a set of pure functions. Equivalent in effect to calling diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 5df549c46c62..862dca2d228a 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -103,6 +103,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/eager_inline.cpp b/test/correctness/eager_inline.cpp new file mode 100644 index 000000000000..4bfd2bce4e90 --- /dev/null +++ b/test/correctness/eager_inline.cpp @@ -0,0 +1,146 @@ +#include "Halide.h" +#include +#include + +using namespace Halide; + +// 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. + +namespace { + +// Does the printed form of `e` contain a direct call to Func `name`? +bool mentions(const Expr &e, const std::string &name) { + std::ostringstream os; + os << e; + return os.str().find(name + "(") != std::string::npos; +} + +} // namespace + +int main(int argc, char **argv) { + 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 + + // Pass the Funcs in "wrong" (callee-before-caller) order: a before b, even + // though b's body calls a. eager_inline() topologically sorts them, so both + // are fully folded regardless of the argument order. + c.eager_inline(a, b); + + Expr c_body = c.function().definition().values()[0]; + Expr c_expected = x * 3 + 3; + internal_assert(Internal::can_prove(c_body == c_expected)) + << "eager_inline chain failed to fold all calls to a and b into c\n" + << "Saw: " << c_body << "\nExpected: " << c_expected << "\n"; + + 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; + } + } + + // A longer chain, also passed in a scrambled order, to exercise the + // topological sort more thoroughly: each Func's body calls the previous one. + { + Func d0{"d0"}, d1{"d1"}, d2{"d2"}, d3{"d3"}, sink{"sink"}; + d0(x) = x + 1; + d1(x) = d0(x) * 2; // calls d0 + d2(x) = d1(x) + 3; // calls d1 + d3(x) = d2(x) * 5; // calls d2 + sink(x) = d3(x) - 4; // calls d3 + + // Scrambled order (not caller-first): the sort must still order them so + // every call gets folded. + sink.eager_inline(d2, d0, d3, d1); + + Expr sink_body = sink.function().definition().values()[0]; + internal_assert(!mentions(sink_body, "d0") && !mentions(sink_body, "d1") && + !mentions(sink_body, "d2") && !mentions(sink_body, "d3")) + << "eager_inline left residual calls after a scrambled-order chain\n" + << "Saw: " << sink_body << "\n"; + + Buffer sout = sink.realize({8}); + for (int i = 0; i < 8; i++) { + int ref = (((i + 1) * 2 + 3) * 5) - 4; + if (sout(i) != ref) { + printf("eager_inline scrambled-chain mismatch at %d: %d vs %d\n", i, sout(i), ref); + return 1; + } + } + } + + // Passing a Func that this stage does not call is silently ignored: there + // are no direct calls to fold, so the definition is left unchanged. + { + Func p{"p"}, unrelated{"unrelated"}, q{"q"}; + p(x) = x + 1; + unrelated(x) = x * 100; // never referenced by q + q(x) = p(x) + 2; // calls p, but not unrelated + + // Mix a reachable Func (p) with an unreachable one (unrelated): p is + // inlined, unrelated is a no-op rather than an error. + q.eager_inline(unrelated, p); + + Expr q_body = q.function().definition().values()[0]; + internal_assert(!mentions(q_body, "p")) + << "eager_inline should have inlined the reachable Func p\n" + << "Saw: " << q_body << "\n"; + + Buffer qout = q.realize({8}); + for (int i = 0; i < 8; i++) { + if (qout(i) != (i + 1) + 2) { + printf("eager_inline unreachable-arg mismatch at %d: %d vs %d\n", i, qout(i), (i + 1) + 2); + return 1; + } + } + } + + // eager_inline() is stage-scoped: inlining into one stage leaves the other + // definitions of the same Func untouched. + { + RDom r(0, 4); + + // Stage-level: inline into the update only; the init definition still + // calls prod. + Func prod{"prod"}, f{"f"}; + prod(x) = x + 1; + f(x) = prod(x); // init definition calls prod + f(x) += prod(x) * r; // update definition also calls prod + + f.update(0).eager_inline(prod); + + internal_assert(mentions(f.function().definition().values()[0], "prod")) + << "Stage::eager_inline on update(0) should not touch the init definition\n"; + internal_assert(!mentions(f.function().update(0).values()[0], "prod")) + << "Stage::eager_inline on update(0) should have inlined prod into the update\n"; + + // Semantics preserved: f(x) = (x+1) + sum_{r=0..3} (x+1)*r = 7*(x+1). + Buffer fout = f.realize({8}); + for (int i = 0; i < 8; i++) { + if (fout(i) != 7 * (i + 1)) { + printf("stage eager_inline mismatch at %d: %d vs %d\n", i, fout(i), 7 * (i + 1)); + return 1; + } + } + + // Func-level: targets the init definition only, leaving updates alone. + Func g{"g"}; + g(x) = prod(x); + g(x) += prod(x) * r; + g.eager_inline(prod); + internal_assert(!mentions(g.function().definition().values()[0], "prod")) + << "Func::eager_inline should inline prod into the init definition\n"; + internal_assert(mentions(g.function().update(0).values()[0], "prod")) + << "Func::eager_inline should not touch update definitions\n"; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt index 1ebe757ee2f6..377f672cb0a7 100644 --- a/test/error/CMakeLists.txt +++ b/test/error/CMakeLists.txt @@ -22,6 +22,8 @@ tests( bad_const_cast.cpp bad_device_api.cpp bad_dimensions.cpp + bad_eager_inline.cpp + bad_eager_inline_undefined.cpp bad_extern_split.cpp bad_fold.cpp bad_func_object.cpp diff --git a/test/error/bad_eager_inline.cpp b/test/error/bad_eager_inline.cpp new file mode 100644 index 000000000000..65341b62d263 --- /dev/null +++ b/test/error/bad_eager_inline.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" +using namespace Halide; + +int main(int argc, char **argv) { + 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); + + // A Func with an update definition is not inlinable, so eager_inline() rejects it. + consumer.eager_inline({reduced}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/bad_eager_inline_undefined.cpp b/test/error/bad_eager_inline_undefined.cpp new file mode 100644 index 000000000000..83a556086db4 --- /dev/null +++ b/test/error/bad_eager_inline_undefined.cpp @@ -0,0 +1,15 @@ +#include "Halide.h" +using namespace Halide; + +int main(int argc, char **argv) { + Var x{"x"}; + Func undefined_producer{"undefined_producer"}; // never given a definition + Func consumer{"consumer"}; + consumer(x) = x; + + // An undefined Func has no body to splice in, so eager_inline() rejects it. + consumer.eager_inline({undefined_producer}); + + printf("Success!\n"); + return 0; +}