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

Avoid copy of pre-existing value to subgraph output (#637)

Browse source 
* Add AllocKind::kShare to allow copying the MLValue for a pre-existing value to a graph output when an Identity node is involved. Ideally we can make this handling for an Identity node more general purpose, however the current logic to free an MLValue during execution doesn't take into account a re-use point also needing a free. Due to that, limit the scope and start with a somewhat ugly hardcoded approach.

Migrate some changes from PR497

The existing Loop unit tests exercise the new code. Also manually stepped through the problematic model to verify the unnecessary copy was avoided.

* Fix build error

* Fix missing switch case in debug output of allocation plan

* Limit optimization to Loop
This commit is contained in:
Scott McKay 2019-03-19 06:55:59 +10:00 committed by GitHub
parent 14d9a2bdc7
commit 971058fc38
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 69 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
include/onnxruntime/core/framework/alloc_kind.h
View file

@ -26,7 +26,8 @@ enum class AllocKind {
kReuse = 1,
kPreExisting = 2,
kAllocateStatically = 3,
kAllocateOutput = 4
kAllocateOutput = 4,
kShare = 5
};
std::ostream& operator<<(std::ostream& out, AllocKind alloc_kind);

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

@ -36,6 +36,9 @@ std::ostream& operator<<(std::ostream& out, AllocKind alloc_kind) {
case AllocKind::kAllocateOutput:
out << "AllocateOutput";
break;
case AllocKind::kShare:
out << "Share";
break;
}
return out;
}
@ -97,7 +100,8 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP
class PlannerImpl {
public:
PlannerImpl(const onnxruntime::GraphViewer& graph_viewer,
PlannerImpl(const Node* parent_node,
const onnxruntime::GraphViewer& graph_viewer,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
@ -106,6 +110,7 @@ class PlannerImpl {
SequentialExecutionPlan& plan)
: context_{context},
plan_{plan},
parent_node_{parent_node},
graph_viewer_{graph_viewer},
outer_scope_node_args_{outer_scope_node_args},
execution_providers_{providers},
@ -119,6 +124,7 @@ class PlannerImpl {
const ISequentialPlannerContext& context_;
SequentialExecutionPlan& plan_;
const Node* parent_node_;
const onnxruntime::GraphViewer& graph_viewer_;
const std::vector<const NodeArg*>& outer_scope_node_args_;
const ExecutionProviders& execution_providers_;
@ -177,7 +183,8 @@ class PlannerImpl {
info.p_def_site = p_def_site;
}
void Reuse(MLValueIndex reused, MLValueIndex reused_for) {
// Reuse/Alias/Share between two MLValue indexes
void Reuse(MLValueIndex reused, MLValueIndex reused_for, AllocKind alloc_kind) {
ORT_ENFORCE(reused != reused_for);
// find original buffer underlying ml-value we want to reuse:
MLValueIndex original = Buffer(reused);
@ -188,7 +195,7 @@ class PlannerImpl {
// update allocation plan (for use at execution-time)
auto& symplan = AllocPlan(reused_for);
symplan.alloc_kind = AllocKind::kReuse;
symplan.alloc_kind = alloc_kind;
symplan.reused_buffer = original;
}
@ -498,17 +505,30 @@ class PlannerImpl {
AllocPlan(current).value_type = utils::GetMLDataType(*node_output);
MLValueIndex reused;
if (std::find(graph_outputs.begin(), graph_outputs.end(), node_output) != graph_outputs.end()) {
// node_output is graph's output, so we can't reuse intermedia buffer
// node_output is graph's output, so we can't reuse intermediate buffer
AllocPlan(current).alloc_kind = AllocKind::kAllocateOutput;
// hacky perf optimization to not copy a pre-existing value to an output if this is a Loop subgraph.
// ideally this is temporary, and a future ONNX change to allow empty variadic inputs means we don't
// have converted models that unnecessarily add loop state variables. if the value is just being
// passed through an implicit input should be used instead.
if (parent_node_ && pnode->OpType() == "Identity" && parent_node_->OpType() == "Loop") {
const auto& input_name = pnode->InputDefs()[0]->Name();
const auto input_index = Index(input_name);
const auto& alloc_plan = AllocPlan(input_index);
if (alloc_plan.alloc_kind == AllocKind::kPreExisting) {
Reuse(input_index, current, AllocKind::kShare);
}
}
} else if (IsNonTensor(*node_output)) {
// we do not try sharing-optimization for non-tensors
AllocPlan(current).alloc_kind = AllocKind::kAllocate;
} else if (FindReusableInput(*pnode, output_arg_num, &reused)) {
// Reuse one of this node's input buffers as the output buffer (for in-place update)
Reuse(reused, current);
Reuse(reused, current, AllocKind::kReuse);
} else if (!context_.EnableParallelExecution() && FindReusableTensor(*node_output, &reused)) {
// Reuse an available (dead) buffer for this output, this is only for sequential execution.
Reuse(reused, current);
Reuse(reused, current, AllocKind::kReuse);
} else {
// otherwise: allocate a new buffer for this output
AllocPlan(current).alloc_kind = AllocKind::kAllocate;
@ -610,7 +630,8 @@ Status PlannerImpl::CreatePlan() {
return Status::OK();
}
Status SequentialPlanner::CreatePlan(const onnxruntime::GraphViewer& graph_viewer,
Status SequentialPlanner::CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph_viewer,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
@ -620,7 +641,7 @@ Status SequentialPlanner::CreatePlan(const onnxruntime::GraphViewer& graph_viewe
// allocate/reset here so we know it's clean
plan = std::make_unique<SequentialExecutionPlan>();
PlannerImpl planner(graph_viewer, outer_scope_node_args,
PlannerImpl planner(parent_node, graph_viewer, outer_scope_node_args,
providers, kernel_registry, mlvalue_name_idx_map, context, *plan);
return planner.CreatePlan();

9
onnxruntime/core/framework/allocation_planner.h
View file

@ -50,7 +50,8 @@ class SequentialPlannerContext : public ISequentialPlannerContext {
class SequentialPlanner {
public:
// This API allows user to provide a custom planner context.
static Status CreatePlan(const onnxruntime::GraphViewer& graph,
static Status CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
@ -60,14 +61,16 @@ class SequentialPlanner {
// This uses a standard planner context and is meant to be the primary API for creating a plan
// as the context is primarily used in test scenarios.
static Status CreatePlan(const onnxruntime::GraphViewer& graph,
static Status CreatePlan(const Node* parent_node,
const onnxruntime::GraphViewer& graph,
const std::vector<const NodeArg*>& outer_scope_node_args,
const ExecutionProviders& providers,
const KernelRegistryManager& kernel_registry,
const MLValueNameIdxMap& mlvalue_name_idx_map,
std::unique_ptr<SequentialExecutionPlan>& plan) {
SequentialPlannerContext context;
return CreatePlan(graph, outer_scope_node_args, providers, kernel_registry, mlvalue_name_idx_map, context, plan);
return CreatePlan(parent_node, graph, outer_scope_node_args, providers, kernel_registry, mlvalue_name_idx_map,
context, plan);
}
};

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

@ -395,6 +395,12 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(MLValue& mlvalue, int mlvalue
per_alloc_plan.create_fence_if_async));
break;
}
case AllocKind::kShare: {
int reuse_mlvalue_index = per_alloc_plan.reused_buffer;
// copy at the MLValue level so the shared_ptr for the data is shared between the two MLValue instances
mlvalue = GetMutableMLValue(reuse_mlvalue_index);
break;
}
default: {
std::ostringstream ostr;
ostr << "Invalid allocation kind: " << static_cast<std::underlying_type<AllocKind>::type>(alloc_kind);

20
onnxruntime/core/framework/session_state_initializer.cc
View file

@ -59,7 +59,8 @@ SessionStateInitializer::SessionStateInitializer(const std::basic_string<PATH_CH
kernel_registry_manager_{kernel_registry_manager},
logger_{session_state.Logger()} {}
common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>& outer_scope_node_args,
common::Status SessionStateInitializer::CreatePlan(const Node* parent_node,
const std::vector<NodeArg*>& outer_scope_node_args,
bool enable_sequential_execution) {
auto graph_viewer = std::make_unique<onnxruntime::GraphViewer>(graph_);
@ -83,7 +84,7 @@ common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>&
// CreatePlan will create a new SequentialExecutionPlan instance that we will
// save into the session state.
ORT_RETURN_IF_ERROR(
SequentialPlanner::CreatePlan(*graph_viewer, valid_outer_scope_node_args, execution_providers_,
SequentialPlanner::CreatePlan(parent_node, *graph_viewer, valid_outer_scope_node_args, execution_providers_,
kernel_registry_manager_, mlvalue_name_idx_map, exec_plan));
session_state_.SetExecutionPlan(std::move(exec_plan));
@ -91,7 +92,7 @@ common::Status SessionStateInitializer::CreatePlan(const std::vector<NodeArg*>&
// Parallel execution still uses same allocation plan, but has limitation of memory buffer reuse.
SequentialPlannerContext context(true /* enable parallel execution */);
ORT_RETURN_IF_ERROR(
SequentialPlanner::CreatePlan(*graph_viewer, valid_outer_scope_node_args, execution_providers_,
SequentialPlanner::CreatePlan(parent_node, *graph_viewer, valid_outer_scope_node_args, execution_providers_,
kernel_registry_manager_, mlvalue_name_idx_map, context, exec_plan));
session_state_.SetExecutionPlan(std::move(exec_plan));
@ -112,12 +113,13 @@ common::Status SessionStateInitializer::InitializeAndSave(const std::vector<Node
// lambda to save initialized tensors into SessionState directly
const Env& env = Env::Default();
ORT_RETURN_IF_ERROR(
SaveInitializedTensors(env, graph_loc_, graph_, exec_plan, execution_providers_, mlvalue_name_idx_map,
session_state_.GetMutableWeightsBuffers(),
[this](int idx, const onnxruntime::MLValue& value, const OrtCallback& d) -> Status {
return session_state_.AddInitializedTensor(idx, value, &d);
},
logger_));
SaveInitializedTensors(
env, graph_loc_, graph_, exec_plan, execution_providers_, mlvalue_name_idx_map,
session_state_.GetMutableWeightsBuffers(),
[this](int idx, const onnxruntime::MLValue& value, const OrtCallback& d) -> Status {
return session_state_.AddInitializedTensor(idx, value, &d);
},
logger_));
// remove weights from the graph now to save memory but in many cases it won't save memory, if the tensor was
// preallocated with the some other tensors in a single 'allocate' call, which is very common.
// TODO: make it better

4
onnxruntime/core/framework/session_state_initializer.h
View file

@ -14,6 +14,7 @@ class Graph;
class GraphTransformerManager;
class InsertCastTransformer;
class KernelRegistryManager;
class Node;
class NodeArg;
class SessionState;
@ -33,7 +34,8 @@ class SessionStateInitializer {
KernelRegistryManager& kernel_registry_manager);
// First perform any transformations and create the execution plan
common::Status CreatePlan(const std::vector<NodeArg*>& outer_scope_node_args,
common::Status CreatePlan(const Node* parent_node,
const std::vector<NodeArg*>& outer_scope_node_args,
bool enable_sequential_execution);
// initialize tensors, and save. save kernels and input/output node mappings

6
onnxruntime/core/providers/cpu/controlflow/loop.cc
View file

@ -103,7 +103,7 @@ class LoopImpl {
private:
void CreateInitialFeeds(std::vector<MLValue>& feeds);
void UpdateFeeds(const std::vector<MLValue>& last_outputs, std::vector<MLValue>& next_inputs);
void SaveOutputsAndUpdateFeeds(const std::vector<MLValue>& last_outputs, std::vector<MLValue>& next_inputs);
// create the single Loop output from a collection of per-iteration outputs
Status ConcatenateLoopOutput(std::vector<MLValue>& per_iteration_output, int output_index);
@ -266,7 +266,7 @@ void LoopImpl::CreateInitialFeeds(std::vector<MLValue>& feeds) {
}
}
void LoopImpl::UpdateFeeds(const std::vector<MLValue>& last_outputs, std::vector<MLValue>& next_inputs) {
void LoopImpl::SaveOutputsAndUpdateFeeds(const std::vector<MLValue>& last_outputs, std::vector<MLValue>& next_inputs) {
// last_output: cond, loop vars..., loop output...
// next_input: iter_num, cond, loop_vars. iter_num is re-used
@ -331,7 +331,7 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca
while (iter_num_value < max_trip_count_ && *condition_mlvalue_.GetMutable<Tensor>()->MutableData<bool>()) {
if (iter_num_value != 0) {
UpdateFeeds(fetches, feeds);
SaveOutputsAndUpdateFeeds(fetches, feeds);
fetches.clear();
}

4
onnxruntime/core/session/inference_session.cc
View file

@ -454,7 +454,7 @@ class InferenceSession::Impl {
SessionStateInitializer initializer{model_location_, subgraph, *subgraph_session_state, execution_providers_,
kernel_registry_manager_};
ORT_RETURN_IF_ERROR(initializer.CreatePlan(node.ImplicitInputDefs(),
ORT_RETURN_IF_ERROR(initializer.CreatePlan(&node, node.ImplicitInputDefs(),
session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(initializer.InitializeAndSave(&node.ImplicitInputDefs()));
@ -522,7 +522,7 @@ class InferenceSession::Impl {
// now that all the transforms are done, call Resolve on the main graph. this will recurse into the subgraphs.
ORT_RETURN_IF_ERROR(graph.Resolve());
ORT_RETURN_IF_ERROR(session_initializer.CreatePlan({}, session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(session_initializer.CreatePlan(nullptr, {}, session_options_.enable_sequential_execution));
ORT_RETURN_IF_ERROR(session_initializer.InitializeAndSave(nullptr));
// handle any subgraphs

2
onnxruntime/test/framework/allocation_planner_test.cc
View file

@ -238,7 +238,7 @@ class PlannerTest : public ::testing::Test {
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
SequentialPlannerTestContext test_context(&shape_map_);
status = SequentialPlanner::CreatePlan(GraphViewer(graph_), outer_scope_node_args, execution_providers,
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph_), outer_scope_node_args, execution_providers,
kernel_registry_manager, mlvalue_name_idx_map, test_context, plan_);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();

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

@ -64,8 +64,8 @@ TEST(ExecutionFrameTest, TensorAllocationTest) {
std::unique_ptr<SequentialExecutionPlan> p_seq_exec_plan;
// TODO below line is for testing only. In production use SequentialPlanner::CreatePlan()
status = SequentialPlanner::CreatePlan(GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
p_seq_exec_plan);
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph), {}, execution_providers, kernel_registry_manager,
mlvalue_name_idx_map, p_seq_exec_plan);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
state.SetExecutionPlan(std::move(p_seq_exec_plan));
@ -214,8 +214,8 @@ TEST(ExecutionFrameTest, MemPatternTest) {
std::vector<float>(6, 1.0f), &v3);
std::unique_ptr<SequentialExecutionPlan> p_seq_exec_plan = std::make_unique<SequentialExecutionPlan>();
status = SequentialPlanner::CreatePlan(GraphViewer(graph), {}, execution_providers, kernel_registry_manager, mlvalue_name_idx_map,
p_seq_exec_plan);
status = SequentialPlanner::CreatePlan(nullptr, GraphViewer(graph), {}, execution_providers, kernel_registry_manager,
mlvalue_name_idx_map, p_seq_exec_plan);
EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
state.SetExecutionPlan(std::move(p_seq_exec_plan));

4
onnxruntime/test/providers/memcpy_test.cc
View file

@ -87,7 +87,7 @@ TEST(MemcpyTest, copy1) {
PutAllNodesOnOneProvider(model.MainGraph(), onnxruntime::kCpuExecutionProvider);
SessionStateInitializer session_initializer{ORT_TSTR(""), model.MainGraph(), s, execution_providers,
kernel_registry_manager};
st = session_initializer.CreatePlan({}, true);
st = session_initializer.CreatePlan(nullptr, {}, true);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();
st = session_initializer.InitializeAndSave(nullptr);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();
@ -104,4 +104,4 @@ TEST(MemcpyTest, copy1) {
st = utils::CopyOneInputAcrossDevices(s, "X", input, output);
ASSERT_TRUE(st.IsOK()) << st.ErrorMessage();
}
} // namespace onnxruntime
} // namespace onnxruntime