Fix print_loop_nest

liqunfu · Jul 13, 2017 · 2aa759d · 2aa759d
1 parent 1bd0550
commit 2aa759d
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/PrintLoopNest.cpp b/src/PrintLoopNest.cpp
@@ -157,8 +157,7 @@ string print_loop_nest(const vector<Function> &output_funcs) {
     // Compute an environment
     map<string, Function> env;
     for (Function f : output_funcs) {
-        map<string, Function> more_funcs = find_transitive_calls(f);
-        env.insert(more_funcs.begin(), more_funcs.end());
+        populate_environment(f, env);
     }
 
     // Create a deep-copy of the entire graph of Funcs.
@@ -170,6 +169,11 @@ string print_loop_nest(const vector<Function> &output_funcs) {
         Func(f).compute_root().store_root();
     }
 
+    // Ensure that all ScheduleParams become well-defined constant Exprs.
+    for (auto &f : env) {
+        f.second.substitute_schedule_param_exprs();
+    }
+
     // Substitute in wrapper Funcs
     env = wrap_func_calls(env);