| 194 | } |
| 195 | |
| 196 | void ThreadPool::ParallelForWithWorkerId( |
| 197 | int64 total, int64 cost_per_unit, |
| 198 | const std::function<void(int64, int64, int)>& fn) { |
| 199 | CHECK_GE(total, 0); |
| 200 | CHECK_EQ(total, (int64)(Eigen::Index)total); |
| 201 | |
| 202 | threadpool_device_->parallelFor(total, |
| 203 | Eigen::TensorOpCost(0, 0, cost_per_unit), |
| 204 | [this, &fn](int64 start, int64 limit) { |
| 205 | // ParallelFor may use the current thread to |
| 206 | // do some work synchronously. When calling |
| 207 | // CurrentThreadId() from outside of the |
| 208 | // thread pool, we get -1, so we can shift |
| 209 | // every id up by 1. |
| 210 | int id = CurrentThreadId() + 1; |
| 211 | fn(start, limit, id); |
| 212 | }); |
| 213 | } |
| 214 | |
| 215 | int ThreadPool::NumThreads() const { |
| 216 | return underlying_threadpool_->NumThreads(); |