Merge pull request halide#3280 from halide/perform_inline_with_order

Use realization order when inlining to avoid extra works

Author: psuriana

src/AutoSchedule.cpp

@@ -2999,10 +2999,10 @@ Partitioner::analyze_spatial_locality(const FStage &stg,
     Definition def = get_stage_definition(stg.func, stg.stage_num);
     // Perform inlining on the all the values and the args in the stage.
     for (auto &val : def.values()) {
-        val = perform_inline(val, dep_analysis.env, inlines);
+        val = perform_inline(val, dep_analysis.env, inlines, dep_analysis.order);
     }
     for (auto &arg : def.args()) {
-        arg = perform_inline(arg, dep_analysis.env, inlines);
+        arg = perform_inline(arg, dep_analysis.env, inlines, dep_analysis.order);
     }
     def.accept(&find);
 
@@ -3433,7 +3433,7 @@ string generate_schedules(const vector<Function> &outputs, const Target &target,
     // Initialize the cost model.
     // Compute the expression costs for each function in the pipeline.
     debug(2) << "Initializing region costs...\n";
-    RegionCosts costs(env);
+    RegionCosts costs(env, order);
     if (debug::debug_level() >= 3) {
         costs.disp_func_costs();
     }
@@ -3468,7 +3468,7 @@ string generate_schedules(const vector<Function> &outputs, const Target &target,
         debug(2) << "Re-computing function value bounds...\n";
         func_val_bounds = compute_function_value_bounds(order, env);
         debug(2) << "Re-initializing region costs...\n";
-        RegionCosts costs(env);
+        RegionCosts costs(env, order);
         debug(2) << "Re-initializing dependence analysis...\n";
         dep_analysis = DependenceAnalysis(env, order, func_val_bounds);
         debug(2) << "Re-computing pipeline bounds...\n";

src/AutoScheduleUtils.cpp

@@ -155,7 +155,8 @@ vector<DimBounds> get_stage_bounds(Function f, const DimBounds &pure_bounds) {
 }
 
 Expr perform_inline(Expr e, const map<string, Function> &env,
-                    const set<string> &inlines) {
+                    const set<string> &inlines,
+                    const vector<string> &order) {
     if (inlines.empty()) {
         return e;
     }
@@ -168,8 +169,24 @@ Expr perform_inline(Expr e, const map<string, Function> &env,
         // Find all the function calls in the current expression.
         FindAllCalls find;
         inlined_expr.accept(&find);
-        const set<string> &calls = find.funcs_called;
+        const set<string> &calls_unsorted = find.funcs_called;
+
+        vector<string> calls(calls_unsorted.begin(), calls_unsorted.end());
+        // Sort 'calls' based on the realization order in descending order
+        // if provided (i.e. last to be realized comes first).
+        if (!order.empty()) {
+            std::sort(calls.begin(), calls.end(),
+                [&order](const string &lhs, const string &rhs){
+                    const auto &iter_lhs = std::find(order.begin(), order.end(), lhs);
+                    const auto &iter_rhs = std::find(order.begin(), order.end(), rhs);
+                    return iter_lhs > iter_rhs;
+                }
+            );
+        }
+
         // Check if any of the calls are in the set of functions to be inlined.
+        // Inline from the last function to be realized to avoid extra
+        // inlining works.
         for (const auto &call : calls) {
             if (inlines.find(call) != inlines.end()) {
                 Function prod_func = env.at(call);

src/AutoScheduleUtils.h

@@ -84,9 +84,12 @@ DimBounds get_stage_bounds(Function f, int stage_num, const DimBounds &pure_boun
 std::vector<DimBounds> get_stage_bounds(Function f, const DimBounds &pure_bounds);
 
 /** Recursively inline all the functions in the set 'inlines' into the
- * expression 'e' and return the resulting expression. */
+ * expression 'e' and return the resulting expression. If 'order' is
+ * passed, inlining will be done in the reverse order of function realization
+ * to avoid extra inlining works. */
 Expr perform_inline(Expr e, const std::map<std::string, Function> &env,
-                    const std::set<std::string> &inlines = std::set<std::string>());
+                    const std::set<std::string> &inlines = std::set<std::string>(),
+                    const std::vector<std::string> &order = std::vector<std::string>());
 
 /** Return all functions that are directly called by a function stage (f, stage). */
 std::set<std::string> get_parents(Function f, int stage);

src/RegionCosts.cpp

@@ -324,7 +324,9 @@ map<string, Expr> compute_expr_detailed_byte_loads(Expr expr) {
 
 } // anonymous namespace
 
-RegionCosts::RegionCosts(const map<string, Function> &_env) : env(_env) {
+RegionCosts::RegionCosts(const map<string, Function> &_env,
+                         const vector<string> &_order)
+        : env(_env), order(_order) {
     for (const auto &kv : env) {
         // Pre-compute the function costs without any inlining.
         func_cost[kv.first] = get_func_cost(kv.second);
@@ -442,7 +444,7 @@ RegionCosts::stage_detailed_load_costs(string func, int stage,
     } else {
         Definition def = get_stage_definition(curr_f, stage);
         for (const auto &e : def.values()) {
-            Expr inlined_expr = perform_inline(e, env, inlines);
+            Expr inlined_expr = perform_inline(e, env, inlines, order);
             inlined_expr = simplify(inlined_expr);
 
             map<string, Expr> expr_load_costs = compute_expr_detailed_byte_loads(inlined_expr);
@@ -553,7 +555,8 @@ RegionCosts::detailed_load_costs(const map<string, Box> &regions,
     return load_costs;
 }
 
-Cost RegionCosts::get_func_stage_cost(const Function &f, int stage, const set<string> &inlines) {
+Cost RegionCosts::get_func_stage_cost(const Function &f, int stage,
+                                      const set<string> &inlines) {
     if (f.has_extern_definition()) {
         return Cost();
     }
@@ -563,7 +566,7 @@ Cost RegionCosts::get_func_stage_cost(const Function &f, int stage, const set<st
     Cost cost(0, 0);
 
     for (const auto &e : def.values()) {
-        Expr inlined_expr = perform_inline(e, env, inlines);
+        Expr inlined_expr = perform_inline(e, env, inlines, order);
         inlined_expr = simplify(inlined_expr);
 
         Cost expr_cost = compute_expr_cost(inlined_expr);
@@ -578,7 +581,7 @@ Cost RegionCosts::get_func_stage_cost(const Function &f, int stage, const set<st
 
     if (!f.is_pure()) {
         for (const auto &arg : def.args()) {
-            Expr inlined_arg = perform_inline(arg, env, inlines);
+            Expr inlined_arg = perform_inline(arg, env, inlines, order);
             inlined_arg = simplify(inlined_arg);
 
             Cost expr_cost = compute_expr_cost(inlined_arg);
@@ -639,7 +642,7 @@ Expr RegionCosts::region_footprint(const map<string, Box> &regions,
         }
     }
 
-    vector<string> order = topological_order(outs, env);
+    vector<string> top_order = topological_order(outs, env);
 
     Expr working_set_size = make_zero(Int(64));
     Expr curr_size = make_zero(Int(64));
@@ -657,7 +660,7 @@ Expr RegionCosts::region_footprint(const map<string, Box> &regions,
         }
     }
 
-    for (const auto &f : order) {
+    for (const auto &f : top_order) {
         if (regions.find(f) != regions.end()) {
             curr_size += get_element(func_sizes, f);
         }

src/RegionCosts.h

@@ -41,6 +41,9 @@ struct Cost {
 struct RegionCosts {
     /** An environment map which contains all functions in the pipeline. */
     std::map<std::string, Function> env;
+    /** Realization order of functions in the pipeline. The first function to
+     * be realized comes first. */
+    std::vector<std::string> order;
     /** A map containing the cost of computing a value in each stage of a
      * function. The number of entries in the vector is equal to the number of
      * stages in the function. */
@@ -130,8 +133,10 @@ struct RegionCosts {
     void disp_func_costs();
 
     /** Construct a region cost object for the pipeline. 'env' is a map of all
-     * functions in the pipeline.*/
-    RegionCosts(const std::map<std::string, Function> &env);
+     * functions in the pipeline. 'order' is the realization order of functions
+     * in the pipeline. The first function to be realized comes first. */
+    RegionCosts(const std::map<std::string, Function> &env,
+                const std::vector<std::string> &order);
 };
 
 /** Return true if the cost of inlining a function is equivalent to the