| 54 | |
| 55 | template<typename T> |
| 56 | void CpuTopK(ep::Stream* /*stream*/, const T* in_ptr, int64_t* indices_ptr, int64_t instance_num, |
| 57 | int64_t instance_size, int64_t k, bool sorted, int64_t* out_ptr) { |
| 58 | const int64_t num_thread = |
| 59 | std::min(instance_num, static_cast<int64_t>(Singleton<ThreadPool>::Get()->thread_num())); |
| 60 | const BalancedSplitter bs(instance_num, num_thread); |
| 61 | BlockingCounter bc(num_thread); |
| 62 | FOR_RANGE(int64_t, thread_id, 0, num_thread) { |
| 63 | const Range range = bs.At(thread_id); |
| 64 | Singleton<ThreadPool>::Get()->AddWork([=, &bc]() { |
| 65 | if (k == 1) { |
| 66 | ComputeTopOne(in_ptr, range, instance_size, out_ptr); |
| 67 | } else { |
| 68 | ComputeTopK(in_ptr, indices_ptr, range, instance_size, k, sorted, out_ptr); |
| 69 | } |
| 70 | bc.Decrease(); |
| 71 | }); |
| 72 | } |
| 73 | bc.WaitForeverUntilCntEqualZero(); |
| 74 | } |
| 75 | |
| 76 | } // namespace |
| 77 |
no test coverage detected