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

Fix node def replacement by checking memory location (#623)

Browse source 
Fix node def replacement by checking memory location
This commit is contained in:
KeDengMS 2019-03-15 01:12:21 -07:00 committed by GitHub
parent 5c09aa863c
commit a362c3bbdf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 27 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
onnxruntime/core/optimizer/transformer_memcpy.cc
View file

@ -21,7 +21,7 @@ class TransformerMemcpyImpl {
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();
void ProcessInitializers(const KernelRegistryManager& kernel_registries);
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TransformerMemcpyImpl);
@ -114,7 +114,7 @@ bool TransformerMemcpyImpl::ModifyGraph(const KernelRegistryManager& kernel_regi
}
// for initializers shared by different providers, create dups
ProcessInitializers();
ProcessInitializers(kernel_registries);
for (auto arg : graph_.GetInputs())
BuildDefsMapping(arg, kernel_registries);
@ -271,7 +271,7 @@ static const onnxruntime::NodeArg* FindNodeArg(const NodeArgSetType& def_set, co
// We duplicate any initializer that is used by both provider nodes and non-provider nodes
// to ensure that provider nodes and non-provider nodes don't share initializers, as they
// need to stay in different memory locations.
void TransformerMemcpyImpl::ProcessInitializers() {
void TransformerMemcpyImpl::ProcessInitializers(const KernelRegistryManager& kernel_registries) {
std::map<const onnxruntime::NodeArg*, onnxruntime::NodeArg*> replacements;
for (const auto& pair : graph_.GetAllInitializedTensors()) {
const auto& name = pair.first;
@ -294,7 +294,30 @@ void TransformerMemcpyImpl::ProcessInitializers() {
}
for (auto p_node : provider_nodes_) {
p_node->ReplaceDefs(replacements);
// make a copy of replacement map as the node may exclude mapping for InputDefs with MemTypeOnCpuExplicitly
auto dup_replacements = replacements;
const KernelCreateInfo* kci = nullptr;
kernel_registries.SearchKernelRegistry(*p_node, &kci);
p_node->ForEachWithIndex(
p_node->InputDefs(),
[kci, &dup_replacements](const onnxruntime::NodeArg& arg, size_t index) {
if (kci && MemTypeOnCpuExplicitly(kci->kernel_def->InputMemoryType(index)))
dup_replacements.erase(&arg);
return Status::OK();
});
// normally initializers are only inputs, but things may change with ops like assign
p_node->ForEachWithIndex(
p_node->OutputDefs(),
[kci, &dup_replacements](const onnxruntime::NodeArg& arg, size_t index) {
if (kci && MemTypeOnCpuExplicitly(kci->kernel_def->OutputMemoryType(index))) {
ORT_ENFORCE(dup_replacements.find(&arg) == dup_replacements.end());
}
return Status::OK();
});
p_node->ReplaceDefs(dup_replacements);
}
}