From 971058fc38b24c8582da032f1421840c8506046f Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Tue, 19 Mar 2019 06:55:59 +1000 Subject: [PATCH] Avoid copy of pre-existing value to subgraph output (#637) * 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 --- .../onnxruntime/core/framework/alloc_kind.h | 3 +- .../core/framework/allocation_planner.cc | 37 +++++++++++++++---- .../core/framework/allocation_planner.h | 9 +++-- onnxruntime/core/framework/execution_frame.cc | 6 +++ .../framework/session_state_initializer.cc | 20 +++++----- .../framework/session_state_initializer.h | 4 +- .../core/providers/cpu/controlflow/loop.cc | 6 +-- onnxruntime/core/session/inference_session.cc | 4 +- .../test/framework/allocation_planner_test.cc | 2 +- .../test/framework/execution_frame_test.cc | 8 ++-- onnxruntime/test/providers/memcpy_test.cc | 4 +- 11 files changed, 69 insertions(+), 34 deletions(-) diff --git a/include/onnxruntime/core/framework/alloc_kind.h b/include/onnxruntime/core/framework/alloc_kind.h index 7ccb01226a..a749e6b26c 100644 --- a/include/onnxruntime/core/framework/alloc_kind.h +++ b/include/onnxruntime/core/framework/alloc_kind.h @@ -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); diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 84583fe838..80f3813d36 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -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& 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& 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& 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(); - 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(); diff --git a/onnxruntime/core/framework/allocation_planner.h b/onnxruntime/core/framework/allocation_planner.h index e220606df7..9115c2f0bc 100644 --- a/onnxruntime/core/framework/allocation_planner.h +++ b/onnxruntime/core/framework/allocation_planner.h @@ -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& 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& outer_scope_node_args, const ExecutionProviders& providers, const KernelRegistryManager& kernel_registry, const MLValueNameIdxMap& mlvalue_name_idx_map, std::unique_ptr& 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); } }; diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index e7ae8805c2..30b828b9e1 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -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::type>(alloc_kind); diff --git a/onnxruntime/core/framework/session_state_initializer.cc b/onnxruntime/core/framework/session_state_initializer.cc index ce72678484..a4020d3399 100644 --- a/onnxruntime/core/framework/session_state_initializer.cc +++ b/onnxruntime/core/framework/session_state_initializer.cc @@ -59,7 +59,8 @@ SessionStateInitializer::SessionStateInitializer(const std::basic_string& outer_scope_node_args, +common::Status SessionStateInitializer::CreatePlan(const Node* parent_node, + const std::vector& outer_scope_node_args, bool enable_sequential_execution) { auto graph_viewer = std::make_unique(graph_); @@ -83,7 +84,7 @@ common::Status SessionStateInitializer::CreatePlan(const std::vector& // 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& // 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 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 diff --git a/onnxruntime/core/framework/session_state_initializer.h b/onnxruntime/core/framework/session_state_initializer.h index 8f575ff0f2..47516c3b75 100644 --- a/onnxruntime/core/framework/session_state_initializer.h +++ b/onnxruntime/core/framework/session_state_initializer.h @@ -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& outer_scope_node_args, + common::Status CreatePlan(const Node* parent_node, + const std::vector& outer_scope_node_args, bool enable_sequential_execution); // initialize tensors, and save. save kernels and input/output node mappings diff --git a/onnxruntime/core/providers/cpu/controlflow/loop.cc b/onnxruntime/core/providers/cpu/controlflow/loop.cc index fbf72956b9..1361ab4d21 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -103,7 +103,7 @@ class LoopImpl { private: void CreateInitialFeeds(std::vector& feeds); - void UpdateFeeds(const std::vector& last_outputs, std::vector& next_inputs); + void SaveOutputsAndUpdateFeeds(const std::vector& last_outputs, std::vector& next_inputs); // create the single Loop output from a collection of per-iteration outputs Status ConcatenateLoopOutput(std::vector& per_iteration_output, int output_index); @@ -266,7 +266,7 @@ void LoopImpl::CreateInitialFeeds(std::vector& feeds) { } } -void LoopImpl::UpdateFeeds(const std::vector& last_outputs, std::vector& next_inputs) { +void LoopImpl::SaveOutputsAndUpdateFeeds(const std::vector& last_outputs, std::vector& 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()->MutableData()) { if (iter_num_value != 0) { - UpdateFeeds(fetches, feeds); + SaveOutputsAndUpdateFeeds(fetches, feeds); fetches.clear(); } diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index f8b7b3d27f..e6eaa2455d 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -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 diff --git a/onnxruntime/test/framework/allocation_planner_test.cc b/onnxruntime/test/framework/allocation_planner_test.cc index 3cc7b0a7a3..6c1910bdcb 100644 --- a/onnxruntime/test/framework/allocation_planner_test.cc +++ b/onnxruntime/test/framework/allocation_planner_test.cc @@ -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(); diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index cb36139464..752edb35ef 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -64,8 +64,8 @@ TEST(ExecutionFrameTest, TensorAllocationTest) { std::unique_ptr 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(6, 1.0f), &v3); std::unique_ptr p_seq_exec_plan = std::make_unique(); - 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)); diff --git a/onnxruntime/test/providers/memcpy_test.cc b/onnxruntime/test/providers/memcpy_test.cc index f8a44f1458..938071c1e0 100644 --- a/onnxruntime/test/providers/memcpy_test.cc +++ b/onnxruntime/test/providers/memcpy_test.cc @@ -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 \ No newline at end of file +} // namespace onnxruntime