From 126c1fd3df628a99dec6c0c29b53345f8c422a63 Mon Sep 17 00:00:00 2001
From: "Tang, Cheng" <souptc@gmail.com>
Date: Sat, 29 Dec 2018 12:10:08 -0800
Subject: [PATCH] do replacement based on node arg usage, instead of provider
 type (#263)

---
 .../core/framework/transformer_memcpy.cc      | 50 ++++++++++++++++---
 .../core/framework/transformer_memcpy.h       |  5 +-
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/onnxruntime/core/framework/transformer_memcpy.cc b/onnxruntime/core/framework/transformer_memcpy.cc
index b04f278599..8c7fe38636 100644
--- a/onnxruntime/core/framework/transformer_memcpy.cc
+++ b/onnxruntime/core/framework/transformer_memcpy.cc
@@ -43,6 +43,12 @@ bool TransformerMemcpyImpl::ModifyGraph(const KernelRegistryManager& kernel_regi
   // for initializers shared by different providers, create dups
   ProcessInitializers();
 
+  for (auto arg : non_provider_input_defs_)
+    BuildDefsMapping(arg, kernel_registries);
+
+  for (auto arg : non_provider_output_defs_)
+    BuildDefsMapping(arg, kernel_registries);
+
   for (auto arg : non_provider_output_defs_)
     if (provider_input_defs_.count(arg)) {
       AddCopyNode(arg, true);
@@ -55,10 +61,6 @@ bool TransformerMemcpyImpl::ModifyGraph(const KernelRegistryManager& kernel_regi
       modified = true;
     }
 
-  for (auto p_node : provider_nodes_) {
-    p_node->ReplaceDefs(replacements_);
-  }
-
   return modified;
 }
 
@@ -113,6 +115,32 @@ void TransformerMemcpyImpl::ProcessDefs(onnxruntime::Node& node, const KernelReg
   }
 }
 
+//for non_provider defs, collect the nodes that expect it is provider tensor as input/output.
+void TransformerMemcpyImpl::BuildDefsMapping(const onnxruntime::NodeArg* arg, const KernelRegistryManager& kernel_registries) {
+  for (auto it = graph_.Nodes().begin(); it != graph_.Nodes().end(); it++) {
+    if (it->OpType() == "MemcpyFromHost" || it->OpType() == "MemcpyToHost")
+      continue;
+    auto input_it = std::find(it->MutableInputDefs().begin(), it->MutableInputDefs().end(), const_cast<onnxruntime::NodeArg*>(arg));
+    auto output_it = std::find(it->MutableOutputDefs().begin(), it->MutableOutputDefs().end(), const_cast<onnxruntime::NodeArg*>(arg));
+    int arg_input_index = input_it != it->MutableInputDefs().end() ? static_cast<int>(input_it - it->MutableInputDefs().begin()) : -1;
+    int arg_output_index = output_it != it->MutableOutputDefs().end() ? static_cast<int>(output_it - it->MutableOutputDefs().begin()) : -1;
+    if (arg_input_index == -1 && arg_output_index == -1)
+      continue;
+    if (it->GetExecutionProviderType() == provider_) {
+      const KernelCreateInfo* kci = nullptr;
+      kernel_registries.SearchKernelRegistry(*it, &kci);
+      if (arg_input_index != -1) {
+        if (!kci || !MemTypeOnCpuExplicitly(kci->kernel_def->InputMemoryType(arg_input_index)))
+          provider_input_nodes_[arg].insert(&*it);
+      }
+      if (arg_output_index != -1) {
+        if (!kci || !MemTypeOnCpuExplicitly(kci->kernel_def->OutputMemoryType(arg_output_index)))
+          provider_output_nodes_[arg].insert(&*it);
+      }
+    }
+  }
+}
+
 void TransformerMemcpyImpl::AddCopyNode(onnxruntime::NodeArg* arg, bool is_input) {
   // create unique name for new def
   std::string new_def_name = graph_.GenerateNodeArgName(arg->Name() + "_" + provider_);
@@ -129,9 +157,17 @@ void TransformerMemcpyImpl::AddCopyNode(onnxruntime::NodeArg* arg, bool is_input
                                   std::vector<onnxruntime::NodeArg*>{src_arg},
                                   std::vector<onnxruntime::NodeArg*>{dst_arg});
   new_node.SetExecutionProviderType(provider_);
-
-  // only add copy-node here; renaming references happens later
-  replacements_.insert(std::make_pair(arg, new_arg));
+  std::map<const onnxruntime::NodeArg*, onnxruntime::NodeArg*> map = {{arg, new_arg}};
+  auto it = provider_input_nodes_.find(arg);
+  if (it != provider_input_nodes_.end()) {
+    for (auto* node : it->second)
+      node->ReplaceDefs(map);
+  }
+  it = provider_output_nodes_.find(arg);
+  if (it != provider_output_nodes_.end()) {
+    for (auto* node : it->second)
+      node->ReplaceDefs(map);
+  }
 }
 
 template <typename NodeArgSetType>
diff --git a/onnxruntime/core/framework/transformer_memcpy.h b/onnxruntime/core/framework/transformer_memcpy.h
index c4459d9d74..d626cda202 100644
--- a/onnxruntime/core/framework/transformer_memcpy.h
+++ b/onnxruntime/core/framework/transformer_memcpy.h
@@ -20,6 +20,7 @@ class TransformerMemcpyImpl {
 
  private:
   void ProcessDefs(onnxruntime::Node& node, const KernelRegistryManager& kernel_registries);
+  void BuildDefsMapping(const onnxruntime::NodeArg* arg, const KernelRegistryManager& kernel_registries);
   void AddCopyNode(onnxruntime::NodeArg* arg, bool is_input);
   void ProcessInitializers();
 
@@ -45,7 +46,9 @@ class TransformerMemcpyImpl {
   std::set<onnxruntime::NodeArg*, NodeArgCompare> non_provider_output_defs_;       // all output defs of non-provider nodes
   std::set<const onnxruntime::NodeArg*, NodeArgCompare> provider_input_defs_;      // all input defs of provider nodes that should be in provider allocator
   std::set<onnxruntime::NodeArg*, NodeArgCompare> provider_output_defs_;           // all output defs of provider nodes that should be in provider allocator
-  std::map<const onnxruntime::NodeArg*, onnxruntime::NodeArg*> replacements_;
+  std::map<const onnxruntime::NodeArg*, std::set<onnxruntime::Node*, NodeCompare>> provider_input_nodes_;
+  std::map<const onnxruntime::NodeArg*, std::set<onnxruntime::Node*, NodeCompare>> provider_output_nodes_;
+
   onnxruntime::Graph& graph_;
   std::string provider_;
 };