diff --git a/src/inference/src/threading/ie_cpu_streams_executor.cpp b/src/inference/src/threading/ie_cpu_streams_executor.cpp index 2e786599a74..146e93687be 100644 --- a/src/inference/src/threading/ie_cpu_streams_executor.cpp +++ b/src/inference/src/threading/ie_cpu_streams_executor.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -231,11 +232,117 @@ struct CPUStreamsExecutor::Impl { #endif }; + // if the thread is created by CPUStreamsExecutor, the Impl::Stream of the thread is stored by tbb Class + // enumerable_thread_specific, the alias is ThreadLocal, the limitations of ThreadLocal please refer to + // https://spec.oneapi.io/versions/latest/elements/oneTBB/source/thread_local_storage/enumerable_thread_specific_cls.html + // if the thread is created by customer, the Impl::Stream of the thread will be stored in variable _stream_map, and + // will be counted by thread_local t_stream_count_map. + // when the customer's thread is destoryed, the stream's count will became 1, + // Call local() will reuse one of them, and release others. + // it's only a workaround for ticket CVS-111490, please be carefully when need to modify + // CustomeThreadLocal::local(), especially like operations that will affect the count of + // CustomThreadLocal::ThreadId + class CustomThreadLocal : public ThreadLocal> { + class ThreadTracker { + public: + explicit ThreadTracker(const std::thread::id& id) + : _id(id), + _count_ptr(std::make_shared(1)) {} + ~ThreadTracker() { + _count_ptr->fetch_sub(1); + } + std::shared_ptr fetch() { + auto new_ptr = std::shared_ptr(new ThreadTracker(*this)); + auto pre_valule = new_ptr.get()->_count_ptr->fetch_add(1); + OPENVINO_ASSERT(pre_valule == 1, "this value must be 1, please check code CustomThreadLocal::local()"); + return new_ptr; + } + const std::thread::id& get_id() const { + return _id; + } + int count() const { + return *(_count_ptr.get()); + } + + private: + // disable all copy and move semantics, user only can use fetch() + // to create a new instance with a shared count num; + ThreadTracker(ThreadTracker const&) = default; + ThreadTracker(ThreadTracker&&) = delete; + ThreadTracker& operator=(ThreadTracker const&) = delete; + ThreadTracker& operator=(ThreadTracker&&) = delete; + std::thread::id _id; + std::shared_ptr _count_ptr; + }; + + public: + CustomThreadLocal(std::function()> callback_construct, Impl* impl) + : ThreadLocal>(callback_construct), + _impl(impl) {} + std::shared_ptr local() { + // maybe there are two CPUStreamsExecutors in the same thread. + static thread_local std::map> t_stream_count_map; + // fix the memory leak issue that CPUStreamsExecutor is already released, + // but still exists CustomThreadLocal::ThreadTracker in t_stream_count_map + for (auto it = t_stream_count_map.begin(); it != t_stream_count_map.end();) { + if (this != it->first && it->second->count() == 1) { + t_stream_count_map.erase(it++); + } else { + it++; + } + } + auto id = std::this_thread::get_id(); + auto search = _thread_ids.find(id); + if (search != _thread_ids.end()) { + return ThreadLocal>::local(); + } + std::lock_guard guard(_stream_map_mutex); + for (auto& item : _stream_map) { + if (item.first->get_id() == id) { + return item.second; + } + } + std::shared_ptr stream = nullptr; + for (auto it = _stream_map.begin(); it != _stream_map.end();) { + if (it->first->count() == 1) { + if (stream == nullptr) { + stream = it->second; + } + _stream_map.erase(it++); + } else { + it++; + } + } + if (stream == nullptr) { + stream = std::make_shared(_impl); + } + auto tracker_ptr = std::make_shared(id); + t_stream_count_map[(void*)this] = tracker_ptr; + auto new_tracker_ptr = tracker_ptr->fetch(); + _stream_map[new_tracker_ptr] = stream; + return stream; + } + + void set_thread_ids_map(std::vector& threads) { + for (auto& thread : threads) { + _thread_ids.insert(thread.get_id()); + } + } + + private: + std::set _thread_ids; + Impl* _impl; + std::map, std::shared_ptr> _stream_map; + std::mutex _stream_map_mutex; + }; + explicit Impl(const Config& config) : _config{config}, - _streams([this] { - return std::make_shared(this); - }) { + _streams( + [this] { + return std::make_shared(this); + }, + this) { _exectorMgr = executorManager(); auto numaNodes = getAvailableNUMANodes(); if (_config._streams != 0) { @@ -296,6 +403,7 @@ struct CPUStreamsExecutor::Impl { } }); } + _streams.set_thread_ids_map(_threads); } void Enqueue(Task task) { @@ -345,7 +453,7 @@ struct CPUStreamsExecutor::Impl { std::queue _taskQueue; bool _isStopped = false; std::vector _usedNumaNodes; - ThreadLocal> _streams; + CustomThreadLocal _streams; #if (IE_THREAD == IE_THREAD_TBB || IE_THREAD == IE_THREAD_TBB_AUTO) // stream id mapping to the core type // stored in the reversed order (so the big cores, with the highest core_type_id value, are populated first)