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

Handle missing initializers in allocation planner to fix crashes with DML provider (#5244)

Browse source 
* Fix memory planning bug with DML EP

* Address PR comments

* Fix typo
This commit is contained in:
Jeff Bloomfield 2020-09-22 19:37:07 -07:00 committed by Tianlei Wu
parent b648fe5f74
commit 389cca7a45
1 changed files with 20 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
onnxruntime/core/framework/allocation_planner.cc
View file

@ -144,7 +144,13 @@ class PlannerImpl {
struct OrtValueInfo {
const onnxruntime::NodeArg* p_def_site; // the (unique) NodeArg corresponding to the MLValue
int usecount = 0; // static reference-count
OrtValueIndex reused_buffer_index; // index of original buffer to reuse
// This is initialized to -1 to ensure that if ProcessDef is somehow not called, planning
// will fail more cleanly. This is also used as a temporary workaround to detect the
// case that the DML provider has removed initilizers from the graph during partitioning.
// Removing initializers is a temporary measure needed to limit the number of copies of
// tensors in GPU memory.
OrtValueIndex reused_buffer_index = -1; // index of original buffer to reuse
};
// ort_value_info_ is indexed by an OrtValueIndex
@ -177,6 +183,12 @@ class PlannerImpl {
}
int& UseCount(const OrtValueName& name) { return UseCount(Index(name)); }
int DecrementUseCount(OrtValueIndex n) {
int& use_count = --UseCount(n);
assert(use_count >= 0);
return use_count;
}
OrtValueIndex& Buffer(OrtValueIndex n) {
ORT_ENFORCE(n >= 0 && static_cast<size_t>(n) < ort_value_info_.size());
return ort_value_info_[n].reused_buffer_index;
@ -643,7 +655,9 @@ class PlannerImpl {
if (node_input->Exists()) {
auto& sym = node_input->Name();
auto original = Buffer(Index(sym));
if (0 == --UseCount(original))
// The index will be -1 if it's an initializer that was removed as part of a temporary workaround.
// See comments in the OrtValueInfo definition.
if ((original != -1) && (0 == DecrementUseCount(original)))
freelist_.push_front(FreeBufferInfo(original, program_counter));
}
}
@ -652,7 +666,9 @@ class PlannerImpl {
if (node_input->Exists()) {
auto& sym = node_input->Name();
auto original = Buffer(Index(sym));
if (0 == --UseCount(original))
// The index will be -1 if it's an initializer that was removed as part of a temporary workaround.
// See comments in the OrtValueInfo definition.
if ((original != -1) && (0 == DecrementUseCount(original)))
freelist_.push_front(FreeBufferInfo(original, program_counter));
}
}
@ -662,7 +678,7 @@ class PlannerImpl {
if (node_output->Exists()) {
auto& sym = node_output->Name();
auto original = Buffer(Index(sym));
if (0 == --UseCount(original))
if (0 == DecrementUseCount(original))
freelist_.push_front(FreeBufferInfo(original, program_counter));
}
}