Add a ParallelForWithWorkerId function.

This allows allocating a partial result for each thread. Writing to the partial result indexed by the worker id is thread safe.
Change: 141209775

1 files changed, 9 insertions, 0 deletions

diff --git a/tensorflow/core/lib/core/threadpool.cc b/tensorflow/core/lib/core/threadpool.cc
index a2245bb28e..59d77a3734 100644
--- a/tensorflow/core/lib/core/threadpool.cc
+++ b/tensorflow/core/lib/core/threadpool.cc
@@ -123,6 +123,15 @@ void ThreadPool::ParallelFor(int64 total, int64 cost_per_unit,
   impl_->ParallelFor(total, cost_per_unit, std::move(fn));
 }
 
+void ThreadPool::ParallelForWithWorkerId(
+    int64 total, int64 cost_per_unit,
+    const std::function<void(int64, int64, int)>& fn) {
+  impl_->ParallelFor(total, cost_per_unit,
+                     [this, &fn](int64 start, int64 limit) {
+                       fn(start, limit, CurrentThreadId());
+                     });
+}
+
 int ThreadPool::NumThreads() const { return impl_->NumThreads(); }
 
 int ThreadPool::CurrentThreadId() const { return impl_->CurrentThreadId(); }