Add ThreadPool::ScheduleToAllThreads

NervanaSystems · Aug 6, 2019 · fd04008 · fd04008
1 parent c94f406
commit fd04008
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/tensorflow/core/lib/core/threadpool.cc b/tensorflow/core/lib/core/threadpool.cc
@@ -204,6 +204,10 @@ void ThreadPool::ScheduleWithHint(std::function<void()> fn, int start,
   impl_->ScheduleWithHint(std::move(fn), start, limit);
 }
 
+void ThreadPool::ScheduleToAllThreads(const std::function<void(int)>& fn) {
+  impl_->ScheduleToAllThreads(fn);
+}
+
 void ThreadPool::SetStealPartitions(
     const std::vector<std::pair<unsigned, unsigned>>& partitions) {
   impl_->SetStealPartitions(partitions);

diff --git a/tensorflow/core/lib/core/threadpool.h b/tensorflow/core/lib/core/threadpool.h
@@ -69,6 +69,7 @@ class ThreadPool {
       const std::vector<std::pair<unsigned, unsigned>>& partitions);
 
   void ScheduleWithHint(std::function<void()> fn, int start, int limit);
+  void ScheduleToAllThreads(const std::function<void(int)>& fn);
   // Requires 0 < block_size <= total.
   // Spawns k threads and calls fn(i*block_size, (i+1)*block_size) from the
   // ith thread (i>=0). When (i+1)*block_size > total, fn(i*block_size, total)