Add an env var to enable the scheduler

ROCm · Mar 26, 2019 · 458ec14 · 458ec14
1 parent 8b6a5cd
commit 458ec14
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/src/include/migraphx/schedule.hpp b/src/include/migraphx/schedule.hpp
@@ -17,6 +17,7 @@ struct program;
 struct schedule
 {
     schedule_model model{};
+    bool enable = true;
     std::string name() const { return "schedule"; }
     void apply(program& p) const;
 };

diff --git a/src/schedule.cpp b/src/schedule.cpp
@@ -341,6 +341,8 @@ struct stream_info
 
 void schedule::apply(program& p) const
 {
+    if (not enable)
+        return;
     stream_info si;
     auto last = std::prev(p.end());
     si.accumulate_weights(last, model);

diff --git a/src/targets/gpu/target.cpp b/src/targets/gpu/target.cpp
@@ -25,6 +25,8 @@ namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_SCHEDULE_PASS)
+
 std::vector<pass> target::get_passes(migraphx::context& gctx) const
 {
     auto& ctx = any_cast<context>(gctx);
@@ -53,7 +55,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx) const
         fuse_ops{&ctx},
         dead_code_elimination{},
         write_literals{&ctx},
-        schedule{gpu::schedule_model{ctx.get_current_device().nstreams()}},
+        schedule{gpu::schedule_model{ctx.get_current_device().nstreams()}, enabled(MIGRAPHX_ENABLE_SCHEDULE_PASS{})},
         memory_coloring{"hip::allocate"},
         dead_code_elimination{},
         eliminate_workspace{},