Skip to content

Fix bug in handling of an initializer that provides a graph output. #3912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Scott McKay (skottmckay) merged 3 commits into from
May 12, 2020
Merged

Fix bug in handling of an initializer that provides a graph output. #3912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
35bdae6
Outputs from model execution should always be returned in a newly all…
skottmckay May 12, 2020
0e7c9bf
Add unit test with and without pre-allocated fetch.
skottmckay May 12, 2020
fafb83a
Add some extra info around why we're handling this special case.
skottmckay May 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 52 additions & 19 deletions onnxruntime/core/framework/execution_frame.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,19 +20,14 @@ using namespace onnxruntime::common;

namespace onnxruntime {

IExecutionFrame::IExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches,
const OrtValueNameIdxMap& ort_value_idx_map, const NodeIndexInfo& node_index_info)
IExecutionFrame::IExecutionFrame(const OrtValueNameIdxMap& ort_value_idx_map,
const NodeIndexInfo& node_index_info,
const std::vector<int>& fetch_mlvalue_idxs)
: node_index_info_(node_index_info),
all_values_size_(static_cast<size_t>(ort_value_idx_map.MaxIdx()) + 1),
fetch_mlvalue_idxs_(fetch_mlvalue_idxs) {
ORT_ENFORCE(feeds.size() == feed_mlvalue_idxs.size());
ORT_ENFORCE(fetches.empty() || fetches.size() == fetch_mlvalue_idxs_.size());
ORT_ENFORCE(node_index_info_.GetMaxMLValueIdx() == ort_value_idx_map.MaxIdx(),
"node_index_info and ort_value_idx_map are out of sync and cannot be used");

Init(feed_mlvalue_idxs, feeds, initializers, fetches);
}

IExecutionFrame::~IExecutionFrame() = default;
Expand Down Expand Up @@ -109,6 +104,9 @@ int IExecutionFrame::GetNodeIdxToMLValueIdx(int index) const {
void IExecutionFrame::Init(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::vector<OrtValue>& fetches) {
ORT_ENFORCE(feeds.size() == feed_mlvalue_idxs.size());
ORT_ENFORCE(fetches.empty() || fetches.size() == fetch_mlvalue_idxs_.size());

// 1. resize the all_value_ vector
all_values_.resize(all_values_size_);

Expand All @@ -123,15 +121,44 @@ void IExecutionFrame::Init(const std::vector<int>& feed_mlvalue_idxs, const std:
}

// 3. handle the weights.
// We do this after the fetches to handle an edge case (possibly dubious) where a Constant is an output.
// The Constant gets lifted to an initializer so there's no Node producing the value as an output during Graph
// execution (i.e. Graph execution won't write the value to all_values_).
// We do this after the fetches to handle an edge case where an initializer is an output.
// e.g. A Constant node gets lifted to an initializer so there's no Node producing the value as an output during
// Graph execution (i.e. Graph execution won't write the value to all_values_).
// A non-empty fetches vector will overwrite the actual weight in all_values_[ort_value_idx] if we did this earlier.
// This makes the ONNX Constant test (onnx\backend\test\data\node\test_constant) happy as that
// involves a graph with a single Constant node.
for (const auto& entry : initializers) {
int ort_value_index = entry.first;
all_values_[ort_value_index] = entry.second;

// if the initializer is an output we need to allocate or use a provided fetch buffer and copy the data
// so it can be returned to the caller.
//
// The alternative to handling this as a special case would be to disallow an initializer providing a graph output.
// There's nothing in the ONNX spec that says a graph output must come from a node output though.
// If we took that approach we'd need to:
// - reject a model with an initializer or Constant node (as we convert those to initializers in Graph::Graph)
// that produces a graph output even though it conforms to the ONNX spec
// - update optimizers to not convert something to an initializer that is a graph output
// (e.g. constant folding)
if (IsOutput(ort_value_index)) {
const Tensor& src = entry.second.Get<Tensor>(); // all initializers in ONNX are tensors
Comment thread
skottmckay marked this conversation as resolved.
Resolved
OrtValue& dest = all_values_[ort_value_index];

if (!dest.IsAllocated()) {
// NOTE: This doesn't need to support ExecutionFrame custom allocators as they only come into play
// for a subgraph with an output of unknown shape that needs to be accumulated by the control flow node.
// If the initializer is providing the output, the shape is known.
AllocatorPtr allocator = GetAllocator(src.Location());

auto p_tensor = onnxruntime::make_unique<Tensor>(src.DataType(), src.Shape(), allocator);
auto ml_tensor = DataTypeImpl::GetType<Tensor>();
dest.Init(p_tensor.release(), ml_tensor, ml_tensor->GetDeleteFunc());
}

ORT_THROW_IF_ERROR(CopyTensor(src, *dest.GetMutable<Tensor>()));
Comment thread
pranavsharma marked this conversation as resolved.
Resolved
} else {
all_values_[ort_value_index] = entry.second;
}
}

// 4. handle feed in values. these can override initializer values so must be last
Expand Down Expand Up @@ -171,11 +198,12 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const
const std::vector<int>& fetch_mlvalue_idxs, const std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
const SessionState& session_state)
: IExecutionFrame(feed_mlvalue_idxs, feeds, session_state.GetInitializedTensors(), fetch_mlvalue_idxs, fetches,
session_state.GetOrtValueNameIdxMap(), session_state.GetNodeIndexInfo()),
: IExecutionFrame(session_state.GetOrtValueNameIdxMap(), session_state.GetNodeIndexInfo(), fetch_mlvalue_idxs),
session_state_(session_state),
mem_patterns_(nullptr),
planner_(nullptr) {
Init(feed_mlvalue_idxs, feeds, session_state.GetInitializedTensors(), fetches);

// map the custom allocators to ort_value_idx entries
if (!fetch_allocators.empty()) {
for (size_t idx = 0, end = fetch_mlvalue_idxs.size(); idx < end; ++idx) {
Expand Down Expand Up @@ -232,6 +260,10 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const

ExecutionFrame::~ExecutionFrame() = default;

Status ExecutionFrame::CopyTensor(const Tensor& src, Tensor& dest) const {
return session_state_.GetDataTransferMgr().CopyTensor(src, dest);
}

Status ExecutionFrame::AllocateMLValueTensorSelfOwnBuffer(OrtValue& ort_value, int ort_value_index,
MLDataType element_type, const OrtMemoryInfo& location,
const TensorShape& shape, bool create_fence) {
Expand Down Expand Up @@ -343,7 +375,7 @@ Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(OrtValue& ort_valu

// be generous and use the buffer if it's large enough. log a warning though as it indicates a bad model
if (buffer_num_elements >= required_num_elements) {
// View Operator is reusing the buffer bigger than the required size.
// View Operator is reusing the buffer bigger than the required size.
// Disabling warning message for now. The op is in the process of being deprecated.
#ifndef ENABLE_TRAINING
LOGS(session_state_.Logger(), WARNING) << message;
Expand Down Expand Up @@ -455,13 +487,13 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(OrtValue& ort_value, int ort_
}
case AllocKind::kReuse: {
int reuse_mlvalue_index = per_alloc_plan.reused_buffer;

// In case OrtRunOptions.only_execute_path_to_fetches == true, it is possible that 'reuse_value'
// is not allocated (its upstream op is not executed due to the option).
// In this case we need to allocate 'reuse_value' and then let 'ort_value' to reuse it.
// In this case we need to allocate 'reuse_value' and then let 'ort_value' to reuse it.
OrtValue& reuse_value = GetMutableMLValue(reuse_mlvalue_index);
if (!reuse_value.IsAllocated()) {
ORT_RETURN_IF_ERROR(AllocateAsPerAllocationPlan(reuse_value, reuse_mlvalue_index, shape, nnz));
ORT_RETURN_IF_ERROR(AllocateAsPerAllocationPlan(reuse_value, reuse_mlvalue_index, shape, nnz));
}
ORT_RETURN_IF_ERROR(AllocateMLValueTensorPreAllocateBuffer(
ort_value, reuse_mlvalue_index, ml_data_type, alloc_info, *shape, per_alloc_plan.create_fence_if_async));
Expand Down Expand Up @@ -497,7 +529,8 @@ AllocatorPtr ExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {

// This method is not thread safe!
// Return S_OK and nullptr if index map to an value that is an unused optional input/output
Status ExecutionFrame::CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) {
Status ExecutionFrame::CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx,
const TensorShape* shape, size_t nnz) {
return AllocateAsPerAllocationPlan(ort_value, ort_value_idx, shape, nnz);
}

Expand Down
23 changes: 14 additions & 9 deletions onnxruntime/core/framework/execution_frame.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,15 @@ class NodeIndexInfo;

class IExecutionFrame {
protected:
IExecutionFrame(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers, const std::vector<int>& fetch_mlvalue_idxs,
const std::vector<OrtValue>& fetches, const OrtValueNameIdxMap& ort_value_idx_map,
const NodeIndexInfo& node_index_info);
// Derived class must call Init in its ctor. We need to use some of the virtual methods in Init and those aren't
// initialized until the derived class is constructed.
IExecutionFrame(const OrtValueNameIdxMap& ort_value_idx_map,
const NodeIndexInfo& node_index_info,
const std::vector<int>& fetch_mlvalue_idxs);

void Init(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::vector<OrtValue>& fetches);

public:
virtual ~IExecutionFrame();
Expand Down Expand Up @@ -72,18 +77,17 @@ class IExecutionFrame {
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IExecutionFrame);

void Init(const std::vector<int>& feed_mlvalue_idxs, const std::vector<OrtValue>& feeds,
const std::unordered_map<int, OrtValue>& initializers,
const std::vector<OrtValue>& fetches);

const OrtValue& GetMLValue(int ort_value_index) const {
ORT_ENFORCE(ort_value_index >= 0 && static_cast<size_t>(ort_value_index) < all_values_size_);
return all_values_[ort_value_index];
}

virtual AllocatorPtr GetAllocatorImpl(const OrtMemoryInfo& info) const = 0;

virtual Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) = 0;
virtual Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape,
size_t nnz) = 0;

virtual Status CopyTensor(const Tensor& src, Tensor& dest) const = 0;

const NodeIndexInfo& node_index_info_;

Expand Down Expand Up @@ -131,6 +135,7 @@ class ExecutionFrame final : public IExecutionFrame {
AllocatorPtr GetAllocatorImpl(const OrtMemoryInfo& info) const override;
Status ReleaseMLValueImpl(int ort_value_idx) override;
Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) override;
Status CopyTensor(const Tensor& src, Tensor& dest) const override;

common::Status AllocateAsPerAllocationPlan(OrtValue& ort_value, int ort_value_index, const TensorShape* shape,
size_t nnz);
Expand Down
15 changes: 10 additions & 5 deletions onnxruntime/core/optimizer/optimizer_execution_frame.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,12 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
size_t cpu_tensor_length;
ORT_RETURN_IF_ERROR(utils::GetSizeInBytesFromTensorProto<0>(tensor_proto, &cpu_tensor_length));
OrtValue ort_value;
const OrtMemoryInfo& info = cpu_execution_provider_->GetAllocator(0, OrtMemTypeDefault)->Info();
std::unique_ptr<char[]> data(new char[cpu_tensor_length]);
std::unique_ptr<Tensor> p_tensor;
OrtCallback d;
ORT_RETURN_IF_ERROR(utils::TensorProtoToMLValue(Env::Default(), nullptr, tensor_proto,
MemBuffer(data.get(), cpu_tensor_length, info), ort_value, d));
MemBuffer(data.get(), cpu_tensor_length, allocator_ptr_->Info()),
ort_value, d));

initializers_[idx] = ort_value;
buffer_for_initialized_tensors_[idx] = std::move(data);
Expand Down Expand Up @@ -83,14 +83,19 @@ std::unique_ptr<const OpKernel> OptimizerExecutionFrame::Info::CreateKernel(cons
// For optimizer, probably no need to pass feed_mlvalue_idxs, feeds to initialize IExecutionFrame.
// If needed, the parameters of OptimizerExecutionFrame ctor can be changed later.
OptimizerExecutionFrame::OptimizerExecutionFrame(const Info& info, const std::vector<int>& fetch_mlvalue_idxs)
: IExecutionFrame(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), fetch_mlvalue_idxs,
std::vector<OrtValue>(), info.GetMLValueNameIdxMap(), info.GetNodeIndexInfo()),
info_(info) {}
: IExecutionFrame(info.GetMLValueNameIdxMap(), info.GetNodeIndexInfo(), fetch_mlvalue_idxs),
info_(info) {
Init(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), std::vector<OrtValue>());
}

AllocatorPtr OptimizerExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {
return info_.GetAllocator(info);
}

Status OptimizerExecutionFrame::CopyTensor(const Tensor& src, Tensor& dest) const {
return info_.GetDataTransferManager().CopyTensor(src, dest);
}

// This method is not thread safe!
// Return S_OK and nullptr if index map to an value that is an unused optional input/output
Status OptimizerExecutionFrame::CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx,
Expand Down
4 changes: 4 additions & 0 deletions onnxruntime/core/optimizer/optimizer_execution_frame.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class OptimizerExecutionFrame final : public IExecutionFrame {

std::unique_ptr<const OpKernel> CreateKernel(const Node* node) const;

const DataTransferManager& GetDataTransferManager() const { return data_transfer_mgr_; }

private:
// The optimizer is running on CPU execution provider by default.
std::unique_ptr<CPUExecutionProvider> cpu_execution_provider_;
Expand Down Expand Up @@ -83,6 +85,8 @@ class OptimizerExecutionFrame final : public IExecutionFrame {

Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) override;

Status CopyTensor(const Tensor& src, Tensor& dest) const override;

const Info& info_;
};

Expand Down
Loading