From: Bing Xu <antinucleon@gmail.com>
Date: Thu, 9 May 2019 04:06:33 +0000 (-0700)
Subject: [BuildModule] Fix AlterLayout Pass (#3155)
X-Git-Tag: upstream/0.7.0~2444
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=181dbd8e94222b1dad5da4c3f15b8c63facc3582;p=platform%2Fupstream%2Ftvm.git

[BuildModule] Fix AlterLayout Pass (#3155)
---

diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index b60a048..67ab750 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -504,7 +504,14 @@ class RelayBuildModule : public runtime::ModuleNode {
     if (cfg.pass_enabled("AlterOpLayout")) {
       if (targets.size() == 1) {
         func = CallPackedFunc("relay._ir_pass.infer_type", func, nullptr);
-        func = CallPackedFunc("relay._ir_pass.AlterOpLayout", func);
+        auto enter_pf = GetPackedFunc("_EnterTargetScope");
+        auto exit_pf = GetPackedFunc("_ExitTargetScope");
+        for (const auto& kv : targets) {
+          auto target = Target::create(kv.second);
+          (*enter_pf)(target);
+          func = CallPackedFunc("relay._ir_pass.AlterOpLayout", func);
+          (*exit_pf)();
+        }
       } else {
         LOG(WARNING) << "AlterOpLayout pass is not enabled for heterogeneous"
                   << " execution yet.";