mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
Avoid the lock for device stream impact the cpu build (#14131)
### Description Introduce a runtime flag in SessionState about whether any EP in current session using stream feature, if no, avoid trigger the lock. This will avoid the impact to CPU build. ### Motivation and Context Currently we use a lock in SessionState when retrieve device stream collection, this is mainly for reusing the device stream for EP like GPU eps, so it shouldn't impact the build which doesn't using stream feature, like cpu build. Instead of play with build flags, this PR introduce a runtime flag in SessionState to indicate whether current session has any EP that using the stream feature. if no, we don't need to trigger the lock. Co-authored-by: Cheng Tang <chenta@microsoft.com@orttrainingdev9.d32nl1ml4oruzj4qz3bqlggovf.px.internal.cloudapp.net>
This commit is contained in:
Tang, Cheng
GitHub
parent
4eac0db3af
commit
90cff21fa7
9 changed files with 127 additions and 83 deletions
|
|
@ -44,7 +44,7 @@ ProgramRegion& PartialGraphExecutionState::GetProgramRegions(const SessionState&
|
|||
PartialGraphExecutionState::~PartialGraphExecutionState() {
|
||||
}
|
||||
|
||||
DeviceStreamCollection& PartialGraphExecutionState::GetDeviceStreamCollection(const SessionState& session_state) {
|
||||
DeviceStreamCollection* PartialGraphExecutionState::GetDeviceStreamCollection(const SessionState& session_state) {
|
||||
if (device_stream_collection_ == nullptr) {
|
||||
device_stream_collection_ = session_state.AcquireDeviceStreamCollection();
|
||||
// the life-time of partial graph execution state is in-consistant with session,
|
||||
|
|
@ -53,7 +53,7 @@ DeviceStreamCollection& PartialGraphExecutionState::GetDeviceStreamCollection(co
|
|||
// so let's always delete the stream collections.
|
||||
// luckily, for ort module, we always running with default stream, so no impact to perf.
|
||||
}
|
||||
return *device_stream_collection_;
|
||||
return device_stream_collection_.get();
|
||||
}
|
||||
|
||||
StreamExecutionContext& PartialGraphExecutionState::GetExecutionContext(gsl::span<const int>& feed_mlvalue_idxs, gsl::span<const OrtValue>& feeds,
|
||||
|
|
@ -61,7 +61,7 @@ StreamExecutionContext& PartialGraphExecutionState::GetExecutionContext(gsl::spa
|
|||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const SessionState& session_state,
|
||||
const logging::Logger& sess_logger,
|
||||
const DeviceStreamCollection& device_streams) {
|
||||
const DeviceStreamCollection* device_streams) {
|
||||
if (execution_context_ == nullptr) {
|
||||
auto* execution_plan = session_state.GetExecutionPlan();
|
||||
LOGS(sess_logger, INFO) << "Number of streams: " << execution_plan->execution_plan.size();
|
||||
|
|
|
|||
|
|
@ -35,8 +35,8 @@ struct PartialGraphExecutionState {
|
|||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const SessionState& session_state,
|
||||
const logging::Logger& sess_logger,
|
||||
const DeviceStreamCollection& device_streams);
|
||||
DeviceStreamCollection& GetDeviceStreamCollection(const SessionState& session_state);
|
||||
const DeviceStreamCollection* device_streams);
|
||||
DeviceStreamCollection* GetDeviceStreamCollection(const SessionState& session_state);
|
||||
|
||||
private:
|
||||
std::unique_ptr<StreamExecutionContext> execution_context_;
|
||||
|
|
|
|||
|
|
@ -500,7 +500,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span<
|
|||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const logging::Logger& logger,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
const DeviceStreamCollection& device_streams,
|
||||
const DeviceStreamCollection* device_streams,
|
||||
#endif
|
||||
const bool& terminate_flag,
|
||||
const bool only_execute_path_to_fetches,
|
||||
|
|
@ -592,7 +592,7 @@ onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl
|
|||
std::vector<OrtValue>& fetches,
|
||||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const logging::Logger& logger,
|
||||
const DeviceStreamCollection& device_streams,
|
||||
const DeviceStreamCollection* device_streams,
|
||||
const bool& terminate_flag,
|
||||
bool single_thread_mode,
|
||||
PartialGraphExecutionState& state,
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span<
|
|||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const logging::Logger& logger,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
const DeviceStreamCollection& device_streams,
|
||||
const DeviceStreamCollection* device_streams,
|
||||
#endif
|
||||
const bool& terminate_flag,
|
||||
const bool only_execute_path_to_fetches,
|
||||
|
|
@ -51,7 +51,7 @@ onnxruntime::Status PartialExecuteThePlan(const SessionState& session_state, gsl
|
|||
std::vector<OrtValue>& fetches,
|
||||
const std::unordered_map<size_t, IExecutor::CustomAllocator>& fetch_allocators,
|
||||
const logging::Logger& logger,
|
||||
const DeviceStreamCollection& device_streams,
|
||||
const DeviceStreamCollection* device_streams,
|
||||
const bool& terminate_flag,
|
||||
bool single_thread_mode,
|
||||
PartialGraphExecutionState& state,
|
||||
|
|
|
|||
|
|
@ -1350,17 +1350,17 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
|
|||
session_options.execution_order,
|
||||
session_options.enable_mem_reuse);
|
||||
auto status = SequentialPlanner::CreatePlan(parent_node, *graph_viewer_, valid_outer_scope_node_args,
|
||||
execution_providers_, kernel_create_info_map_,
|
||||
subgraphs_kernel_create_info_maps,
|
||||
outer_scope_node_arg_to_location_map,
|
||||
ort_value_name_idx_map_, context,
|
||||
execution_providers_, kernel_create_info_map_,
|
||||
subgraphs_kernel_create_info_maps,
|
||||
outer_scope_node_arg_to_location_map,
|
||||
ort_value_name_idx_map_, context,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
GetStreamHandleRegistryInstance(),
|
||||
GetStreamHandleRegistryInstance(),
|
||||
#endif
|
||||
session_options.config_options.GetConfigOrDefault(
|
||||
kNodePartitionConfigFile, ""),
|
||||
Logger(),
|
||||
p_seq_exec_plan_);
|
||||
session_options.config_options.GetConfigOrDefault(
|
||||
kNodePartitionConfigFile, ""),
|
||||
Logger(),
|
||||
p_seq_exec_plan_);
|
||||
ORT_RETURN_IF_ERROR(status);
|
||||
|
||||
// Record the allocation plan
|
||||
|
|
@ -1410,6 +1410,23 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
|
|||
|
||||
#endif
|
||||
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
// set the has_device_stream_enabled_ep_ flag
|
||||
has_device_stream_enabled_ep_ = false;
|
||||
if (p_seq_exec_plan_.has_value()) {
|
||||
auto& execution_plan = (*p_seq_exec_plan_).execution_plan;
|
||||
for (size_t i = 0; i < execution_plan.size(); ++i) {
|
||||
auto& logic_stream = execution_plan[i];
|
||||
if (logic_stream->steps_.size() > 0) {
|
||||
auto create_stream_fn = GetStreamHandleRegistryInstance().GetCreateStreamFn(logic_stream->device_.Type());
|
||||
if (create_stream_fn) {
|
||||
has_device_stream_enabled_ep_ = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
ORT_RETURN_IF_ERROR(
|
||||
session_state_utils::SaveInitializedTensors(
|
||||
Env::Default(), graph_location, *graph_viewer_,
|
||||
|
|
@ -1526,21 +1543,31 @@ static void BindToDeviceStream(const SequentialExecutionPlan& execution_plan,
|
|||
}
|
||||
|
||||
std::unique_ptr<DeviceStreamCollection> SessionState::AcquireDeviceStreamCollection() const {
|
||||
std::lock_guard<onnxruntime::OrtMutex> lock(device_stream_pool_mutex_);
|
||||
if (!device_stream_pool_.empty()) {
|
||||
auto device_stream = std::move(device_stream_pool_.back());
|
||||
device_stream_pool_.pop_back();
|
||||
return device_stream;
|
||||
if (has_device_stream_enabled_ep_) {
|
||||
std::lock_guard<onnxruntime::OrtMutex> lock(device_stream_pool_mutex_);
|
||||
if (!device_stream_pool_.empty()) {
|
||||
auto device_stream = std::move(device_stream_pool_.back());
|
||||
device_stream_pool_.pop_back();
|
||||
return device_stream;
|
||||
} else {
|
||||
auto device_stream = std::make_unique<DeviceStreamCollection>(this->GetExecutionPlan()->execution_plan.size(), *this);
|
||||
BindToDeviceStream(*this->GetExecutionPlan(), *device_stream, *stream_handles_registry_);
|
||||
return device_stream;
|
||||
}
|
||||
} else {
|
||||
auto device_stream = std::make_unique<DeviceStreamCollection>(this->GetExecutionPlan()->execution_plan.size(), *this);
|
||||
BindToDeviceStream(*this->GetExecutionPlan(), *device_stream, *stream_handles_registry_);
|
||||
return device_stream;
|
||||
// no reusing of device stream is needed, just return nullptr, the caller will handle it
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
void SessionState::RecycleDeviceStreamCollection(std::unique_ptr<DeviceStreamCollection> device_stream_collection) const {
|
||||
std::lock_guard<onnxruntime::OrtMutex> lock(device_stream_pool_mutex_);
|
||||
device_stream_pool_.push_back(std::move(device_stream_collection));
|
||||
// if no need to reuse the device stream, don't perform the recycle
|
||||
if (has_device_stream_enabled_ep_) {
|
||||
std::lock_guard<onnxruntime::OrtMutex> lock(device_stream_pool_mutex_);
|
||||
device_stream_pool_.push_back(std::move(device_stream_collection));
|
||||
} else {
|
||||
device_stream_collection.reset(nullptr);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
|
|
|
|||
|
|
@ -555,6 +555,8 @@ class SessionState {
|
|||
// lock for the device stream pool
|
||||
mutable OrtMutex device_stream_pool_mutex_;
|
||||
mutable std::vector<std::unique_ptr<DeviceStreamCollection>> device_stream_pool_;
|
||||
// flag to indicate whether current session using any EP that create device stream dynamically.
|
||||
bool has_device_stream_enabled_ep_ = false;
|
||||
#endif
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ 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,
|
||||
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,
|
||||
|
|
@ -26,14 +26,14 @@ StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state,
|
|||
fetches,
|
||||
fetch_allocators,
|
||||
sess_state,
|
||||
device_stream_map.GetStreams()),
|
||||
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_.GetStream(notification_owners[i]);
|
||||
auto* stream = device_stream_map_ ? device_stream_map_->GetStream(notification_owners[i]) : nullptr;
|
||||
if (stream)
|
||||
notifications_.emplace_back(stream->CreateNotification(/*TODO: calculate num of consumers*/ 0));
|
||||
else
|
||||
|
|
@ -69,8 +69,12 @@ bool StreamExecutionContext ::DecCountDownBarrier(size_t barrier_id) {
|
|||
}
|
||||
|
||||
Stream* StreamExecutionContext ::GetDeviceStream(size_t idx) {
|
||||
ORT_ENFORCE(idx < device_stream_map_.NumStreams());
|
||||
return device_stream_map_.GetStream(idx);
|
||||
if (device_stream_map_) {
|
||||
ORT_ENFORCE(idx < device_stream_map_->NumStreams());
|
||||
return device_stream_map_->GetStream(idx);
|
||||
} else {
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
#else
|
||||
|
|
@ -110,15 +114,15 @@ StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state,
|
|||
}
|
||||
|
||||
synchronize::Notification* StreamExecutionContext ::GetNotification(size_t /*idx*/) {
|
||||
ORT_THROW("Try to get notification in a build which doesn't enable Stream!");
|
||||
ORT_THROW("Try to get notification in a build which doesn't enable Stream!");
|
||||
}
|
||||
|
||||
bool StreamExecutionContext ::DecCountDownBarrier(size_t /*barrier_id*/) {
|
||||
ORT_THROW("Try to decrease barrier in a build which doesn't enable Stream!");
|
||||
ORT_THROW("Try to decrease barrier in a build which doesn't enable Stream!");
|
||||
}
|
||||
|
||||
Stream* StreamExecutionContext ::GetDeviceStream(size_t /*idx*/) {
|
||||
return nullptr;
|
||||
return nullptr;
|
||||
}
|
||||
#endif
|
||||
|
||||
|
|
|
|||
|
|
@ -61,7 +61,7 @@ class StreamExecutionContext {
|
|||
#ifdef ORT_ENABLE_STREAM
|
||||
gsl::span<const size_t> notification_owners,
|
||||
size_t num_barriers,
|
||||
const DeviceStreamCollection& device_stream_map,
|
||||
const DeviceStreamCollection* device_stream_map,
|
||||
#endif
|
||||
gsl::span<const int> feed_mlvalue_idxs,
|
||||
gsl::span<const OrtValue> feeds, gsl::span<const int> fetch_mlvalue_idxs,
|
||||
|
|
@ -171,7 +171,8 @@ class StreamExecutionContext {
|
|||
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
InlinedVector<std::unique_ptr<synchronize::Notification>> notifications_;
|
||||
const DeviceStreamCollection& device_stream_map_;
|
||||
// if it is nullptr, means current session doesn't have any EP using stream feature
|
||||
const DeviceStreamCollection* device_stream_map_;
|
||||
|
||||
std::vector<CountDownBarrier> count_down_barriers_;
|
||||
#endif
|
||||
|
|
|
|||
|
|
@ -566,12 +566,15 @@ static common::Status CopyOutputsAcrossDevices(const SessionState& session_state
|
|||
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
struct DeviceStreamCollectionHolder {
|
||||
DeviceStreamCollectionHolder(const SessionState& session_state) : session_state_(session_state),
|
||||
p_(session_state.AcquireDeviceStreamCollection()) {
|
||||
DeviceStreamCollectionHolder(
|
||||
const SessionState& session_state) : session_state_(session_state),
|
||||
p_(session_state.AcquireDeviceStreamCollection()) {
|
||||
}
|
||||
|
||||
~DeviceStreamCollectionHolder() {
|
||||
session_state_.RecycleDeviceStreamCollection(std::move(p_));
|
||||
if (p_) {
|
||||
session_state_.RecycleDeviceStreamCollection(std::move(p_));
|
||||
}
|
||||
}
|
||||
|
||||
const SessionState& session_state_;
|
||||
|
|
@ -611,7 +614,7 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
ExecutionMode execution_mode, const bool& terminate_flag,
|
||||
const logging::Logger& logger,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
DeviceStreamCollection& device_stream_collection,
|
||||
DeviceStreamCollection* device_stream_collection,
|
||||
#endif
|
||||
const bool only_execute_path_to_fetches = false,
|
||||
Stream* parent_stream = nullptr) {
|
||||
|
|
@ -619,7 +622,8 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
const auto& device_copy_checks = feeds_fetches_manager.GetDeviceCopyChecks();
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
auto* execution_plan = session_state.GetExecutionPlan();
|
||||
UpdateWithParentStream(device_stream_collection, parent_stream);
|
||||
if (device_stream_collection)
|
||||
UpdateWithParentStream(*device_stream_collection, parent_stream);
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(parent_stream);
|
||||
#endif
|
||||
|
|
@ -639,16 +643,16 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
if (device_copy_checks.status == DeviceCopyCheck::NoCopy) {
|
||||
// no device copies are needed so simple execute
|
||||
auto status = (ExecuteThePlan(session_state,
|
||||
feeds_fetches_info.feeds_mlvalue_idxs, feeds,
|
||||
feeds_fetches_info.fetches_mlvalue_idxs, fetches, fetch_allocators,
|
||||
logger,
|
||||
feeds_fetches_info.feeds_mlvalue_idxs, feeds,
|
||||
feeds_fetches_info.fetches_mlvalue_idxs, fetches, fetch_allocators,
|
||||
logger,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
device_stream_collection,
|
||||
device_stream_collection,
|
||||
#endif
|
||||
terminate_flag,
|
||||
only_execute_path_to_fetches,
|
||||
// single thread mode
|
||||
single_thread_mode));
|
||||
terminate_flag,
|
||||
only_execute_path_to_fetches,
|
||||
// single thread mode
|
||||
single_thread_mode));
|
||||
ORT_RETURN_IF_ERROR(status);
|
||||
} else {
|
||||
auto feeds_to_use = feeds;
|
||||
|
|
@ -663,15 +667,16 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
// TODO: we can pre-calculate the stream index for graph inputs in execution plan
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
for (auto& copy_info : feed_copy_info) {
|
||||
|
||||
auto& device = copy_info.target_device;
|
||||
auto streams = device_stream_collection.GetStreams();
|
||||
bool found = false;
|
||||
for (auto* stream : streams) {
|
||||
if (stream && stream->GetDevice().Type() == device.Type()) {
|
||||
feed_streams.push_back(stream);
|
||||
found = true;
|
||||
break;
|
||||
if (device_stream_collection) {
|
||||
auto streams = device_stream_collection->GetStreams();
|
||||
for (auto* stream : streams) {
|
||||
if (stream && stream->GetDevice().Type() == device.Type()) {
|
||||
feed_streams.push_back(stream);
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!found)
|
||||
|
|
@ -707,15 +712,15 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
|
||||
// no device copies are needed so simple execute
|
||||
auto status = (ExecuteThePlan(session_state,
|
||||
feeds_fetches_info.feeds_mlvalue_idxs, feeds_to_use,
|
||||
feeds_fetches_info.fetches_mlvalue_idxs, *p_fetches, fetch_allocators,
|
||||
logger,
|
||||
feeds_fetches_info.feeds_mlvalue_idxs, feeds_to_use,
|
||||
feeds_fetches_info.fetches_mlvalue_idxs, *p_fetches, fetch_allocators,
|
||||
logger,
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
device_stream_collection,
|
||||
device_stream_collection,
|
||||
#endif
|
||||
terminate_flag,
|
||||
only_execute_path_to_fetches,
|
||||
single_thread_mode));
|
||||
terminate_flag,
|
||||
only_execute_path_to_fetches,
|
||||
single_thread_mode));
|
||||
ORT_RETURN_IF_ERROR(status);
|
||||
InlinedVector<Stream*> fetches_streams;
|
||||
fetches_streams.reserve(feeds_fetches_info.fetches_mlvalue_idxs.size());
|
||||
|
|
@ -724,7 +729,7 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
for (auto fetch_idx : feeds_fetches_info.fetches_mlvalue_idxs) {
|
||||
auto it = value_to_stream_map.find(fetch_idx);
|
||||
if (it != value_to_stream_map.end()) {
|
||||
fetches_streams.push_back(device_stream_collection.GetStream(it->second));
|
||||
fetches_streams.push_back(device_stream_collection ? device_stream_collection->GetStream(it->second) : nullptr);
|
||||
} else {
|
||||
// for subgraph, it is possible the graph is empty,
|
||||
// the fetches are come from parent graph.
|
||||
|
|
@ -737,7 +742,6 @@ ExecuteGraphImpl(const SessionState& session_state,
|
|||
}
|
||||
#endif
|
||||
|
||||
|
||||
if (device_copy_checks.output_copy_needed == DeviceCopyCheck::Copy) {
|
||||
ORT_RETURN_IF_ERROR(CopyOutputsAcrossDevices(session_state, *p_fetches, fetches, fetch_copy_info, fetches_streams));
|
||||
}
|
||||
|
|
@ -758,13 +762,14 @@ common::Status ExecuteGraph(const SessionState& session_state,
|
|||
FinalizeFeedFetchCopyInfo(feeds_fetches_manager, feeds, fetches);
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
DeviceStreamCollectionHolder device_stream_collection_holder(session_state);
|
||||
DeviceStreamCollection& device_stream_collection = *device_stream_collection_holder.p_;
|
||||
DeviceStreamCollection* device_stream_collection = device_stream_collection_holder.p_.get();
|
||||
auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, {},
|
||||
execution_mode, terminate_flag, logger,
|
||||
device_stream_collection,
|
||||
only_execute_path_to_fetches,
|
||||
parent_stream);
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection.CleanUp(sync_execution_provider));
|
||||
if (device_stream_collection)
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(sync_execution_provider));
|
||||
return retval;
|
||||
#else
|
||||
ORT_UNUSED_PARAMETER(sync_execution_provider);
|
||||
|
|
@ -806,7 +811,7 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF
|
|||
gsl::span<const OrtValue> feeds, std::vector<OrtValue>& fetches,
|
||||
const logging::Logger& logger, PartialGraphExecutionState& state,
|
||||
const OrtValueCachePtr& cache, const bool& terminate_flag,
|
||||
DeviceStreamCollection& device_stream_collection,
|
||||
DeviceStreamCollection* device_stream_collection,
|
||||
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);
|
||||
|
|
@ -816,8 +821,8 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF
|
|||
bool single_thread_mode = true;
|
||||
|
||||
auto* execution_plan = session_state.GetExecutionPlan();
|
||||
|
||||
UpdateWithParentStream(device_stream_collection, parent_stream);
|
||||
if (device_stream_collection)
|
||||
UpdateWithParentStream(*device_stream_collection, parent_stream);
|
||||
|
||||
// see if we can skip copies due to the types of execution providers available
|
||||
if (device_copy_checks.status == DeviceCopyCheck::NoCopy) {
|
||||
|
|
@ -845,13 +850,16 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF
|
|||
// TODO: we can pre-calculate the stream index for graph inputs in execution plan
|
||||
for (auto& copy_info : feed_copy_info) {
|
||||
auto& device = copy_info.target_device;
|
||||
auto streams = device_stream_collection.GetStreams();
|
||||
bool found = false;
|
||||
for (auto* stream : streams) {
|
||||
if (stream && stream->GetDevice().Type() == device.Type()) {
|
||||
feed_streams.push_back(stream);
|
||||
found = true;
|
||||
break;
|
||||
if (device_stream_collection) {
|
||||
auto streams = device_stream_collection->GetStreams();
|
||||
|
||||
for (auto* stream : streams) {
|
||||
if (stream && stream->GetDevice().Type() == device.Type()) {
|
||||
feed_streams.push_back(stream);
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!found)
|
||||
|
|
@ -897,7 +905,7 @@ common::Status ExecutePartialGraphImpl(const SessionState& session_state, FeedsF
|
|||
for (auto fetch_idx : feeds_fetches_info.fetches_mlvalue_idxs) {
|
||||
auto it = value_to_stream_map.find(fetch_idx);
|
||||
if (it != value_to_stream_map.end()) {
|
||||
fetches_streams.push_back(device_stream_collection.GetStream(it->second));
|
||||
fetches_streams.push_back(device_stream_collection ? device_stream_collection->GetStream(it->second) : nullptr);
|
||||
} else {
|
||||
// for subgraph, it is possible the graph is empty,
|
||||
// the fetches are come from parent graph.
|
||||
|
|
@ -921,11 +929,12 @@ common::Status ExecutePartialGraph(const SessionState& session_state, FeedsFetch
|
|||
// TODO: handle the partial_graph_index
|
||||
int32_t /*partial_graph_index*/,
|
||||
Stream* parent_stream) {
|
||||
DeviceStreamCollection& device_stream_collection = state.GetDeviceStreamCollection(session_state);
|
||||
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);
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection.CleanUp(false));
|
||||
if (device_stream_collection)
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(false));
|
||||
return retval;
|
||||
}
|
||||
#endif
|
||||
|
|
@ -938,11 +947,12 @@ common::Status ExecuteSubgraph(const SessionState& session_state, const FeedsFet
|
|||
bool sync_subgraph_fetches) {
|
||||
#ifdef ORT_ENABLE_STREAM
|
||||
DeviceStreamCollectionHolder device_stream_collection_holder(session_state);
|
||||
DeviceStreamCollection& device_stream_collection = *device_stream_collection_holder.p_;
|
||||
DeviceStreamCollection* device_stream_collection = device_stream_collection_holder.p_.get();
|
||||
|
||||
auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators,
|
||||
execution_mode, terminate_flag, logger, device_stream_collection, false, parent_stream);
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection.CleanUp(false));
|
||||
if (device_stream_collection)
|
||||
ORT_CHECK_AND_SET_RETVAL(device_stream_collection->CleanUp(false));
|
||||
#else
|
||||
auto retval = ExecuteGraphImpl(session_state, feeds_fetches_manager, feeds, fetches, fetch_allocators,
|
||||
execution_mode, terminate_flag, logger, false, parent_stream);
|
||||
|
|
|
|||
Loading…
Reference in a new issue