diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 84583fe83819e..457cfd99c1502 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -97,7 +97,8 @@ std::ostream& operator<<(std::ostream& out, std::pair& outer_scope_node_args, const ExecutionProviders& providers, const KernelRegistryManager& kernel_registry, @@ -106,6 +107,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 +121,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_; @@ -193,8 +196,7 @@ class PlannerImpl { } // Find if there exists some input tensor that we can use in-place for output_arg - bool FindReusableInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* reusable_input) { - auto p_output_arg = node.OutputDefs()[output_arg_num]; + bool FindAliasInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* alias_input) { const KernelCreateInfo* ci; Status st = kernel_registry_.SearchKernelRegistry(node, &ci); if (!st.IsOK() || ci == nullptr || ci->kernel_def == nullptr) { @@ -209,14 +211,26 @@ class PlannerImpl { if ((0 <= pair.first) && (static_cast(pair.first) < input_args.size())) { auto p_input_arg = input_args[pair.first]; if (p_input_arg->Exists()) { - *reusable_input = Index(p_input_arg->Name()); + *alias_input = Index(p_input_arg->Name()); return true; } } } } + return false; + } + + bool FindInplaceInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* inplace_input) { + auto p_output_arg = node.OutputDefs()[output_arg_num]; + const KernelCreateInfo* ci; + Status st = kernel_registry_.SearchKernelRegistry(node, &ci); + if (!st.IsOK() || ci == nullptr || ci->kernel_def == nullptr) { + return false; + } + const std::vector>& inplace_map = ci->kernel_def->MayInplace(); + auto& input_args = node.InputDefs(); for (auto pair : inplace_map) { if (pair.second == output_arg_num) { if ((0 <= pair.first) && (static_cast(pair.first) < input_args.size())) { @@ -227,7 +241,7 @@ class PlannerImpl { if (1 == UseCount(original)) { if (SameSize(*p_input_arg, *p_output_arg)) { // we can reuse this input since it is its last use and permitted for in-place update - *reusable_input = input_arg_index; // or original; both should be okay + *inplace_input = input_arg_index; // or original; both should be okay return true; } } @@ -235,6 +249,37 @@ class PlannerImpl { } } } + + return false; + } + + bool FindReusableGraphInput(const onnxruntime::Node& node, int output_arg_num, MLValueIndex* reusable_input) { + // skip if this is main graph + if (parent_node_) { + // only check a single layer. ie. an Identity/Dropout node connecting graph's input and output + auto& graph_inputs = graph_viewer_.GetInputs(); + auto& graph_outputs = graph_viewer_.GetOutputs(); + + if (FindAliasInput(node, output_arg_num, reusable_input)) { + auto& arg = ml_value_info_.at(Buffer(*reusable_input)).p_def_site; + auto it = std::find(graph_inputs.cbegin(), graph_inputs.cend(), arg); + if (it != graph_inputs.end()) { + auto graph_input_index = std::distance(graph_inputs.cbegin(), it); + ORT_ENFORCE(graph_input_index >= 0 && static_cast(graph_input_index) < graph_inputs.size()); + auto graph_output_index = std::distance(graph_outputs.cbegin(), + std::find(graph_outputs.cbegin(), graph_outputs.cend(), node.OutputDefs()[output_arg_num])); + ORT_ENFORCE(graph_output_index >= 0 && static_cast(graph_output_index) < graph_outputs.size()); + + if (parent_node_->OpType() == "Loop") { + // the matching carried parameters + return graph_input_index >= 2; + } else { + // TODO: other operators like Scan + } + } + } + } + return false; } @@ -497,13 +542,15 @@ class PlannerImpl { auto current = Index(node_output->Name()); AllocPlan(current).value_type = utils::GetMLDataType(*node_output); MLValueIndex reused; - if (std::find(graph_outputs.begin(), graph_outputs.end(), node_output) != graph_outputs.end()) { + if (std::find(graph_outputs.begin(), graph_outputs.end(), node_output) != graph_outputs.end() && + !FindReusableGraphInput(*pnode, output_arg_num, &reused)) { // node_output is graph's output, so we can't reuse intermedia buffer AllocPlan(current).alloc_kind = AllocKind::kAllocateOutput; } 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)) { + } else if (FindAliasInput(*pnode, output_arg_num, &reused) || + FindInplaceInput(*pnode, output_arg_num, &reused)) { // Reuse one of this node's input buffers as the output buffer (for in-place update) Reuse(reused, current); } else if (!context_.EnableParallelExecution() && FindReusableTensor(*node_output, &reused)) { @@ -610,7 +657,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 +668,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 e220606df7f58..af61984c70622 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,15 @@ 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/session_state_initializer.cc b/onnxruntime/core/framework/session_state_initializer.cc index ce7267848465e..b0857910a035d 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)); diff --git a/onnxruntime/core/framework/session_state_initializer.h b/onnxruntime/core/framework/session_state_initializer.h index 8f575ff0f2138..47516c3b7513f 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 fbf72956b98e8..ace90fa408965 100644 --- a/onnxruntime/core/providers/cpu/controlflow/loop.cc +++ b/onnxruntime/core/providers/cpu/controlflow/loop.cc @@ -328,8 +328,9 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca CreateInitialFeeds(feeds); auto& iter_num_value = *iter_num_mlvalue_.GetMutable()->MutableData(); + auto current_condition_mlvalue = condition_mlvalue_; - while (iter_num_value < max_trip_count_ && *condition_mlvalue_.GetMutable()->MutableData()) { + while (iter_num_value < max_trip_count_ && *current_condition_mlvalue.GetMutable()->MutableData()) { if (iter_num_value != 0) { UpdateFeeds(fetches, feeds); fetches.clear(); @@ -353,7 +354,7 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca ORT_RETURN_IF_ERROR(status); - condition_mlvalue_ = fetches[0]; + current_condition_mlvalue = fetches[0]; ++iter_num_value; } diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index a925d52ff2980..b48fe69c0458f 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -344,8 +344,8 @@ class InferenceSession::Impl { SessionStateInitializer initializer{model_location_, subgraph, *subgraph_session_state, execution_providers_, kernel_registry_manager_}; - ORT_RETURN_IF_ERROR(initializer.CreatePlan(node.ImplicitInputDefs(), - session_options_.enable_sequential_execution)); + ORT_RETURN_IF_ERROR(initializer.CreatePlan(&node, node.ImplicitInputDefs(), + session_options_.enable_sequential_execution)); ORT_RETURN_IF_ERROR(initializer.InitializeAndSave(&node.ImplicitInputDefs())); @@ -412,7 +412,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 3cc7b0a7a33aa..6c1910bdcb8bd 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 cb36139464f29..4c95b4778357f 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -64,7 +64,7 @@ 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, + 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,7 +214,7 @@ 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, + 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(); diff --git a/onnxruntime/test/providers/memcpy_test.cc b/onnxruntime/test/providers/memcpy_test.cc index f8a44f1458dd2..55eadcc6334f3 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();