saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix memory profiler (#14695)

Browse source 
### 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
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
pengwa 2023-02-23 18:05:21 +08:00 committed by GitHub
parent 62ee0c8110
commit b392d5351f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 145 additions and 106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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();
}

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

@ -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<std::string> 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));
}
}
}

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

@ -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<std::string>& GetEvents() { return events; }
const std::vector<std::string>& GetEvents() { return events_; }
size_t GetAndIncreasePid() {
size_t val = pid_++;
@ -272,7 +272,7 @@ struct MemoryProfiler {
"cq_build_attempt_failed",
};
std::vector<std::string> events;
std::vector<std::string> events_;
// Key: the hash function of device+map_type. Value: (key: The time step. value: The allocation information)
std::unordered_map<size_t, std::unordered_map<size_t, MemoryInfo::AllocationSummary> > summary_;

4
onnxruntime/core/framework/partial_graph_execution_state.cc
View file

@ -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"

121
onnxruntime/core/framework/sequential_executor.cc
View file

@ -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<const int> feed_mlvalue_idxs,
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
const std::unordered_map<size_t, IExecutor::CustomAllocator>&
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();

3
onnxruntime/core/framework/sequential_executor.h
View file

@ -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

5
onnxruntime/core/framework/session_state.cc
View file

@ -1381,10 +1381,9 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
// Uncomment the below to dump the allocation plan to std::cout
// LOGS(logger_, VERBOSE) << std::make_pair(p_seq_exec_plan_.get(), this);
#if !defined(ORT_MINIMAL_BUILD) && defined(ORT_MEMORY_PROFILE)
auto mem_profiler = std::make_unique<MemoryProfiler>();
mem_profiler->Init(GetExecutionPlan(), GetOrtValueNameIdxMap());
SetMemoryProfiler(mem_profiler.release());
GetMemoryProfiler()->Init(GetExecutionPlan(), GetOrtValueNameIdxMap());
#endif
// Note: For Training Prepacking should be always disabled.

90
onnxruntime/core/framework/stream_execution_context.cc
View file

@ -9,28 +9,30 @@
namespace onnxruntime {
#ifdef ORT_ENABLE_STREAM
StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state,
int32_t num_streams,
gsl::span<const size_t> notification_owners,
size_t num_barriers,
const DeviceStreamCollection* device_stream_map,
gsl::span<const int> feed_mlvalue_idxs,
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& 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<Stream*>({})),
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<const size_t> notification_owners,
size_t num_barriers,
const DeviceStreamCollection* device_stream_map,
gsl::span<const int> feed_mlvalue_idxs,
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>&
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<Stream*>({})),
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<const int> feed_mlvalue_idxs,
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>& 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<const int> feed_mlvalue_idxs,
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, IExecutor::CustomAllocator>&
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<int>(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";
}
}
}

12
onnxruntime/core/framework/utils.cc
View file

@ -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<OrtValue>* 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<Stream*> 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<const OrtValue> feeds, std::vector<OrtValue>& 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;