saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-28 20:11:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

do replacement based on node arg usage, instead of provider type (#263)

Browse source
This commit is contained in:
Tang, Cheng 2018-12-29 12:10:08 -08:00 committed by Changming Sun
parent 6a090985fb
commit 126c1fd3df
2 changed files with 47 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
onnxruntime/core/framework/transformer_memcpy.cc
View file

@ -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>

5
onnxruntime/core/framework/transformer_memcpy.h
View file

@ -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_;
};