From b392d5351fe98c87f53b5ab1d9552f706839cea4 Mon Sep 17 00:00:00 2001 From: pengwa Date: Thu, 23 Feb 2023 18:05:21 +0800 Subject: [PATCH] Fix memory profiler (#14695) ### Fix memory profiler A follow up fix for PR https://github.com/microsoft/onnxruntime/pull/13495 In ORTModule training, `PartialExecuteThePlan` is called twice, we need create log event after the backward graph run complete to collect the whole training graph's activations infos. Also change some log level to verbose, to avoid too many logs in > verbose log level. ### Motivation and Context --- onnxruntime/core/framework/execution_steps.cc | 6 +- onnxruntime/core/framework/memory_info.cc | 6 +- onnxruntime/core/framework/memory_info.h | 4 +- .../partial_graph_execution_state.cc | 4 +- .../core/framework/sequential_executor.cc | 121 +++++++++++------- .../core/framework/sequential_executor.h | 3 +- onnxruntime/core/framework/session_state.cc | 5 +- .../framework/stream_execution_context.cc | 90 ++++++------- onnxruntime/core/framework/utils.cc | 12 +- 9 files changed, 145 insertions(+), 106 deletions(-) diff --git a/onnxruntime/core/framework/execution_steps.cc b/onnxruntime/core/framework/execution_steps.cc index 929b0187fa..af3b1b9af3 100644 --- a/onnxruntime/core/framework/execution_steps.cc +++ b/onnxruntime/core/framework/execution_steps.cc @@ -38,7 +38,7 @@ Status WaitOnEPStep::Execute(StreamExecutionContext& ctx, if (ctx.GetDeviceStream(stream_idx)) { ctx.GetDeviceStream(stream_idx)->UpdateStreamClock(ctx.GetNotification(notification_idx_)->GetStreamSyncTable()); } - LOGS(ctx.GetLogger(), INFO) << "stream " << stream_idx << " wait on Notification with id: " << notification_idx_; + LOGS(ctx.GetLogger(), VERBOSE) << "stream " << stream_idx << " wait on Notification with id: " << notification_idx_; continue_flag = true; return Status::OK(); } @@ -85,8 +85,8 @@ Status ActivateNotificationStep::Execute(StreamExecutionContext& ctx, if (ctx.GetNotification(notification_idx_)) { ctx.GetNotification(notification_idx_)->ActivateAndUpdate(); } - LOGS(ctx.GetLogger(), INFO) << "stream " << stream_idx - << " activate notification with index " << notification_idx_; + LOGS(ctx.GetLogger(), VERBOSE) << "stream " << stream_idx + << " activate notification with index " << notification_idx_; continue_flag = true; return Status::OK(); } diff --git a/onnxruntime/core/framework/memory_info.cc b/onnxruntime/core/framework/memory_info.cc index 5b588bc272..41c8191778 100644 --- a/onnxruntime/core/framework/memory_info.cc +++ b/onnxruntime/core/framework/memory_info.cc @@ -266,7 +266,7 @@ void MemoryProfiler::CreateEvents(const std::string& p_name, const OrtDevice::DeviceType device_type) { // Metadata. std::string pid_name_internal = "device_" + std::to_string(device_type) + "_" + p_name + group_name; - events.push_back(CreateMetadataEvent(pid_name_internal, pid)); + events_.push_back(CreateMetadataEvent(pid_name_internal, pid)); size_t summary_size = 10; std::hash str_hash; @@ -350,14 +350,14 @@ void MemoryProfiler::CreateEvents(const std::string& p_name, size_t offset = IsStaticType(map_type) ? map.GetPlannedAddress(live_tensor) : map.GetAllocAddress(live_tensor); size_t size = IsStaticType(map_type) ? map.GetPlannedSize(live_tensor) : map.GetAllocSize(live_tensor); alloc_size_for_pattern += size; - events.push_back(CreateMemoryEvent(pid, item.first, name, offset, size, cname)); + events_.push_back(CreateMemoryEvent(pid, item.first, name, offset, size, cname)); has_other_tensors = true; } // If for that time steps, we have other tensors to plot other than just summary, add summary events. // These will show up visually as 10 % of the max width in a section. if (has_other_tensors) { summary_size = std::max(summary_size, item.second.total_size / 10); - events.push_back(CreateSummaryEvent(pid, item.first, item.second, summary_size, alloc_size_for_pattern)); + events_.push_back(CreateSummaryEvent(pid, item.first, item.second, summary_size, alloc_size_for_pattern)); } } } diff --git a/onnxruntime/core/framework/memory_info.h b/onnxruntime/core/framework/memory_info.h index d3d7db7f89..4b11dcaf7a 100644 --- a/onnxruntime/core/framework/memory_info.h +++ b/onnxruntime/core/framework/memory_info.h @@ -228,7 +228,7 @@ struct MemoryProfiler { const std::string& group_name, const size_t top_k, const OrtDevice::DeviceType device_t = OrtDevice::GPU); - const std::vector& GetEvents() { return events; } + const std::vector& GetEvents() { return events_; } size_t GetAndIncreasePid() { size_t val = pid_++; @@ -272,7 +272,7 @@ struct MemoryProfiler { "cq_build_attempt_failed", }; - std::vector events; + std::vector events_; // Key: the hash function of device+map_type. Value: (key: The time step. value: The allocation information) std::unordered_map > summary_; diff --git a/onnxruntime/core/framework/partial_graph_execution_state.cc b/onnxruntime/core/framework/partial_graph_execution_state.cc index e7235cac74..7b87c7fb60 100644 --- a/onnxruntime/core/framework/partial_graph_execution_state.cc +++ b/onnxruntime/core/framework/partial_graph_execution_state.cc @@ -1,6 +1,8 @@ -//// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. + #ifdef ENABLE_TRAINING + #include "core/framework/partial_graph_execution_state.h" #include "core/framework/session_state.h" #include "core/framework/stream_execution_context.h" diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index 8005fbf055..8c29ad8dbc 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -57,7 +57,8 @@ LARGE_INTEGER perf_freq = OrtGetPerformanceFrequency(); namespace onnxruntime { static void CalculateTotalOutputSizes(OpKernelContextInternal* op_kernel_context, - size_t& total_output_sizes, const std::string& node_name, std::string& output_type_shape) { + size_t& total_output_sizes, const std::string& node_name, + std::string& output_type_shape) { // Calculate total output sizes for this operation. std::stringstream ss; int added_type_shapes = 0; @@ -153,24 +154,27 @@ std::string ComposeSeriesName(const GraphViewer& graph_viewer) { class SessionScope { public: friend class KernelScope; - SessionScope(const SessionState& session_state, const ExecutionFrame& frame) : session_state_(session_state) + SessionScope(const SessionState& session_state, const ExecutionFrame& frame) + : session_state_(session_state) #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) - , - frame_(frame) + , + frame_(frame) #endif #ifdef CONCURRENCY_VISUALIZER - , - series_(ComposeSeriesName(session_state.GetGraphViewer()) + , + series_(ComposeSeriesName(session_state.GetGraphViewer())) #endif #ifdef ENABLE_NVTX_PROFILE - , - session_tag_(profile::Context::GetInstance().GetThreadTagOrDefault(std::this_thread::get_id())), - forward_range_("Batch-" + session_tag_ + " Forward", profile::Color::White), - backward_range_("Batch-" + session_tag_ + " Backward", profile::Color::Black) + , + session_tag_(profile::Context::GetInstance().GetThreadTagOrDefault(std::this_thread::get_id())), + forward_range_("Batch-" + session_tag_ + " Forward", profile::Color::White), + backward_range_("Batch-" + session_tag_ + " Backward", profile::Color::Black) #endif #ifdef DEBUG_NODE_INPUTS_OUTPUTS - , - dump_context_ {session_state_.GetGraphExecutionCounter(), 0} + , + dump_context_ { + session_state_.GetGraphExecutionCounter(), 0 + } #endif { if (session_state_.Profiler().IsEnabled()) { @@ -178,7 +182,7 @@ class SessionScope { } auto& logger = session_state_.Logger(); - LOGS(logger, INFO) << "Begin execution"; + LOGS(logger, VERBOSE) << "Begin execution"; const SequentialExecutionPlan& seq_exec_plan = *session_state_.GetExecutionPlan(); const auto& exec_plan_vec = seq_exec_plan.execution_plan; VLOGS(logger, 1) << "Size of execution plan vector: " << exec_plan_vec.size(); @@ -213,10 +217,12 @@ class SessionScope { #endif #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) - session_state_.GetMemoryProfiler()->CreateEvents( - "dynamic activations_" + std::to_string(session_state_.GetMemoryProfiler()->GetMemoryInfo().GetIteration()), - session_state_.GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0); - session_state_.GetMemoryProfiler()->Clear(); + if (flush_memory_info_) { + session_state_.GetMemoryProfiler()->CreateEvents( + "dynamic activations_" + std::to_string(session_state_.GetMemoryProfiler()->GetMemoryInfo().GetIteration()), + session_state_.GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0); + session_state_.GetMemoryProfiler()->Clear(); + } #endif if (session_state_.Profiler().IsEnabled()) { @@ -225,21 +231,31 @@ class SessionScope { #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) auto& logger = session_state_.Logger(); for (auto i : frame_.GetStaticMemorySizeInfo()) { - LOGS(logger, INFO) << "[Memory] ExecutionFrame statically allocates " - << i.second << " bytes for " << i.first << std::endl; + LOGS(logger, VERBOSE) << "[Memory] ExecutionFrame statically allocates " + << i.second << " bytes for " << i.first << std::endl; } for (auto i : frame_.GetDynamicMemorySizeInfo()) { - LOGS(logger, INFO) << "[Memory] ExecutionFrame dynamically allocates " - << i.second << " bytes for " << i.first << std::endl; + LOGS(logger, VERBOSE) << "[Memory] ExecutionFrame dynamically allocates " + << i.second << " bytes for " << i.first << std::endl; } #endif } + +#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) + void SetFlushMemoryInfoFlag(bool flush_memory_info) { + flush_memory_info_ = flush_memory_info; + } +#endif + private: const SessionState& session_state_; TimePoint session_start_; #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) const ExecutionFrame& frame_; + // Whether memory profiler need create events and flush to file. + // For partial graph run, when the last subgraph of the whole graph is executing, we need flush to file. + bool flush_memory_info_ = true; #endif #ifdef CONCURRENCY_VISUALIZER @@ -261,27 +277,30 @@ class KernelScope { public: KernelScope(SessionScope& session_scope, OpKernelContextInternal& kernel_context, - const OpKernel& kernel) : session_scope_(session_scope), - session_state_(session_scope_.session_state_), - kernel_context_(kernel_context), - kernel_(kernel) + const OpKernel& kernel) + : session_scope_(session_scope), + session_state_(session_scope_.session_state_), + kernel_context_(kernel_context), + kernel_(kernel) #ifdef CONCURRENCY_VISUALIZER - , - span_(session_scope_.series_, "%s.%d", kernel_.Node().OpType().c_str(), kernel_.Node().Index()) + , + span_(session_scope_.series_, "%s.%d", kernel_.Node().OpType().c_str(), kernel_.Node().Index()) #endif #ifdef ENABLE_NVTX_PROFILE - , - node_compute_range_(MakeString(kernel_.Node().OpType(), - ".", - kernel_.Node().Index(), - "(", - kernel_.Node().Name(), - ")"), - profile::Color::Yellow) + , + node_compute_range_(MakeString(kernel_.Node().OpType(), + ".", + kernel_.Node().Index(), + "(", + kernel_.Node().Name(), + ")"), + profile::Color::Yellow) #endif #ifdef DEBUG_NODE_INPUTS_OUTPUTS - , - dump_context_{session_scope_.dump_context_.iteration, kernel_.Node().Index()} + , + dump_context_ { + session_scope_.dump_context_.iteration, kernel_.Node().Index() + } #endif { #ifdef CONCURRENCY_VISUALIZER @@ -295,7 +314,8 @@ class KernelScope { if (node.Description() != "Backward pass" && !forward_range.IsBeginCalled()) { // Start timing forward pass when encountering the first forward node. forward_range.Begin(); - } else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled() && forward_range.IsBeginCalled()) { + } else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled() && + forward_range.IsBeginCalled()) { // Start timing backward pass when encountering the first backward node. // In the meanwhile, forward range ends. forward_range.End(); @@ -358,7 +378,8 @@ class KernelScope { {"output_size", std::to_string(total_output_sizes_)}, {"input_type_shape", input_type_shape_}, {"output_type_shape", output_type_shape_}, - {"thread_scheduling_stats", concurrency::ThreadPool::StopProfiling(session_state_.GetThreadPool())}, + {"thread_scheduling_stats", + concurrency::ThreadPool::StopProfiling(session_state_.GetThreadPool())}, }); auto sync_time_begin = profiler.Start(); profiler.EndTimeAndRecordEvent(profiling::NODE_EVENT, @@ -487,15 +508,16 @@ onnxruntime::Status ExecuteKernel(StreamExecutionContext& ctx, << "' Status Message: " << status.ErrorMessage(); // If the computation failed, we still can record the memory consumption #if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) - ctx.GetSessionState().GetMemoryProfiler()->CreateEvents("dynamic activations_" + std::to_string(ctx.GetSessionState().GetMemoryProfiler()->GetMemoryInfo().GetIteration()), - ctx.GetSessionState().GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0); + ctx.GetSessionState().GetMemoryProfiler()->CreateEvents( + "dynamic activations_" + std::to_string(ctx.GetSessionState().GetMemoryProfiler()->GetMemoryInfo().GetIteration()), + ctx.GetSessionState().GetMemoryProfiler()->GetAndIncreasePid(), MemoryInfo::MapType::DynamicActivation, "", 0); #endif const auto msg_string = ss.str(); LOGS(logger, ERROR) << msg_string; return Status(status.Category(), status.Code(), msg_string); } ctx.RecycleNodeInputs(idx); - LOGS(logger, INFO) << "stream " << stream_idx << " launch kernel with idx " << idx; + LOGS(logger, VERBOSE) << "stream " << stream_idx << " launch kernel with idx " << idx; return Status::OK(); } @@ -511,7 +533,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< const bool only_execute_path_to_fetches, bool single_thread_mode) { auto* execution_plan = session_state.GetExecutionPlan(); - LOGS(logger, INFO) << "Number of streams: " << execution_plan->execution_plan.size(); + LOGS(logger, VERBOSE) << "Number of streams: " << execution_plan->execution_plan.size(); int32_t valid_streams = 0; for (auto& stream : execution_plan->execution_plan) { if (stream && stream->steps_.size() > 0) @@ -595,13 +617,15 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl::span feed_mlvalue_idxs, gsl::span feeds, gsl::span fetch_mlvalue_idxs, std::vector& fetches, - const std::unordered_map& fetch_allocators, + const std::unordered_map& + fetch_allocators, const logging::Logger& logger, const DeviceStreamCollection* device_streams, const bool& terminate_flag, bool single_thread_mode, PartialGraphExecutionState& state, - const OrtValueCachePtr& cache) { + const OrtValueCachePtr& cache, + int32_t partial_graph_index) { auto& ctx = state.GetExecutionContext(feed_mlvalue_idxs, feeds, fetch_mlvalue_idxs, fetches, fetch_allocators, session_state, logger, device_streams); auto* plan = session_state.GetExecutionPlan(); @@ -610,6 +634,13 @@ onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl SessionScope session_scope(session_state, ctx.GetExecutionFrame()); +#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE) + // Only flush memory info for the 2nd partial graph execution (since ORTModule runs this function twice). + session_scope.SetFlushMemoryInfoFlag(partial_graph_index == 1); +#else + ORT_UNUSED_PARAMETER(partial_graph_index); +#endif + ctx.SetOrtValueCache(cache); auto* tp = single_thread_mode ? nullptr : session_state.GetInterOpThreadPool(); diff --git a/onnxruntime/core/framework/sequential_executor.h b/onnxruntime/core/framework/sequential_executor.h index c363513dcc..9bef4be34f 100644 --- a/onnxruntime/core/framework/sequential_executor.h +++ b/onnxruntime/core/framework/sequential_executor.h @@ -58,6 +58,7 @@ onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl const bool& terminate_flag, bool single_thread_mode, PartialGraphExecutionState& state, - const OrtValueCachePtr& cache); + const OrtValueCachePtr& cache, + int32_t partial_graph_index); #endif } // namespace onnxruntime diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index 7204683912..66699c907f 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -1381,10 +1381,9 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string(); - mem_profiler->Init(GetExecutionPlan(), GetOrtValueNameIdxMap()); - SetMemoryProfiler(mem_profiler.release()); + GetMemoryProfiler()->Init(GetExecutionPlan(), GetOrtValueNameIdxMap()); #endif // Note: For Training Prepacking should be always disabled. diff --git a/onnxruntime/core/framework/stream_execution_context.cc b/onnxruntime/core/framework/stream_execution_context.cc index 560358ec99..8919adac6c 100644 --- a/onnxruntime/core/framework/stream_execution_context.cc +++ b/onnxruntime/core/framework/stream_execution_context.cc @@ -9,28 +9,30 @@ namespace onnxruntime { #ifdef ORT_ENABLE_STREAM -StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state, - int32_t num_streams, - gsl::span notification_owners, - size_t num_barriers, - const DeviceStreamCollection* device_stream_map, - gsl::span feed_mlvalue_idxs, - gsl::span feeds, gsl::span fetch_mlvalue_idxs, - std::vector& fetches, - const std::unordered_map& fetch_allocators, - const logging::Logger& sess_logger, - bool single_thread_mode) : session_state_(&sess_state), - frame_(feed_mlvalue_idxs, - feeds, - fetch_mlvalue_idxs, - fetches, - fetch_allocators, - sess_state, - device_stream_map ? device_stream_map->GetStreams() : gsl::span({})), - logger_(&sess_logger), - single_thread_mode_(single_thread_mode), - device_stream_map_(device_stream_map), - count_down_barriers_(num_barriers) { +StreamExecutionContext::StreamExecutionContext(const SessionState& sess_state, + int32_t num_streams, + gsl::span notification_owners, + size_t num_barriers, + const DeviceStreamCollection* device_stream_map, + gsl::span feed_mlvalue_idxs, + gsl::span feeds, gsl::span fetch_mlvalue_idxs, + std::vector& fetches, + const std::unordered_map& + fetch_allocators, + const logging::Logger& sess_logger, + bool single_thread_mode) + : session_state_(&sess_state), + frame_(feed_mlvalue_idxs, + feeds, + fetch_mlvalue_idxs, + fetches, + fetch_allocators, + sess_state, + device_stream_map ? device_stream_map->GetStreams() : gsl::span({})), + logger_(&sess_logger), + single_thread_mode_(single_thread_mode), + device_stream_map_(device_stream_map), + count_down_barriers_(num_barriers) { notifications_.reserve(notification_owners.size()); for (size_t i = 0; i < notification_owners.size(); ++i) { auto* stream = device_stream_map_ ? device_stream_map_->GetStream(notification_owners[i]) : nullptr; @@ -49,7 +51,7 @@ StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state, #pragma warning(pop) #endif - // init barreris + // init barriers for (size_t i = 0; i < num_barriers; ++i) { count_down_barriers_[i].Set(2); } @@ -78,23 +80,25 @@ Stream* StreamExecutionContext ::GetDeviceStream(size_t idx) { } #else -StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state, - int32_t num_streams, - gsl::span feed_mlvalue_idxs, - gsl::span feeds, gsl::span fetch_mlvalue_idxs, - std::vector& fetches, - const std::unordered_map& fetch_allocators, - const logging::Logger& sess_logger, - bool single_thread_mode) : session_state_(&sess_state), - frame_(feed_mlvalue_idxs, - feeds, - fetch_mlvalue_idxs, - fetches, - fetch_allocators, - sess_state, - {}), - logger_(&sess_logger), - single_thread_mode_(single_thread_mode) { +StreamExecutionContext::StreamExecutionContext(const SessionState& sess_state, + int32_t num_streams, + gsl::span feed_mlvalue_idxs, + gsl::span feeds, gsl::span fetch_mlvalue_idxs, + std::vector& fetches, + const std::unordered_map& + fetch_allocators, + const logging::Logger& sess_logger, + bool single_thread_mode) + : session_state_(&sess_state), + frame_(feed_mlvalue_idxs, + feeds, + fetch_mlvalue_idxs, + fetches, + fetch_allocators, + sess_state, + {}), + logger_(&sess_logger), + single_thread_mode_(single_thread_mode) { #ifdef _WIN32 #pragma warning(push) #pragma warning(disable : 26409 26400) @@ -158,14 +162,14 @@ void StreamExecutionContext ::SetStatus(Status& status) { task_status_ = status; } -StreamExecutionContext ::~StreamExecutionContext() {} +StreamExecutionContext::~StreamExecutionContext() {} -void StreamExecutionContext ::RecycleNodeInputs(onnxruntime::NodeIndex node_index) { +void StreamExecutionContext::RecycleNodeInputs(onnxruntime::NodeIndex node_index) { auto* execution_plan = session_state_->GetExecutionPlan(); for (auto idx : execution_plan->node_release_list[node_index]) { if (--release_plan_[idx] == 0) { ORT_ENFORCE(frame_.ReleaseMLValue(static_cast(execution_plan->release_actions[idx].value_index)).IsOK()); - LOGS(*logger_, INFO) << "ort value " << execution_plan->release_actions[idx].value_index << " released"; + LOGS(*logger_, VERBOSE) << "ort value " << execution_plan->release_actions[idx].value_index << " released"; } } } diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 40b566b098..0fa25657c0 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -822,6 +822,7 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF const logging::Logger& logger, PartialGraphExecutionState& state, const OrtValueCachePtr& cache, const bool& terminate_flag, DeviceStreamCollection* device_stream_collection, + int32_t partial_graph_index, Stream* parent_stream) { // finalize the copy info using the provided feeds and fetches. will update device_copy_checks in the background FinalizeFeedFetchCopyInfo(feeds_fetches_manager, feeds, fetches); @@ -846,7 +847,8 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF // single thread mode single_thread_mode, state, - cache)); + cache, + partial_graph_index)); } else { auto p_feeds = feeds; std::vector* p_fetches = &fetches; @@ -907,7 +909,8 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF // single thread mode single_thread_mode, state, - cache)); + cache, + partial_graph_index)); InlinedVector fetches_streams; fetches_streams.reserve(feeds_fetches_info.fetches_mlvalue_idxs.size()); @@ -936,13 +939,12 @@ common::Status ExecutePartialGraph(const SessionState& session_state, FeedsFetch gsl::span feeds, std::vector& fetches, const logging::Logger& logger, PartialGraphExecutionState& state, const OrtValueCachePtr& cache, const bool& terminate_flag, - // TODO: handle the partial_graph_index - int32_t /*partial_graph_index*/, + int32_t partial_graph_index, Stream* parent_stream) { DeviceStreamCollection* device_stream_collection = state.GetDeviceStreamCollection(session_state); auto retval = ExecutePartialGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, logger, state, cache, terminate_flag, device_stream_collection, - parent_stream); + partial_graph_index, parent_stream); if (device_stream_collection) ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(false)); return retval;