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,6 +25,7 @@
 #include "IROperator.h"
 #include "IRPrinter.h"
 #include "ImageParam.h"
+#include "Inline.h"
 #include "LLVM_Output.h"
 #include "Lower.h"
 #include "Param.h"
@@ -3232,6 +3233,31 @@ Func &Func::compute_inline() {
     return compute_at(LoopLevel::inlined());
 }
 
+Stage &Stage::eager_inline(const std::vector<Func> &fs) {
+    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 stage's definition 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(definition, f.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, 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 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, 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). 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,89 @@
+#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
+
+    // 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);
+
+    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;
+        }
+    }
+
+    // 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;
+}