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

Break dependency on SessionState for ExecutionFrame and OpKernelContext so optimizers can execute a node with a minimal setup (#498)

Browse source 
* Break dependency on SessionState for ExecutionFrame and OpKernelContext so optimizers can execute a node with a minimal setup.

- Create IExecutionFrame
  - split out core logic and interface from extended logic used in full Graph execution (that uses allocation plan and memory pattern planner)
- Update NodeIndexInfo to allow contruction from a subset of nodes
  - split out logic from GraphNodes into a re-usable template so it can be used with a vector of const Node* as well as a vector of unique_ptr<Node>
- Remove SessionState from OpKernelContext
- Misc cleanups
  - move AllocPlanPerValue out of SequentialExecutionPlan as it's used in a more generic manner that isn't specific to a sequential execution plan

NOTE: I manually tested the new paths, especially NodeIndexInfo. There will shortly be optimizers added that use the new infrastucture so they'll get test coverage as part of those changes.

* Fix linux build issue.
Handle graph with no nodes in NodeIndexInfo.
This commit is contained in:
Scott McKay 2019-02-27 15:46:50 -08:00 committed by GitHub
parent dfa21af302
commit 6c7099a18e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 524 additions and 434 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

101
include/onnxruntime/core/framework/op_kernel.h
View file

@ -20,7 +20,7 @@
#include "onnx/defs/schema.h"
namespace onnxruntime {
class ExecutionFrame;
class IExecutionFrame;
class OpKernelContext;
class OpKernelWrapper;
@ -35,7 +35,7 @@ class OpKernel {
return op_kernel_info_.node();
}
const ::onnxruntime::KernelDef& KernelDef() const {
const onnxruntime::KernelDef& KernelDef() const {
return op_kernel_info_.GetKernelDef();
}
@ -61,7 +61,7 @@ class OpKernelContext {
public:
using ArgMap = std::unordered_map<std::string, size_t>;
explicit OpKernelContext(ExecutionFrame* frame,
explicit OpKernelContext(IExecutionFrame* frame,
const OpKernel* kernel,
const logging::Logger& logger);
@ -147,20 +147,21 @@ class OpKernelContext {
protected:
onnxruntime::NodeIndex GetNodeIndex() const;
const SessionState& GetSessionState() const;
const MLValue* GetInputMLValue(int index) const;
const MLValue* GetImplicitInputMLValue(int index) const;
MLValue* GetOutputMLValue(int index);
private:
ORT_DISALLOW_COPY_AND_ASSIGNMENT(OpKernelContext);
Status GetOrCreateOutputMLValue(int index, MLValue*& value);
int GetInputArgIndex(int index) const;
int GetImplicitInputArgIndex(int index) const;
int GetOutputArgIndex(int index) const;
ExecutionFrame* execution_frame_{nullptr};
IExecutionFrame* execution_frame_{nullptr};
const OpKernel* kernel_{nullptr};
const logging::Logger* logger_{nullptr};
@ -221,18 +222,18 @@ KernelCreateInfo BuildKernelCreateInfo();
#define ONNX_CPU_OPERATOR_ML_KERNEL(name, ver, builder, ...) \
ONNX_OPERATOR_KERNEL_EX(name, kMLDomain, ver, kCpuExecutionProvider, builder, __VA_ARGS__)
#define ONNX_OPERATOR_KERNEL_EX(name, domain, ver, provider, builder, ...) \
class ONNX_OPERATOR_KERNEL_CLASS_NAME(provider, domain, ver, name); \
template <> \
KernelCreateInfo \
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(provider, domain, ver, name)>() { \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(ver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
#define ONNX_OPERATOR_KERNEL_EX(name, domain, ver, provider, builder, ...) \
class ONNX_OPERATOR_KERNEL_CLASS_NAME(provider, domain, ver, name); \
template <> \
KernelCreateInfo \
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(provider, domain, ver, name)>() { \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(ver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
}
#define ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(provider, domain, startver, endver, name) \
@ -244,18 +245,18 @@ KernelCreateInfo BuildKernelCreateInfo();
#define ONNX_CPU_OPERATOR_VERSIONED_ML_KERNEL(name, startver, endver, builder, ...) \
ONNX_OPERATOR_VERSIONED_KERNEL_EX(name, kMLDomain, startver, endver, kCpuExecutionProvider, builder, __VA_ARGS__)
#define ONNX_OPERATOR_VERSIONED_KERNEL_EX(name, domain, startver, endver, provider, builder, ...) \
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(provider, domain, startver, endver, name); \
template <> \
KernelCreateInfo \
#define ONNX_OPERATOR_VERSIONED_KERNEL_EX(name, domain, startver, endver, provider, builder, ...) \
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(provider, domain, startver, endver, name); \
template <> \
KernelCreateInfo \
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(provider, domain, startver, endver, name)>() { \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(startver, endver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(startver, endver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
}
#define ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(provider, domain, ver, type, name) \
@ -270,18 +271,18 @@ KernelCreateInfo BuildKernelCreateInfo();
#define ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(name, ver, type, builder, ...) \
ONNX_OPERATOR_TYPED_KERNEL_EX(name, kMSDomain, ver, type, kCpuExecutionProvider, builder, __VA_ARGS__)
#define ONNX_OPERATOR_TYPED_KERNEL_EX(name, domain, ver, type, provider, builder, ...) \
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(provider, domain, ver, type, name); \
template <> \
KernelCreateInfo \
#define ONNX_OPERATOR_TYPED_KERNEL_EX(name, domain, ver, type, provider, builder, ...) \
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(provider, domain, ver, type, name); \
template <> \
KernelCreateInfo \
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(provider, domain, ver, type, name)>() { \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(ver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(ver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
}
#define ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(provider, domain, startver, endver, type, name) \
@ -296,18 +297,18 @@ KernelCreateInfo BuildKernelCreateInfo();
#define ONNX_CPU_OPERATOR_VERSIONED_TYPED_MS_KERNEL(name, startver, endver, type, builder, ...) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(name, kMSDomain, startver, endver, type, kCpuExecutionProvider, builder, __VA_ARGS__)
#define ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(name, domain, startver, endver, type, provider, builder, ...) \
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(provider, domain, startver, endver, type, name); \
template <> \
KernelCreateInfo \
#define ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(name, domain, startver, endver, type, provider, builder, ...) \
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(provider, domain, startver, endver, type, name); \
template <> \
KernelCreateInfo \
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(provider, domain, startver, endver, type, name)>() { \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(startver, endver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
return KernelCreateInfo( \
builder.SetName(#name) \
.SetDomain(domain) \
.SinceVersion(startver, endver) \
.Provider(provider) \
.Build(), \
[](const OpKernelInfo& info) -> OpKernel* { return new __VA_ARGS__(info); }); \
}
} // namespace onnxruntime

26
include/onnxruntime/core/graph/graph_nodes.h
View file

@ -12,23 +12,23 @@ namespace onnxruntime {
class Node;
/**
Class that provides iteration over all valid nodes in the Graph.
Class to filter out null entries from either a vector of unique_ptr<Node> or a vector of [const] Node* and
provide an iterator interface that returns [const] Node& for the valid entries.
*/
class GraphNodes {
using TNodesContainer = std::vector<std::unique_ptr<Node>>;
template <typename TNodesContainer>
class ValidNodes {
public:
template <typename TIterator>
class NodeIterator;
/**
Construct a GraphNodes instance to provide iteration over all valid nodes in the Graph
Construct a ValidNodes instance to provide iteration over all valid nodes in the TNodesCollection
@param[in] nodes Nodes to iterate, skipping invalid entries.
*/
explicit GraphNodes(TNodesContainer& nodes) noexcept : nodes_(nodes) {}
explicit ValidNodes(TNodesContainer& nodes) noexcept : nodes_(nodes) {}
using ConstNodeIterator = NodeIterator<TNodesContainer::const_iterator>;
using MutableNodeIterator = NodeIterator<TNodesContainer::iterator>;
using ConstNodeIterator = NodeIterator<typename TNodesContainer::const_iterator>;
using MutableNodeIterator = NodeIterator<typename TNodesContainer::iterator>;
ConstNodeIterator cbegin() const noexcept {
return {nodes_.cbegin(), nodes_.cend()};
@ -54,6 +54,8 @@ class GraphNodes {
return {nodes_.end(), nodes_.end()};
}
bool empty() const noexcept { return nodes_.empty(); }
/**
@class NodeIterator
Iterator to provide const and non-const access to valid Node instances in a Graph.
@ -127,4 +129,12 @@ class GraphNodes {
TNodesContainer& nodes_;
};
/**
Class that provides iteration over all valid nodes in the Graph.
*/
class GraphNodes : public ValidNodes<std::vector<std::unique_ptr<Node>>> {
public:
GraphNodes(std::vector<std::unique_ptr<Node>>& nodes) : ValidNodes(nodes) {}
};
} // namespace onnxruntime

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

@ -161,11 +161,11 @@ class PlannerImpl {
MLValueIndex& Buffer(MLValueIndex n) { return ml_value_info_.at(n).reused_buffer_index; }
SequentialExecutionPlan::AllocPlanPerValue& AllocPlan(MLValueIndex n) {
AllocPlanPerValue& AllocPlan(MLValueIndex n) {
return plan_.allocation_plan.at(n);
}
SequentialExecutionPlan::AllocPlanPerValue& AllocPlan(const MLValueName& name) {
AllocPlanPerValue& AllocPlan(const MLValueName& name) {
return AllocPlan(Index(name));
}
@ -437,7 +437,7 @@ class PlannerImpl {
if (!weights.count(def_name)) return Status::OK();
auto wt_index = Index(def_name);
SequentialExecutionPlan::AllocPlanPerValue& thisplan = AllocPlan(wt_index);
AllocPlanPerValue& thisplan = AllocPlan(wt_index);
auto* p_provider = execution_providers_.Get(node);
ORT_ENFORCE(p_provider);
@ -466,7 +466,7 @@ class PlannerImpl {
auto setup_preexisting = [this](const NodeArg* node_arg) {
auto input_index = Index(node_arg->Name());
SequentialExecutionPlan::AllocPlanPerValue& thisplan = AllocPlan(input_index);
AllocPlanPerValue& thisplan = AllocPlan(input_index);
thisplan.alloc_kind = AllocKind::kPreExisting;
thisplan.value_type = utils::GetMLDataType(*node_arg);
};

501
onnxruntime/core/framework/execution_frame.cc
View file

@ -16,24 +16,177 @@ using namespace onnxruntime::common;
namespace onnxruntime {
IExecutionFrame::IExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::unordered_map<int, MLValue>& initializers,
const std::vector<int>& fetch_mlvalue_idxs,
std::vector<MLValue>& fetches,
const MLValueNameIdxMap& mlvalue_idx_map,
const NodeIndexInfo& node_index_info)
: node_index_info_{node_index_info}, fetch_mlvalue_idxs_{fetch_mlvalue_idxs} {
ORT_ENFORCE(feeds.size() == feed_mlvalue_idxs.size());
ORT_ENFORCE(fetches.empty() || fetches.size() == fetch_mlvalue_idxs.size());
Init(feed_mlvalue_idxs, feeds, initializers, fetch_mlvalue_idxs, fetches, mlvalue_idx_map);
}
IExecutionFrame::~IExecutionFrame() = default;
// Return nullptr if index map to an value that is an unused optional input/output
const MLValue* IExecutionFrame::GetNodeInputOrOutputMLValue(int index) const {
int mlvalue_idx = GetNodeIdxToMLValueIdx(index);
return mlvalue_idx != NodeIndexInfo::kInvalidEntry ? &all_values_[mlvalue_idx] : nullptr;
}
MLValue* IExecutionFrame::GetMutableNodeInputOrOutputMLValue(int index) {
return const_cast<MLValue*>(GetNodeInputOrOutputMLValue(index));
}
// TO DO: make it thread safe
// 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 IExecutionFrame::GetOrCreateNodeOutputMLValue(int index, const TensorShape* shape, MLValue*& p_mlvalue) {
auto status = Status::OK();
int mlvalue_idx = GetNodeIdxToMLValueIdx(index);
// return nullptr if it is optional
if (mlvalue_idx == NodeIndexInfo::kInvalidEntry) {
p_mlvalue = nullptr;
} else {
p_mlvalue = &all_values_[mlvalue_idx];
if (p_mlvalue->IsAllocated()) {
// already allocated. verify shape matches if tensor.
if (p_mlvalue->IsTensor()) {
const Tensor& tensor = p_mlvalue->Get<Tensor>();
ORT_ENFORCE(shape && tensor.Shape() == *shape,
"MLValue shape verification failed. Current shape:", tensor.Shape(),
" Requested shape:", shape ? shape->ToString() : "null");
}
} else {
status = CreateNodeOutputMLValueImpl(*p_mlvalue, mlvalue_idx, shape);
}
}
return status;
}
AllocatorPtr IExecutionFrame::GetAllocator(const OrtAllocatorInfo& info) const {
return GetAllocatorImpl(info);
}
Status IExecutionFrame::ReleaseMLValue(int mlvalue_idx) {
return ReleaseMLValueImpl(mlvalue_idx);
}
Status IExecutionFrame::ReleaseMLValueImpl(int mlvalue_idx) {
if (mlvalue_idx == NodeIndexInfo::kInvalidEntry || static_cast<size_t>(mlvalue_idx) >= all_values_.size()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index ", mlvalue_idx);
}
all_values_[mlvalue_idx] = MLValue();
return Status::OK();
}
int IExecutionFrame::GetNodeIdxToMLValueIdx(int index) const {
int mlvalue_idx = node_index_info_.GetMLValueIndex(index);
ORT_ENFORCE(mlvalue_idx == NodeIndexInfo::kInvalidEntry ||
(mlvalue_idx >= 0 && static_cast<size_t>(mlvalue_idx) < all_values_.size()));
return mlvalue_idx;
}
void IExecutionFrame::Init(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::unordered_map<int, MLValue>& initializers,
const std::vector<int>& fetch_mlvalue_idxs,
const std::vector<MLValue>& fetches,
const MLValueNameIdxMap& mlvalue_idx_map) {
// 1. resize the all_value_ vector
all_values_.resize(mlvalue_idx_map.MaxIdx() + 1);
// 2. Handle non-empty output vector
if (!fetches.empty()) {
auto num_fetches = fetch_mlvalue_idxs.size();
for (size_t idx = 0; idx < num_fetches; ++idx) {
int mlvalue_idx = fetch_mlvalue_idxs[idx];
all_values_[mlvalue_idx] = fetches[idx];
}
}
// 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_).
// A non-empty fetches vector will overwrite the actual weight in all_values_[mlvalue_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) {
auto mlvalue_index = entry.first;
all_values_[mlvalue_index] = entry.second;
}
// 4. handle feed in values. these can override initializer values so must be last
for (size_t idx = 0, end = feed_mlvalue_idxs.size(); idx < end; ++idx) {
int mlvalue_idx = feed_mlvalue_idxs[idx];
// we are sharing the underline tensor/object for MLValue
all_values_[mlvalue_idx] = feeds[idx];
}
}
Status IExecutionFrame::GetOutputs(std::vector<MLValue>& fetches) {
auto num_fetches = fetch_mlvalue_idxs_.size();
if (fetches.empty()) {
fetches.resize(num_fetches);
} else {
// if there's a mismatch things are out so sync so fail
if (fetches.size() != num_fetches) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Fetches vector passed to GetOutputs contains ", fetches.size(),
" entries which doesn't match the number of fetches the frame was initialized with of ",
num_fetches);
}
}
for (size_t idx = 0; idx < num_fetches; ++idx) {
fetches[idx] = GetMLValue(fetch_mlvalue_idxs_[idx]);
}
return Status::OK();
}
bool IExecutionFrame::IsOutput(int mlvalue_idx) const {
return std::find(fetch_mlvalue_idxs_.begin(), fetch_mlvalue_idxs_.end(), mlvalue_idx) != fetch_mlvalue_idxs_.end();
}
ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::vector<int>& fetch_mlvalue_idxs,
std::vector<MLValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
const SessionState& session_state)
: node_index_info_(session_state.GetNodeIndexInfo()),
session_state_(session_state),
mem_patterns_(nullptr),
planner_(nullptr),
fetch_mlvalue_idxs_{fetch_mlvalue_idxs} {
Init(feed_mlvalue_idxs, feeds, fetch_mlvalue_idxs, fetches, fetch_allocators);
: IExecutionFrame(feed_mlvalue_idxs, feeds, session_state.GetInitializedTensors(), fetch_mlvalue_idxs, fetches,
session_state.GetMLValueNameIdxMap(), session_state.GetNodeIndexInfo()),
session_state_{session_state},
mem_patterns_{nullptr},
planner_{nullptr} {
// map the custom allocators to mlvalue_idx entries
if (!fetch_allocators.empty()) {
for (size_t idx = 0, end = fetch_mlvalue_idxs.size(); idx < end; ++idx) {
int mlvalue_idx = fetch_mlvalue_idxs[idx];
auto custom_alloc_entry = fetch_allocators.find(idx);
if (custom_alloc_entry != fetch_allocators.cend()) {
custom_allocators_[mlvalue_idx] = custom_alloc_entry->second;
}
}
}
// If the session enable memory pattern optimization
// and we have execution plan generated, try to setup
// memory pattern optimization.
if (session_state.GetEnableMemoryPattern() &&
session_state.GetExecutionPlan()) {
if (session_state.GetEnableMemoryPattern() && session_state.GetExecutionPlan()) {
std::vector<TensorShape> input_shapes;
bool all_tensors = true;
for (const auto& feed : feeds) {
@ -44,8 +197,8 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
auto& tensor = feed.Get<Tensor>();
input_shapes.push_back(tensor.Shape());
}
// if there is some traditional ml value type in inputs
// disable the memory pattern optimization.
// if there are some traditional ml value type in inputs disable the memory pattern optimization.
if (all_tensors) {
mem_patterns_ = session_state.GetMemoryPatternGroup(input_shapes);
// if no existing patterns, generate one in this executionframe
@ -57,7 +210,9 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
for (size_t i = 0; i < mem_patterns_->locations.size(); i++) {
ORT_ENFORCE(buffers_.find(mem_patterns_->locations[i]) == buffers_.end());
AllocatorPtr alloc = GetAllocator(mem_patterns_->locations[i]);
void* buffer = mem_patterns_->patterns[i].PeakSize() > 0 ? alloc->Alloc(mem_patterns_->patterns[i].PeakSize()) : nullptr;
void* buffer = mem_patterns_->patterns[i].PeakSize() > 0
? alloc->Alloc(mem_patterns_->patterns[i].PeakSize())
: nullptr;
buffers_[mem_patterns_->locations[i]] = BufferUniquePtr(buffer, alloc);
}
}
@ -67,45 +222,43 @@ ExecutionFrame::ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
ExecutionFrame::~ExecutionFrame() = default;
Status ExecutionFrame::AllocateMLValueTensorSelfOwnBuffer(int mlvalue_index,
const DataTypeImpl* element_type,
Status ExecutionFrame::AllocateMLValueTensorSelfOwnBuffer(MLValue& mlvalue,
int mlvalue_index,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence) {
ORT_ENFORCE(mlvalue_index >= 0 && static_cast<size_t>(mlvalue_index) < all_values_.size());
return AllocateMLValueTensorSelfOwnBufferHelper(mlvalue_index, element_type, location, shape, create_fence);
return AllocateMLValueTensorSelfOwnBufferHelper(mlvalue, mlvalue_index, element_type, location, shape, create_fence);
}
Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(int mlvalue_index,
const DataTypeImpl* element_type,
Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(MLValue& mlvalue,
int mlvalue_index,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence) {
if (mlvalue_index == NodeIndexInfo::kInvalidEntry)
if (mlvalue_index == NodeIndexInfo::kInvalidEntry) {
return Status(ONNXRUNTIME, FAIL, "Trying to allocate memory for unused optional inputs/outputs");
}
auto p_mlvalue = &all_values_[mlvalue_index];
if (p_mlvalue->IsAllocated()) {
return Status::OK();
}
auto alloc = GetAllocator(location);
size_t size;
{
int64_t len = shape.Size();
if (len < 0) {
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Tensor shape cannot contain any negative value");
}
if (!IAllocator::CalcMemSizeForArrayWithAlignment<64>(len, element_type->Size(), &size)) {
return Status(ONNXRUNTIME, FAIL, "size overflow");
}
int64_t len = shape.Size();
if (len < 0) {
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Tensor shape cannot contain any negative value");
}
if (!IAllocator::CalcMemSizeForArrayWithAlignment<64>(len, element_type->Size(), &size)) {
return Status(ONNXRUNTIME, FAIL, "size overflow");
}
auto alloc = GetAllocator(location);
// create fence if needed
if (create_fence) {
ORT_ENFORCE(p_mlvalue->Fence() == nullptr);
FencePtr f = alloc->CreateFence(&GetSessionState());
ORT_ENFORCE(mlvalue.Fence() == nullptr);
FencePtr f = alloc->CreateFence(&session_state_);
// it is OK to have fence been nullptr if the execution provider has no async execution,
// and allocator::CreateFence returns nullptr
p_mlvalue->SetFence(f);
mlvalue.SetFence(f);
}
// if we have pre-calculated memory pattern, and the mlvalue is not output mlvalue
@ -122,122 +275,88 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(int mlvalue_inde
if (it != buffers_.end() && block->size_ == size) {
void* buffer = it->second.get();
auto status = AllocateTensorWithPreAllocateBufferHelper(
p_mlvalue, static_cast<void*>(static_cast<char*>(buffer) + block->offset_),
mlvalue, static_cast<void*>(static_cast<char*>(buffer) + block->offset_),
element_type, location, shape);
return status;
}
if (block->size_ != size) {
LOGS_DEFAULT(WARNING) << "For mlvalue with index: " << mlvalue_index << ", block in memory pattern size is: "
<< block->size_ << " but the actually size is: " << size << ", fall back to default allocation behavior";
<< block->size_ << " but the actually size is: " << size
<< ", fall back to default allocation behavior";
} else if (it == buffers_.end()) {
LOGS_DEFAULT(WARNING) << "For mlvalue with index: " << mlvalue_index << ", block not found in target loation. "
" fall back to default allocation behavior";
LOGS_DEFAULT(WARNING) << "For mlvalue with index: " << mlvalue_index
<< ", block not found in target location. fall back to default allocation behavior";
}
}
}
}
//no memory pattern, or the pattern is not correct.
void* buffer = size == 0 ? nullptr : alloc->Alloc(size);
std::unique_ptr<Tensor> p_tensor = std::make_unique<Tensor>(element_type,
shape,
buffer,
location,
alloc);
auto p_tensor = std::make_unique<Tensor>(element_type, shape, buffer, location, alloc);
p_mlvalue->Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
mlvalue.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
// trace the memory allocation.
// don't trace the memory allocation on string tensors, as it need
// placement new, we don't support it in memory pattern optimization.
if (element_type != DataTypeImpl::GetType<std::string>())
if (element_type != DataTypeImpl::GetType<std::string>()) {
TraceAllocate(mlvalue_index, size);
}
return Status::OK();
}
void ExecutionFrame::TraceAllocate(int mlvalue_idx, size_t size) {
// don't trace the output tensors.
auto& allocation_plan = GetAllocationPlan(mlvalue_idx);
if (planner_ && allocation_plan.alloc_kind != AllocKind::kAllocateOutput) {
auto status = planner_->TraceAllocation(mlvalue_idx, size);
if (!status.IsOK())
LOGS(session_state_.Logger(), WARNING) << "TraceAllocation for mlvalue_idx=" << mlvalue_idx << " size=" << size
<< " failed: " << status.ErrorMessage();
}
}
Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(int mlvalue_index_to_allocate,
Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(MLValue& mlvalue,
int mlvalue_index_reuse,
const DataTypeImpl* element_type,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence) {
ORT_ENFORCE(mlvalue_index_to_allocate >= 0 && mlvalue_index_to_allocate < all_values_.size());
MLValue* p_mlvalue = &all_values_[mlvalue_index_to_allocate];
MLValue& mlvalue_reuse = GetMutableMLValue(mlvalue_index_reuse);
ORT_ENFORCE(mlvalue_index_reuse >= 0 && mlvalue_index_reuse < all_values_.size());
MLValue* p_mlvalue_reuse = &all_values_[mlvalue_index_reuse];
auto* reuse_tensor = p_mlvalue_reuse->GetMutable<Tensor>();
auto* reuse_tensor = mlvalue_reuse.GetMutable<Tensor>();
void* reuse_buffer = reuse_tensor->MutableDataRaw();
// create fence on reused mlvalue if needed
// TODO: differentiate reuse and alias, by add AllocKind::kAlias?
if (create_fence && p_mlvalue_reuse->Fence() == nullptr) {
FencePtr f = GetAllocator(location)->CreateFence(&GetSessionState());
p_mlvalue_reuse->SetFence(f);
if (create_fence && mlvalue_reuse.Fence() == nullptr) {
FencePtr f = GetAllocator(location)->CreateFence(&session_state_);
mlvalue_reuse.SetFence(f);
}
// reused MLValue share the same fence
p_mlvalue->ShareFenceWith(*p_mlvalue_reuse);
return AllocateTensorWithPreAllocateBufferHelper(p_mlvalue, reuse_buffer, element_type, location, shape);
mlvalue.ShareFenceWith(mlvalue_reuse);
return AllocateTensorWithPreAllocateBufferHelper(mlvalue, reuse_buffer, element_type, location, shape);
}
Status ExecutionFrame::AllocateTensorWithPreAllocateBufferHelper(MLValue* p_mlvalue,
Status ExecutionFrame::AllocateTensorWithPreAllocateBufferHelper(MLValue& mlvalue,
void* pBuffer,
const DataTypeImpl* element_type,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape) {
if (p_mlvalue->IsAllocated()) {
return Status::OK();
}
std::unique_ptr<Tensor> p_tensor = std::make_unique<Tensor>(element_type,
shape,
pBuffer,
location);
p_mlvalue->Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
auto p_tensor = std::make_unique<Tensor>(element_type, shape, pBuffer, location);
mlvalue.Init(p_tensor.release(),
DataTypeImpl::GetType<Tensor>(),
DataTypeImpl::GetType<Tensor>()->GetDeleteFunc());
return Status::OK();
}
Status AllocateTraditionalMLValue(MLValue* p_mlvalue,
const NonTensorTypeBase* type,
const MLValueAllocationParameters& parameters) {
// right now we don't need any parameter for ml value creation,
// keep it in api for extensibility
ORT_UNUSED_PARAMETER(parameters);
auto creator = type->GetCreateFunc();
p_mlvalue->Init(creator(),
type,
type->GetDeleteFunc());
static Status AllocateTraditionalMLValue(MLValue& mlvalue, const NonTensorTypeBase& type) {
auto creator = type.GetCreateFunc();
mlvalue.Init(creator(), &type, type.GetDeleteFunc());
return Status::OK();
}
// This method is not thread safe!
Status ExecutionFrame::AllocateAsPerAllocationPlan(int mlvalue_index,
const MLValueAllocationParameters& parameters) {
if (mlvalue_index == NodeIndexInfo::kInvalidEntry || mlvalue_index >= all_values_.size())
return Status(ONNXRUNTIME, INVALID_ARGUMENT,
"Tried to allocated with invalid mlvalue index: " + std::to_string(mlvalue_index));
Status ExecutionFrame::AllocateAsPerAllocationPlan(MLValue& mlvalue, int mlvalue_index, const TensorShape* shape) {
// if there is a custom allocator for this mlvalue_index, call it to do the allocation
auto custom_alloc_entry = custom_allocators_.find(mlvalue_index);
if (custom_alloc_entry != custom_allocators_.cend()) {
return (custom_alloc_entry->second)(parameters.GetTensorShape(), all_values_[mlvalue_index]);
ORT_ENFORCE(shape, "We don't expect custom allocators for non-tensor types, so a shape is mandatory here.");
return (custom_alloc_entry->second)(*shape, mlvalue);
}
const SequentialExecutionPlan* p_seq_exec_plan = session_state_.GetExecutionPlan();
@ -250,12 +369,13 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(int mlvalue_index,
if (ml_type == nullptr)
return Status(ONNXRUNTIME, INVALID_ARGUMENT,
"Tried to allocate without valid type information, mlvalue index=" + std::to_string(mlvalue_index));
if (!ml_type->IsTensorType()) {
return AllocateTraditionalMLValue(&all_values_[mlvalue_index],
static_cast<const NonTensorTypeBase*>(ml_type),
parameters);
return AllocateTraditionalMLValue(mlvalue, *static_cast<const NonTensorTypeBase*>(ml_type));
}
ORT_ENFORCE(shape, "Allocation of tensor types requires a shape.");
// tensors
auto ml_data_type = static_cast<const TensorTypeBase*>(ml_type)->GetElementType();
@ -265,20 +385,14 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(int mlvalue_index,
// In the future we may want to have different way to handle it.
case AllocKind::kAllocateOutput:
case AllocKind::kAllocate: {
ORT_RETURN_IF_ERROR(AllocateMLValueTensorSelfOwnBuffer(mlvalue_index,
ml_data_type,
alloc_info,
parameters.GetTensorShape(),
ORT_RETURN_IF_ERROR(AllocateMLValueTensorSelfOwnBuffer(mlvalue, mlvalue_index, ml_data_type, alloc_info, *shape,
per_alloc_plan.create_fence_if_async));
break;
}
case AllocKind::kReuse: {
int reuse_mlvalue_index = per_alloc_plan.reused_buffer;
ORT_RETURN_IF_ERROR(AllocateMLValueTensorPreAllocateBuffer(mlvalue_index,
reuse_mlvalue_index,
ml_data_type,
alloc_info,
parameters.GetTensorShape(),
ORT_RETURN_IF_ERROR(AllocateMLValueTensorPreAllocateBuffer(mlvalue, reuse_mlvalue_index,
ml_data_type, alloc_info, *shape,
per_alloc_plan.create_fence_if_async));
break;
}
@ -292,56 +406,43 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(int mlvalue_index,
return Status::OK();
}
void ExecutionFrame::Init(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::vector<int>& fetch_mlvalue_idxs,
std::vector<MLValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators) {
auto& mlvalue_idx_map = session_state_.GetMLValueNameIdxMap();
AllocatorPtr ExecutionFrame::GetAllocatorImpl(const OrtAllocatorInfo& info) const {
return utils::GetAllocator(session_state_, info);
}
// 1. resize the all_value_ vector
all_values_.resize(mlvalue_idx_map.MaxIdx() + 1);
// 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(MLValue& mlvalue, int mlvalue_idx, const TensorShape* shape) {
return AllocateAsPerAllocationPlan(mlvalue, mlvalue_idx, shape);
}
// 2. Handle non-empty output vector
if (!fetches.empty()) {
// setup output_indices_, we don't want to generate mem plan on output tensors.
auto num_fetches = fetch_mlvalue_idxs.size();
Status ExecutionFrame::ReleaseMLValueImpl(int mlvalue_idx) {
IExecutionFrame::ReleaseMLValueImpl(mlvalue_idx);
TraceFree(mlvalue_idx);
return Status::OK();
}
for (size_t idx = 0; idx < num_fetches; ++idx) {
int mlvalue_idx = fetch_mlvalue_idxs[idx];
all_values_[mlvalue_idx] = fetches[idx];
const AllocPlanPerValue& ExecutionFrame::GetAllocationPlan(int mlvalue_idx) {
const SequentialExecutionPlan* p_seq_exec_plan = session_state_.GetExecutionPlan();
const auto& alloc_plan = p_seq_exec_plan->allocation_plan;
ORT_ENFORCE(mlvalue_idx != NodeIndexInfo::kInvalidEntry && mlvalue_idx < alloc_plan.size());
return alloc_plan[mlvalue_idx];
}
auto custom_alloc_entry = fetch_allocators.find(idx);
if (custom_alloc_entry != fetch_allocators.cend()) {
custom_allocators_[mlvalue_idx] = custom_alloc_entry->second;
}
}
}
// 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_).
// A non-empty fetches vector will overwrite the actual weight in all_values_[mlvalue_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 : session_state_.GetInitializedTensors()) {
auto mlvalue_index = entry.first;
all_values_[mlvalue_index] = entry.second;
}
// 4. handle feed in values. these can override initializer values so must be last
for (size_t idx = 0, end = feed_mlvalue_idxs.size(); idx < end; ++idx) {
int mlvalue_idx = feed_mlvalue_idxs[idx];
// we are sharing the underline tensor/object for MLValue
all_values_[mlvalue_idx] = feeds[idx];
void ExecutionFrame::TraceAllocate(int mlvalue_idx, size_t size) {
// don't trace the output tensors.
auto& allocation_plan = GetAllocationPlan(mlvalue_idx);
if (planner_ && allocation_plan.alloc_kind != AllocKind::kAllocateOutput) {
auto status = planner_->TraceAllocation(mlvalue_idx, size);
if (!status.IsOK())
LOGS(session_state_.Logger(), WARNING) << "TraceAllocation for mlvalue_idx=" << mlvalue_idx << " size=" << size
<< " failed: " << status.ErrorMessage();
}
}
void ExecutionFrame::TraceFree(int mlvalue_idx) {
// don't trace free on output tensors.
if (planner_ &&
std::find(fetch_mlvalue_idxs_.begin(), fetch_mlvalue_idxs_.end(), mlvalue_idx) == fetch_mlvalue_idxs_.end()) {
if (planner_ && !IsOutput(mlvalue_idx)) {
const SequentialExecutionPlan* p_seq_exec_plan = session_state_.GetExecutionPlan();
const auto& alloc_plan = p_seq_exec_plan->allocation_plan;
const auto& per_alloc_plan = alloc_plan.at(mlvalue_idx);
@ -373,98 +474,4 @@ Status ExecutionFrame::GeneratePatterns(MemoryPatternGroup* out) const {
return planner_->GeneratePatterns(out);
}
int ExecutionFrame::GetNodeOffset(onnxruntime::NodeIndex node_index) const {
return node_index_info_.GetNodeOffset(node_index);
}
// Return nullptr if index map to an value that is an unused optional input/output
const MLValue* ExecutionFrame::GetNodeInputOrOutputMLValue(int index) const {
int mlvalue_idx = node_index_info_.GetMLValueIndex(index);
return mlvalue_idx != NodeIndexInfo::kInvalidEntry ? &all_values_[mlvalue_idx] : nullptr;
}
// Return nullptr if index map to an value that is an unused optional input/output
MLValue* ExecutionFrame::GetMutableNodeInputOrOutputMLValue(int index) {
return const_cast<MLValue*>(GetNodeInputOrOutputMLValue(index));
}
AllocatorPtr ExecutionFrame::GetAllocator(const OrtAllocatorInfo& info) {
return utils::GetAllocator(session_state_, info);
}
static inline void VerifyShape(const MLValue* p_mlvalue,
const MLValueAllocationParameters& parameters) {
if (p_mlvalue->IsTensor()) {
const Tensor* tensor = &p_mlvalue->Get<Tensor>();
ORT_ENFORCE(tensor->Shape() == parameters.GetTensorShape(),
"MLValue shape verification failed. Current shape:", tensor->Shape(),
" Requested shape:", parameters.GetTensorShape());
}
}
// 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::GetOrCreateNodeOutputMLValue(int index,
const MLValueAllocationParameters& parameters,
MLValue*& p_mlvalue) {
int mlvalue_idx = node_index_info_.GetMLValueIndex(index);
// return nullptr if it is optional
if (mlvalue_idx == NodeIndexInfo::kInvalidEntry) {
p_mlvalue = nullptr;
return Status::OK();
}
p_mlvalue = &all_values_.at(mlvalue_idx);
if (p_mlvalue->IsAllocated()) {
// The ml has already been allocated.
// Now only tensor need to be check.
VerifyShape(p_mlvalue, parameters); // TODO find a better way to do this
return Status::OK();
}
// It's not allocated, then allocate it with given shape and return.
// Perform allocation based on the allocation plan
ORT_RETURN_IF_ERROR(AllocateAsPerAllocationPlan(mlvalue_idx, parameters));
return Status::OK();
}
Status ExecutionFrame::GetOutputs(std::vector<MLValue>& fetches) {
auto num_fetches = fetch_mlvalue_idxs_.size();
if (fetches.empty()) {
fetches.resize(num_fetches);
} else {
// if there's a mismatch things are out so sync so fail
if (fetches.size() != num_fetches) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Fetches vector passed to GetOutputs contains ", fetches.size(),
" entries which doesn't match the number of fetches the frame was initialized with of ",
num_fetches);
}
}
for (size_t idx = 0; idx < num_fetches; ++idx) {
fetches[idx] = GetMLValue(fetch_mlvalue_idxs_[idx]);
}
return Status::OK();
}
Status ExecutionFrame::ReleaseMLValue(int mlvalue_idx) {
if (mlvalue_idx == NodeIndexInfo::kInvalidEntry || static_cast<size_t>(mlvalue_idx) >= all_values_.size()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "invalid index ", mlvalue_idx);
}
all_values_[mlvalue_idx] = MLValue();
TraceFree(mlvalue_idx);
return Status::OK();
}
const SequentialExecutionPlan::AllocPlanPerValue& ExecutionFrame::GetAllocationPlan(int mlvalue_idx) {
const SequentialExecutionPlan* p_seq_exec_plan = session_state_.GetExecutionPlan();
const auto& alloc_plan = p_seq_exec_plan->allocation_plan;
ORT_ENFORCE(mlvalue_idx != NodeIndexInfo::kInvalidEntry && mlvalue_idx < alloc_plan.size());
return alloc_plan[mlvalue_idx];
}
} // namespace onnxruntime

161
onnxruntime/core/framework/execution_frame.h
View file

@ -10,6 +10,7 @@
#include "core/common/status.h"
#include "core/framework/iexecutor.h"
#include "core/framework/ml_value.h"
#include "core/framework/node_index_info.h"
#include "core/framework/sequential_execution_plan.h"
#include "core/framework/tensor.h"
#include "core/graph/graph_viewer.h"
@ -17,26 +18,87 @@
namespace onnxruntime {
class SessionState;
class MLValueNameIdxMap;
class MLValuePatternPlanner;
struct MemoryPatternGroup;
class NodeIndexInfo;
struct MLValueAllocationParameters {
MLValueAllocationParameters() = default;
MLValueAllocationParameters(const TensorShape* shape)
: tensor_shape{shape} {}
class IExecutionFrame {
protected:
IExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::unordered_map<int, MLValue>& initializers,
const std::vector<int>& fetch_mlvalue_idxs,
std::vector<MLValue>& fetches,
const MLValueNameIdxMap& mlvalue_idx_map,
const NodeIndexInfo& node_index_info);
const TensorShape& GetTensorShape() const {
static const TensorShape s_empty_tensor_shape;
return tensor_shape != nullptr ? *tensor_shape : s_empty_tensor_shape;
public:
virtual ~IExecutionFrame();
// Get the index for the first entry of the given node.
int GetNodeOffset(NodeIndex index) const {
return node_index_info_.GetNodeOffset(index);
}
// Return nullptr if index map to an value that is an unused optional input/output
const MLValue* GetNodeInputOrOutputMLValue(int index) const;
MLValue* GetMutableNodeInputOrOutputMLValue(int index);
// TO DO: make it thread safe
// This method is not thread safe!
// Return S_OK and nullptr if index map to an value that is an unused optional input/output
// Shape is required for tensors but not traditional ML values.
Status GetOrCreateNodeOutputMLValue(int index, const TensorShape* shape, MLValue*& p_mlvalue);
// write the output values to the 'fetches' vector
Status GetOutputs(std::vector<MLValue>& fetches);
AllocatorPtr GetAllocator(const OrtAllocatorInfo& info) const;
Status ReleaseMLValue(int mlvalue_idx);
protected:
// get the mlvalue_idx from NodeIndexInfo
int GetNodeIdxToMLValueIdx(int index) const;
MLValue& GetMutableMLValue(int mlvalue_index) {
return const_cast<MLValue&>(GetMLValue(mlvalue_index));
}
virtual Status ReleaseMLValueImpl(int mlvalue_idx);
// returns true if the mlvalue_idx is an output from the graph
bool IsOutput(int mlvalue_idx) const;
private:
const TensorShape* tensor_shape{};
// todo: is there any parameter needed for ml types?
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IExecutionFrame);
void Init(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::unordered_map<int, MLValue>& initializers,
const std::vector<int>& fetch_mlvalue_idxs,
const std::vector<MLValue>& fetches,
const MLValueNameIdxMap& mlvalue_idx_map);
const MLValue& GetMLValue(int mlvalue_index) const {
ORT_ENFORCE(mlvalue_index >= 0 && static_cast<size_t>(mlvalue_index) < all_values_.size());
return all_values_[mlvalue_index];
}
virtual AllocatorPtr GetAllocatorImpl(const OrtAllocatorInfo& info) const = 0;
virtual Status CreateNodeOutputMLValueImpl(MLValue& mlvalue, int mlvalue_idx, const TensorShape* shape) = 0;
const NodeIndexInfo& node_index_info_;
// All the intermediate values for the entire graph.
// Input and Output values are passed in by executors
std::vector<MLValue> all_values_;
const std::vector<int> fetch_mlvalue_idxs_;
};
class ExecutionFrame {
class ExecutionFrame final : public IExecutionFrame {
public:
ExecutionFrame(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
@ -51,105 +113,60 @@ class ExecutionFrame {
// TODO: These two AllocateMLValue... methods are in the API purely for unit test usage.
// Fix the unit tests so they set an execution plan that results in these methods being called by
// GetOrCreateNodeOutputMLValue instead
Status AllocateMLValueTensorSelfOwnBuffer(int mlvalue_index,
Status AllocateMLValueTensorSelfOwnBuffer(MLValue& mlvalue,
int mlvalue_index,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence = false);
Status AllocateMLValueTensorPreAllocateBuffer(int mlvalue_index_to_allocate,
Status AllocateMLValueTensorPreAllocateBuffer(MLValue& mlvalue,
int mlvalue_index_reuse,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence = false);
const MLValue& GetMLValue(int mlvalue_index) const {
ORT_ENFORCE(mlvalue_index >= 0 && static_cast<size_t>(mlvalue_index) < all_values_.size());
return all_values_[mlvalue_index];
}
MLValue& GetMutableMLValue(int mlvalue_index) {
ORT_ENFORCE(mlvalue_index >= 0 && static_cast<size_t>(mlvalue_index) < all_values_.size());
return all_values_[mlvalue_index];
}
// Get the index for the first entry of the given node.
int GetNodeOffset(NodeIndex index) const;
// Return nullptr if index map to an value that is an unused optional input/output
const MLValue* GetNodeInputOrOutputMLValue(int index) const;
MLValue* GetMutableNodeInputOrOutputMLValue(int index);
// TO DO: make it thread safe
// 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 GetOrCreateNodeOutputMLValue(int index,
const MLValueAllocationParameters& parameters,
MLValue*& p_mlvalue);
AllocatorPtr GetAllocator(const OrtAllocatorInfo& info);
// write the output values to the 'fetches' vector
Status GetOutputs(std::vector<MLValue>& fetches);
Status ReleaseMLValue(int mlvalue_idx);
const SessionState& GetSessionState() const {
return session_state_;
}
Status GeneratePatterns(MemoryPatternGroup* out) const;
bool HasPlan() const {
bool HasMemoryPatternPlanner() const {
return planner_ != nullptr;
}
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExecutionFrame);
void Init(const std::vector<int>& feed_mlvalue_idxs,
const std::vector<MLValue>& feeds,
const std::vector<int>& fetch_mlvalue_idxs,
std::vector<MLValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators);
AllocatorPtr GetAllocatorImpl(const OrtAllocatorInfo& info) const override;
Status ReleaseMLValueImpl(int mlvalue_idx) override;
Status CreateNodeOutputMLValueImpl(MLValue& mlvalue, int mlvalue_idx, const TensorShape* shape) override;
common::Status AllocateAsPerAllocationPlan(int mlvalue_index,
const MLValueAllocationParameters& parameters);
common::Status AllocateAsPerAllocationPlan(MLValue& mlvalue,
int mlvalue_index,
const TensorShape* shape);
Status AllocateMLValueTensorSelfOwnBufferHelper(int mlvalue_index,
Status AllocateMLValueTensorSelfOwnBufferHelper(MLValue& mlvalue,
int mlvalue_index,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape,
bool create_fence);
Status AllocateTensorWithPreAllocateBufferHelper(MLValue* p_mlvalue,
Status AllocateTensorWithPreAllocateBufferHelper(MLValue& mlvalue,
void* pBuffer,
MLDataType element_type,
const OrtAllocatorInfo& location,
const TensorShape& shape);
void TraceAllocate(int mlvalue_idx, size_t size);
void TraceFree(int mlvalue_idx);
const SequentialExecutionPlan::AllocPlanPerValue& GetAllocationPlan(int mlvalue_idx);
const AllocPlanPerValue& GetAllocationPlan(int mlvalue_idx);
Status status_;
const NodeIndexInfo& node_index_info_;
// All the intermediate values for the entire graph.
// Input and Output values are passed in by executors
std::vector<MLValue> all_values_;
// i-th kernel is still waiting for pending_counts_[i] inputs.
std::vector<int> pending_counts_; // not used currently
const SessionState& session_state_;
// map of index to custom allocator
std::unordered_map<int, IExecutor::CustomAllocator> custom_allocators_;
const SessionState& session_state_;
// If we already have cached memory pattern on these input shapes
// Use this mem pattern that create a big chunk for all the internal
// kernel's input/output tensors.
@ -159,10 +176,6 @@ class ExecutionFrame {
// use this planner_ to trace the memory allocation in current executor.
std::unique_ptr<MLValuePatternPlanner> planner_;
// Record the ml value indices for output values. we won't include those
// values' allocation in memory pattern, as they can't be shared.
const std::vector<int>& fetch_mlvalue_idxs_;
// Big chunks on different locations that will be used by mem_pattern.
std::map<OrtAllocatorInfo, BufferUniquePtr> buffers_;
};

52
onnxruntime/core/framework/node_index_info.cc
View file

@ -9,13 +9,55 @@
namespace onnxruntime {
// if we have a full GraphViewer, assume the min node index is 0
NodeIndexInfo::NodeIndexInfo(const GraphViewer& graph_viewer, const MLValueNameIdxMap& mlvalue_idx_map)
: min_node_index_{0}, max_mlvalue_idx_{mlvalue_idx_map.MaxIdx()} {
Init(graph_viewer.Nodes(), graph_viewer.MaxNodeIndex(), mlvalue_idx_map);
}
NodeIndexInfo::NodeIndexInfo(const GraphNodes& nodes, const MLValueNameIdxMap& mlvalue_idx_map)
: max_mlvalue_idx_{mlvalue_idx_map.MaxIdx()} {
Init(nodes, 0, mlvalue_idx_map);
}
NodeIndexInfo::NodeIndexInfo(const std::vector<const Node*>& nodes, const MLValueNameIdxMap& mlvalue_idx_map)
: max_mlvalue_idx_{mlvalue_idx_map.MaxIdx()} {
Init(ValidNodes<const std::vector<const Node*>>(nodes), 0, mlvalue_idx_map);
}
template <typename TValidNodes>
static void FindMinAndMaxNodeIndex(const TValidNodes& nodes, NodeIndex& min, NodeIndex& max) {
min = std::numeric_limits<NodeIndex>::max();
max = 0;
std::for_each(nodes.cbegin(), nodes.cend(), [&min, &max](const Node& node) {
auto idx = node.Index();
if (idx > max) max = idx;
if (idx >= 0 && idx < min) min = idx;
});
// match GraphViewer::MaxNodeIndex() which returns nodes_.size(), so is actually the max used value + 1.
// if we didn't do this, we'd have to add 1 when calling node_offsets_.resize, which would give max used value + 2
// if Init was called with max_node_index from GraphViewer::MaxNodeIndex(), which would create an extra invalid entry
// in node_offsets_ at the end.
max += 1;
}
template <typename TValidNodes>
void NodeIndexInfo::Init(const TValidNodes& nodes, NodeIndex max_node_index, const MLValueNameIdxMap& mlvalue_idx_map) {
if (nodes.empty()) {
// fairly stupid edge case to handle unit test for Constant. the Constant node becomes an initializer, leaving
// the graph with no nodes.
return;
}
std::size_t total_def_count{};
const bool include_missing_optional_defs = true;
bool include_missing_optional_defs = true;
if (max_node_index == 0) {
FindMinAndMaxNodeIndex(nodes, min_node_index_, max_node_index);
}
for (const auto& node : graph_viewer.Nodes()) {
for (const auto& node : nodes) {
node.ForEachDef(
[&](const onnxruntime::NodeArg& /*arg*/, bool /*is_input*/) {
++total_def_count;
@ -24,12 +66,12 @@ NodeIndexInfo::NodeIndexInfo(const GraphViewer& graph_viewer, const MLValueNameI
}
// init all to kInvalidEntry
node_offsets_.resize(graph_viewer.MaxNodeIndex(), kInvalidEntry);
node_offsets_.resize(GetNodeOffsetsIndex(max_node_index), kInvalidEntry);
node_values_.resize(total_def_count, kInvalidEntry);
int cur_idx = 0;
for (auto& node : graph_viewer.Nodes()) {
node_offsets_[node.Index()] = cur_idx;
for (auto& node : nodes) {
node_offsets_[GetNodeOffsetsIndex(node.Index())] = cur_idx;
node.ForEachDef(
[&](const onnxruntime::NodeArg& node_arg, bool /*is_input*/) {

25
onnxruntime/core/framework/node_index_info.h
View file

@ -9,22 +9,30 @@
#include "core/framework/ml_value.h"
namespace onnxruntime {
class GraphNodes;
class GraphViewer;
class MLValueNameIdxMap;
class Node;
class NodeIndexInfo final {
public:
// construct from a GraphViewer.
NodeIndexInfo(const GraphViewer& graph_viewer, const MLValueNameIdxMap& mlvalue_idx_map);
// construct from a subset of nodes. The min and max NodeIndex values will be calculated by iterating 'nodes'.
NodeIndexInfo(const GraphNodes& nodes, const MLValueNameIdxMap& mlvalue_idx_map);
NodeIndexInfo(const std::vector<const Node*>& nodes, const MLValueNameIdxMap& mlvalue_idx_map);
enum { kInvalidEntry = -1 };
// Index to the first argument of the given Node.
// The Node will have (num inputs + num implicit inputs + num outputs) entries, in that order, starting at the
// offset that is returned. Use the offset in calls to GetMLValueIndex.
// Returns kInvalidEntry if the Node with the given node_index did not exist when the NodeIndexInfo was created.
int GetNodeOffset(onnxruntime::NodeIndex node_index) const {
ORT_ENFORCE(node_index < node_offsets_.size());
return node_offsets_[node_index];
int GetNodeOffset(NodeIndex node_index) const {
auto node_offsets_index = GetNodeOffsetsIndex(node_index);
ORT_ENFORCE(node_offsets_index < node_offsets_.size());
return node_offsets_[node_offsets_index];
}
// Get the mlvalue index value.
@ -39,12 +47,19 @@ class NodeIndexInfo final {
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(NodeIndexInfo);
template <typename TValidNodes>
void Init(const TValidNodes& nodes, NodeIndex max_node_index, const MLValueNameIdxMap& mlvalue_idx_map);
// This vector contains the indices from the MLValueNameIdxMap in the SessionState for each Node's input/outputs.
// Order is node inputs, implicit inputs, outputs.
std::vector<int> node_values_;
// The entry at node_offset_[Node::Index()] contains the index in node_values_ where the information for the Node
// begins.
// the minimum NodeIndex. we use this to minimize the size of node_offsets_.
NodeIndex min_node_index_ = 0;
// The entry at node_offsets_[GetNodeOffsetsIndex(Node::Index())] contains the index in node_values_
// where the information for the Node begins.
size_t GetNodeOffsetsIndex(NodeIndex node_index) const { return node_index - min_node_index_; }
std::vector<int> node_offsets_;
const int max_mlvalue_idx_;

15
onnxruntime/core/framework/op_kernel.cc
View file

@ -9,7 +9,7 @@
using namespace ::onnxruntime::common;
namespace onnxruntime {
OpKernelContext::OpKernelContext(ExecutionFrame* frame,
OpKernelContext::OpKernelContext(IExecutionFrame* frame,
const OpKernel* kernel,
const logging::Logger& logger)
: execution_frame_(frame),
@ -27,16 +27,12 @@ Tensor* OpKernelContext::Output(int index, const TensorShape& shape) {
if (index < 0 || index >= OutputCount())
return nullptr;
// In this case, it's assumed that the tensor hasn't been allocated yet,
// so that it's calling ExecutionFrame to create a tensor in the given position with given shape.
MLValueAllocationParameters parameters{&shape};
//: Though we don't need to give 'ret' an initial value, GCC would generate a warning if we don't do that
//"error: 'ret' may be used uninitialized in this function"
//This warning only exists in Release build.
//I believe it's a false alarm.
MLValue* p_ml_value = nullptr;
Status status = execution_frame_->GetOrCreateNodeOutputMLValue(GetOutputArgIndex(index), parameters, p_ml_value);
Status status = execution_frame_->GetOrCreateNodeOutputMLValue(GetOutputArgIndex(index), &shape, p_ml_value);
ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
return p_ml_value ? p_ml_value->GetMutable<Tensor>() : nullptr;
}
@ -97,8 +93,7 @@ Fence_t OpKernelContext::OutputFence(int index) const {
Status OpKernelContext::GetOrCreateOutputMLValue(int index, MLValue*& p_value) {
auto output_arg_index = GetOutputArgIndex(index);
MLValueAllocationParameters parameters;
ORT_ENFORCE(execution_frame_->GetOrCreateNodeOutputMLValue(output_arg_index, parameters, p_value).IsOK());
ORT_ENFORCE(execution_frame_->GetOrCreateNodeOutputMLValue(output_arg_index, nullptr, p_value).IsOK());
return Status::OK();
}
@ -118,10 +113,6 @@ onnxruntime::NodeIndex OpKernelContext::GetNodeIndex() const {
return kernel_->Node().Index();
}
const SessionState& OpKernelContext::GetSessionState() const {
return execution_frame_->GetSessionState();
}
const MLValue* OpKernelContext::GetInputMLValue(int index) const {
if (index < 0 || index >= InputCount())
return nullptr;

8
onnxruntime/core/framework/op_kernel_context_internal.h
View file

@ -11,21 +11,24 @@
namespace onnxruntime {
class SessionState;
class ExecutionFrame;
class OpKernelContextInternal : public OpKernelContext {
public:
explicit OpKernelContextInternal(ExecutionFrame& frame,
explicit OpKernelContextInternal(const SessionState& session_state,
IExecutionFrame& frame,
const OpKernel& kernel,
const logging::Logger& logger,
const std::vector<NodeArg*>& implicit_inputs,
const bool& terminate_flag)
: OpKernelContext(&frame, &kernel, logger),
session_state_{session_state},
implicit_inputs_{implicit_inputs},
terminate_flag_{terminate_flag} {
}
const SessionState* SubgraphSessionState(const std::string& attribute_name) {
return GetSessionState().GetSubgraphSessionState(GetNodeIndex(), attribute_name);
return session_state_.GetSubgraphSessionState(GetNodeIndex(), attribute_name);
}
const MLValue* GetInputMLValue(int index) const {
@ -51,6 +54,7 @@ class OpKernelContextInternal : public OpKernelContext {
const bool& GetTerminateFlag() const noexcept { return terminate_flag_; }
private:
const SessionState& session_state_;
const std::vector<NodeArg*>& implicit_inputs_;
const bool& terminate_flag_;
};

4
onnxruntime/core/framework/parallel_executor.cc
View file

@ -67,7 +67,7 @@ Status ParallelExecutor::Execute(const SessionState& session_state,
ORT_RETURN_IF_ERROR(root_frame_->GetOutputs(fetches));
VLOGS(logger, 1) << "Done execution.";
if (root_frame_->HasPlan()) {
if (root_frame_->HasMemoryPatternPlanner()) {
std::vector<TensorShape> input_shapes;
bool all_tensors = true;
for (const auto& feed : feeds) {
@ -131,7 +131,7 @@ void ParallelExecutor::RunNodeAsyncInternal(size_t p_node_index,
graph_viewer->GetNode(node_index)->Name());
}
OpKernelContextInternal op_kernel_context(*root_frame_, *p_op_kernel, logger,
OpKernelContextInternal op_kernel_context(session_state, *root_frame_, *p_op_kernel, logger,
p_op_kernel->Node().ImplicitInputDefs(),
terminate_flag_);

35
onnxruntime/core/framework/sequential_execution_plan.h
View file

@ -17,6 +17,24 @@ using MLValueIndex = int;
using MLValueName = std::string;
class SessionState;
// AllocPlanPerValue: (a simplified form of AllocationPlanPerValue above)
// Captures information required to allocate/reuse buffer for a ml-value
struct AllocPlanPerValue {
AllocKind alloc_kind{AllocKind::kAllocate};
MLDataType value_type{nullptr};
OrtAllocatorInfo location;
// reused_buffer is valid only if alloc_kind == kReuse. It indicates
// which MLValue's buffer must be reused for this MLValue.
MLValueIndex reused_buffer{0};
// if the value is used in async kernel, a fence object would be created
// note the fence object would be shared between MLValues reusing the same buffer
bool create_fence_if_async{false};
public:
AllocPlanPerValue() : location(CPU, OrtArenaAllocator) {}
};
// SequentialExecutionPlan: This is the data that is produced by a static
// planner for a sequential execution, to be used by a SequentialExecutor.
struct SequentialExecutionPlan {
@ -24,23 +42,6 @@ struct SequentialExecutionPlan {
// ExecutionFrame::GetOrCreateTensor() should use the following information
// to decide whether to allocate a new buffer or reuse an existing buffer
// AllocPlanPerValue: (a simplified form of AllocationPlanPerValue above)
// Captures information required to allocate/reuse buffer for a ml-value
struct AllocPlanPerValue {
AllocKind alloc_kind{AllocKind::kAllocate};
MLDataType value_type{nullptr};
OrtAllocatorInfo location;
// reused_buffer is valid only if alloc_kind == kReuse. It indicates
// which MLValue's buffer must be reused for this MLValue.
MLValueIndex reused_buffer{0};
// if the value is used in async kernel, a fence object would be created
// note the fence object would be shared between MLValues reusing the same buffer
bool create_fence_if_async{false};
public:
AllocPlanPerValue() : location(CPU, OrtArenaAllocator) {}
};
// The following vector is indexed by MLValueIndex
std::vector<AllocPlanPerValue> allocation_plan;

6
onnxruntime/core/framework/sequential_executor.cc
View file

@ -63,8 +63,8 @@ Status SequentialExecutor::Execute(const SessionState& session_state,
// construct OpKernelContext
// TODO: log kernel inputs?
OpKernelContextInternal op_kernel_context(frame, *p_op_kernel, logger, p_op_kernel->Node().ImplicitInputDefs(),
terminate_flag_);
OpKernelContextInternal op_kernel_context(session_state, frame, *p_op_kernel, logger,
p_op_kernel->Node().ImplicitInputDefs(), terminate_flag_);
// TODO: log kernel outputs?
if (f_profiler_enabled) {
sync_time_begin = session_state.Profiler().StartTime();
@ -162,7 +162,7 @@ Status SequentialExecutor::Execute(const SessionState& session_state,
ORT_RETURN_IF_ERROR(frame.GetOutputs(fetches));
VLOGS(logger, 1) << "Done with execution.";
if (frame.HasPlan()) {
if (frame.HasMemoryPatternPlanner()) {
std::vector<TensorShape> input_shapes;
bool all_tensors = true;
for (const auto& feed : feeds) {

16
onnxruntime/test/framework/execution_frame_test.cc
View file

@ -78,7 +78,8 @@ TEST(ExecutionFrameTest, TensorAllocationTest) {
EXPECT_EQ(start_index, 0);
TensorShape shape(std::vector<int64_t>{2, 3});
status = frame.AllocateMLValueTensorSelfOwnBuffer(start_index, DataTypeImpl::GetType<float>(),
MLValue& mlvalue0 = *frame.GetMutableNodeInputOrOutputMLValue(start_index);
status = frame.AllocateMLValueTensorSelfOwnBuffer(mlvalue0, start_index, DataTypeImpl::GetType<float>(),
execution_providers.Get(xp_typ)->GetAllocator(0, OrtMemTypeDefault)->Info(), shape);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
@ -90,7 +91,8 @@ TEST(ExecutionFrameTest, TensorAllocationTest) {
//test share memory from tensor
TensorShape shape2(std::vector<int64_t>{3, 2});
status = frame.AllocateMLValueTensorPreAllocateBuffer(start_index + 1,
MLValue& mlvalue1 = *frame.GetMutableNodeInputOrOutputMLValue(start_index + 1);
status = frame.AllocateMLValueTensorPreAllocateBuffer(mlvalue1,
start_index,
DataTypeImpl::GetType<float>(),
p_tensor->Location(),
@ -226,19 +228,23 @@ TEST(ExecutionFrameTest, MemPatternTest) {
vector<MLValue> outputs;
ExecutionFrame frame({x1_idx, x2_idx, x3_idx}, {v1, v2, v3}, {t3_idx}, outputs, {}, state);
status = frame.AllocateMLValueTensorSelfOwnBuffer(3,
MLValue& mlvalue3 = *frame.GetMutableNodeInputOrOutputMLValue(3);
MLValue& mlvalue4 = *frame.GetMutableNodeInputOrOutputMLValue(4);
MLValue& mlvalue5 = *frame.GetMutableNodeInputOrOutputMLValue(5);
status = frame.AllocateMLValueTensorSelfOwnBuffer(mlvalue3, 3,
DataTypeImpl::GetType<float>(),
cpu_allocator->Info(),
TensorShape(std::vector<int64_t>{2, 2}));
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
status = frame.AllocateMLValueTensorSelfOwnBuffer(4,
status = frame.AllocateMLValueTensorSelfOwnBuffer(mlvalue4, 4,
DataTypeImpl::GetType<float>(),
cpu_allocator->Info(),
TensorShape(std::vector<int64_t>{2, 3}));
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
status = frame.AllocateMLValueTensorSelfOwnBuffer(5,
status = frame.AllocateMLValueTensorSelfOwnBuffer(mlvalue5, 5,
DataTypeImpl::GetType<float>(),
cpu_allocator->Info(),
TensorShape(std::vector<int64_t>{2, 3}));