1 /** 2 * Copyright 2020-2022 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_INCLUDE_COMMON_THREAD_POOL_H_ 18 #define MINDSPORE_CCSRC_INCLUDE_COMMON_THREAD_POOL_H_ 19 20 #include <mutex> 21 #include <condition_variable> 22 #include <thread> 23 #include <vector> 24 #include <queue> 25 #include <string> 26 #include <atomic> 27 #include <memory> 28 #include <utility> 29 #include <functional> 30 #include <iostream> 31 #include "utils/log_adapter.h" 32 #include "include/common/visible.h" 33 34 #ifdef PARALLEL_INFERENCE 35 #define BIND_CORE 36 #endif 37 #ifdef __ANDROID__ 38 #define BIND_CORE 39 #include <sched.h> 40 #endif 41 #ifdef __linux__ 42 #define BIND_CORE 43 #endif 44 45 namespace mindspore { 46 namespace common { 47 enum Status { FAIL = -1, SUCCESS = 0 }; 48 using Task = std::function<Status()>; 49 50 struct ThreadContext { 51 std::mutex mutex; 52 std::unique_ptr<std::condition_variable> cond_var{nullptr}; 53 const Task *task{nullptr}; 54 ThreadContextThreadContext55 ThreadContext() { 56 cond_var = std::make_unique<std::condition_variable>(); 57 MS_EXCEPTION_IF_NULL(cond_var); 58 } 59 }; 60 61 class COMMON_EXPORT ThreadPool { 62 public: 63 ~ThreadPool(); 64 ThreadPool(const ThreadPool &) = delete; 65 ThreadPool &operator=(const ThreadPool &) = delete; 66 static ThreadPool &GetInstance(); 67 bool SyncRun(const std::vector<Task> &tasks, const std::vector<int> &core_list = {}); GetSyncRunThreadNum()68 size_t GetSyncRunThreadNum() const { return max_thread_num_; } 69 void ClearThreadPool(); 70 void ChildAfterFork(); 71 72 private: 73 ThreadPool(); 74 void SyncRunLoop(const std::shared_ptr<ThreadContext> &context); 75 bool SetCpuAffinity(const std::vector<int> &core_list); 76 bool FreeScheduleThreads(const std::vector<int> &core_list); 77 #ifdef _WIN32 78 bool ThreadPool::SetAffinity() const; 79 #elif defined(BIND_CORE) 80 bool SetAffinity(const pthread_t &thread_id, cpu_set_t *cpu_set); 81 #endif 82 83 size_t max_thread_num_{1}; 84 std::mutex pool_mtx_; 85 std::atomic_bool exit_run_ = {false}; 86 std::vector<std::unique_ptr<std::thread>> sync_run_threads_{}; 87 std::vector<std::shared_ptr<ThreadContext>> contexts_; 88 }; 89 } // namespace common 90 } // namespace mindspore 91 92 #endif // MINDSPORE_CCSRC_INCLUDE_COMMON_THREAD_POOL_H_ 93