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);