From 90cff21fa72bd103b3666fc3a72c7ba7504acebe Mon Sep 17 00:00:00 2001 From: "Tang, Cheng" Date: Thu, 5 Jan 2023 09:01:33 -0800 Subject: [PATCH] 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 --- .../partial_graph_execution_state.cc | 6 +- .../framework/partial_graph_execution_state.h | 4 +- .../core/framework/sequential_executor.cc | 4 +- .../core/framework/sequential_executor.h | 4 +- onnxruntime/core/framework/session_state.cc | 65 ++++++++---- onnxruntime/core/framework/session_state.h | 2 + .../framework/stream_execution_context.cc | 20 ++-- .../core/framework/stream_execution_context.h | 5 +- onnxruntime/core/framework/utils.cc | 100 ++++++++++-------- 9 files changed, 127 insertions(+), 83 deletions(-) diff --git a/onnxruntime/core/framework/partial_graph_execution_state.cc b/onnxruntime/core/framework/partial_graph_execution_state.cc index e4c0cc01e1..e7235cac74 100644 --- a/onnxruntime/core/framework/partial_graph_execution_state.cc +++ b/onnxruntime/core/framework/partial_graph_execution_state.cc @@ -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& feed_mlvalue_idxs, gsl::span& feeds, @@ -61,7 +61,7 @@ StreamExecutionContext& PartialGraphExecutionState::GetExecutionContext(gsl::spa const std::unordered_map& 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(); diff --git a/onnxruntime/core/framework/partial_graph_execution_state.h b/onnxruntime/core/framework/partial_graph_execution_state.h index 8a78e72da3..80f7e52a9c 100644 --- a/onnxruntime/core/framework/partial_graph_execution_state.h +++ b/onnxruntime/core/framework/partial_graph_execution_state.h @@ -35,8 +35,8 @@ struct PartialGraphExecutionState { const std::unordered_map& 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 execution_context_; diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index 537b78414a..0736625807 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -500,7 +500,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< const std::unordered_map& 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& fetches, const std::unordered_map& fetch_allocators, const logging::Logger& logger, - const DeviceStreamCollection& device_streams, + const DeviceStreamCollection* device_streams, const bool& terminate_flag, bool single_thread_mode, PartialGraphExecutionState& state, diff --git a/onnxruntime/core/framework/sequential_executor.h b/onnxruntime/core/framework/sequential_executor.h index a4943f5366..ce5f12ab08 100644 --- a/onnxruntime/core/framework/sequential_executor.h +++ b/onnxruntime/core/framework/sequential_executor.h @@ -39,7 +39,7 @@ onnxruntime::Status ExecuteThePlan(const SessionState& session_state, gsl::span< const std::unordered_map& 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& fetches, const std::unordered_map& fetch_allocators, const logging::Logger& logger, - const DeviceStreamCollection& device_streams, + const DeviceStreamCollection* device_streams, const bool& terminate_flag, bool single_thread_mode, PartialGraphExecutionState& state, diff --git a/onnxruntime/core/framework/session_state.cc b/onnxruntime/core/framework/session_state.cc index fdd2b85315..4cc41dc25c 100644 --- a/onnxruntime/core/framework/session_state.cc +++ b/onnxruntime/core/framework/session_state.cc @@ -1350,17 +1350,17 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_stringsteps_.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 SessionState::AcquireDeviceStreamCollection() const { - std::lock_guard 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 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(this->GetExecutionPlan()->execution_plan.size(), *this); + BindToDeviceStream(*this->GetExecutionPlan(), *device_stream, *stream_handles_registry_); + return device_stream; + } } else { - auto device_stream = std::make_unique(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 device_stream_collection) const { - std::lock_guard 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 lock(device_stream_pool_mutex_); + device_stream_pool_.push_back(std::move(device_stream_collection)); + } else { + device_stream_collection.reset(nullptr); + } } #endif diff --git a/onnxruntime/core/framework/session_state.h b/onnxruntime/core/framework/session_state.h index 6d0022065c..8c5ca41767 100644 --- a/onnxruntime/core/framework/session_state.h +++ b/onnxruntime/core/framework/session_state.h @@ -555,6 +555,8 @@ class SessionState { // lock for the device stream pool mutable OrtMutex device_stream_pool_mutex_; mutable std::vector> 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 }; diff --git a/onnxruntime/core/framework/stream_execution_context.cc b/onnxruntime/core/framework/stream_execution_context.cc index bdbcdc1f82..560358ec99 100644 --- a/onnxruntime/core/framework/stream_execution_context.cc +++ b/onnxruntime/core/framework/stream_execution_context.cc @@ -13,7 +13,7 @@ StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state, int32_t num_streams, gsl::span notification_owners, size_t num_barriers, - const DeviceStreamCollection& device_stream_map, + const DeviceStreamCollection* device_stream_map, gsl::span feed_mlvalue_idxs, gsl::span feeds, gsl::span fetch_mlvalue_idxs, std::vector& 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({})), 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 diff --git a/onnxruntime/core/framework/stream_execution_context.h b/onnxruntime/core/framework/stream_execution_context.h index 4f0951d9f9..ede5056209 100644 --- a/onnxruntime/core/framework/stream_execution_context.h +++ b/onnxruntime/core/framework/stream_execution_context.h @@ -61,7 +61,7 @@ class StreamExecutionContext { #ifdef ORT_ENABLE_STREAM gsl::span notification_owners, size_t num_barriers, - const DeviceStreamCollection& device_stream_map, + const DeviceStreamCollection* device_stream_map, #endif gsl::span feed_mlvalue_idxs, gsl::span feeds, gsl::span fetch_mlvalue_idxs, @@ -171,7 +171,8 @@ class StreamExecutionContext { #ifdef ORT_ENABLE_STREAM InlinedVector> 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 count_down_barriers_; #endif diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index 102bd15f06..da523f300f 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -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 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 feeds, std::vector& 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);