Add eager_inline() directive for schedule-time inlining

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> &)) & Func::eager_inline, py::arg("fs"))
+            .def("eager_inline", [](Func &func, const py::args &args) -> Func & {
+                return func.eager_inline(args_to_vector<Func>(args));
+            })
             .def("compute_root", &Func::compute_root)
             .def("store_root", &Func::store_root)
 

python_bindings/src/halide/halide_/PyStage.cpp

@@ -24,6 +24,11 @@ void define_stage(py::module &m) {
             .def("rfactor", static_cast<Func (Stage::*)(const RVar &, const Var &)>(&Stage::rfactor),
                  py::arg("r"), py::arg("v"))
 
+            .def("eager_inline", (Stage & (Stage::*)(const std::vector<Func> &)) & Stage::eager_inline, py::arg("fs"))
+            .def("eager_inline", [](Stage &stage, const py::args &args) -> Stage & {
+                return stage.eager_inline(args_to_vector<Func>(args));
+            })
+
             .def("split_vars", [](const Stage &stage) -> py::list {
                 auto vars = stage.split_vars();
                 py::list result;

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<Func> &fs) {
+    vector<Function> funcs;
+    map<string, Func> 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<string, Function> env = Internal::build_environment(funcs);
+    vector<string> 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<Func> &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();

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<Func> &fs);
+
+    template<typename... Args>
+    HALIDE_NO_USER_CODE_INLINE std::enable_if_t<Internal::all_are_convertible<Func, Args...>::value, Stage &>
+    eager_inline(const Func &first, Args &&...args) {
+        std::vector<Func> collected_args{first, std::forward<Args>(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<Func> &fs);
+
+    template<typename... Args>
+    HALIDE_NO_USER_CODE_INLINE std::enable_if_t<Internal::all_are_convertible<Func, Args...>::value, Func &>
+    eager_inline(const Func &first, Args &&...args) {
+        std::vector<Func> collected_args{first, std::forward<Args>(args)...};
+        return eager_inline(collected_args);
+    }
+    // @}
+
     /** Get a handle on an update step for the purposes of scheduling
      * it. */
     Stage update(int idx = 0);

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

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

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

test/correctness/eager_inline.cpp

@@ -0,0 +1,146 @@
+#include "Halide.h"
+#include <sstream>
+#include <string>
+
+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<int> 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<int> 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<int> 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<int> 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;
+}

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

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;
+}

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;
+}