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