From 9a46ad2e7b8224099af7e1a96952acf698f7d6d2 Mon Sep 17 00:00:00 2001
From: Andrey Malyshev <elvin.nnov@gmail.com>
Date: Tue, 19 Jul 2022 19:52:47 +0300
Subject: [PATCH] Remove extra call of PlanDevice in OptimizeImpl

---
 src/relay/backend/build_module.cc                | 15 ++-------------
 src/relay/transforms/annotate_texture_storage.cc |  6 +++++-
 2 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index a372229c7e55..917ca8a4afa0 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -345,6 +345,7 @@ class RelayBuildModule : public runtime::ModuleNode {
 
     // Fuse the operations if it is needed.
     pass_seqs.push_back(transform::FuseOps());
+    pass_seqs.push_back(transform::PlanDevices(config_));
 
     // Create a sequential pass and perform optimizations.
     transform::Pass seq = transform::Sequential(pass_seqs);
@@ -396,20 +397,8 @@ class RelayBuildModule : public runtime::ModuleNode {
     relay_module = transform::Inline()(relay_module);
     relay_module = transform::InferType()(relay_module);
     relay_module = transform::LabelOps()(relay_module);
-
     relay_module = transform::AnnotateMemoryScope(config_)(relay_module);
-    pass_seqs = GetPassPrefix(
-        /*is_homogenous=*/config_->optional_homogeneous_target.defined(), /*is_vm=*/false);
-    pass_seqs.push_back(transform::PlanDevices(config_));
-    // Create a sequential pass and perform optimizations.
-    seq = transform::Sequential(pass_seqs);
-    if (config_->optional_homogeneous_target.defined()) {
-      With<Target> tctx(config_->optional_homogeneous_target);
-      relay_module = seq(relay_module);
-    } else {
-      relay_module = seq(relay_module);
-    }
-    relay_module = transform::InferType()(relay_module);
+
     ICHECK(relay_module.defined());
 
     return relay_module;
diff --git a/src/relay/transforms/annotate_texture_storage.cc b/src/relay/transforms/annotate_texture_storage.cc
index a7eca1ebe70c..66b4523f84e2 100644
--- a/src/relay/transforms/annotate_texture_storage.cc
+++ b/src/relay/transforms/annotate_texture_storage.cc
@@ -504,7 +504,11 @@ Map<Expr, Array<String>> CollectStorageInfo(const Expr& expr) {
 
 Expr AnnotateMemoryScopeExpr(const Expr& expr, const IRModule& mod, CompilationConfig config) {
   auto storage_scope = CollectStorageInfo(expr);
-  return VDRewriter(storage_scope).Rewrite(expr);
+  if (storage_scope.size()) {
+    return VDRewriter(storage_scope).Rewrite(expr);
+  } else {
+    return expr;
+  }
 }
 
 namespace transform {