diff --git a/src/PrintLoopNest.cpp b/src/PrintLoopNest.cpp index bb4fa211efc5..0299a3cfe4fb 100644 --- a/src/PrintLoopNest.cpp +++ b/src/PrintLoopNest.cpp @@ -157,8 +157,7 @@ string print_loop_nest(const vector &output_funcs) { // Compute an environment map env; for (Function f : output_funcs) { - map 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 &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);