| 60 | } |
| 61 | } |
| 62 | void ThreadPool::add_task(const TaskElem& task_elem) { |
| 63 | //! Make sure the main thread have bind |
| 64 | if (m_main_affinity_flag && m_core_binding_function != nullptr) { |
| 65 | std::lock_guard<std::mutex> lock(m_mutex_task); |
| 66 | m_core_binding_function(m_nr_threads - 1); |
| 67 | m_main_affinity_flag = false; |
| 68 | } |
| 69 | size_t parallelism = task_elem.nr_parallelism; |
| 70 | //! If only one thread or one task, execute directly |
| 71 | if (task_elem.nr_parallelism == 1 || m_nr_threads == 1) { |
| 72 | for (size_t i = 0; i < parallelism; i++) { |
| 73 | task_elem.task(i, 0); |
| 74 | } |
| 75 | return; |
| 76 | } else { |
| 77 | std::lock_guard<std::mutex> lock(m_mutex_task); |
| 78 | mgb_assert( |
| 79 | m_task_iter.load(std::memory_order_acquire) <= 0, |
| 80 | "The init value of m_all_sub_task is not zero."); |
| 81 | active(); |
| 82 | //! Set the task number, task iter and task |
| 83 | m_nr_parallelism = parallelism; |
| 84 | m_task_iter.exchange(parallelism, std::memory_order_relaxed); |
| 85 | m_task = [&task_elem](size_t index, size_t thread_id) { |
| 86 | task_elem.task(index, thread_id); |
| 87 | }; |
| 88 | //! Set flag to start thread working |
| 89 | for (uint32_t i = 0; i < m_nr_threads - 1; i++) { |
| 90 | m_workers[i]->work_flag = true; |
| 91 | } |
| 92 | //! Main thread working |
| 93 | int index = -1; |
| 94 | while ((index = m_task_iter.fetch_sub(1, std::memory_order_acq_rel)) && |
| 95 | (index > 0)) { |
| 96 | m_task(static_cast<size_t>(m_nr_parallelism - index), m_nr_threads - 1); |
| 97 | } |
| 98 | //! make sure all threads done |
| 99 | sync(); |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | void ThreadPool::set_affinity(AffinityCallBack affinity_cb) { |
| 104 | mgb_assert(affinity_cb, "The affinity callback must not be nullptr"); |