Add CurrentThreadId and NumThreads methods to Eigen threadpools and TensorDeviceThreadPool.

This commit is contained in:
Rasmus Munk Larsen
2016-06-03 16:28:58 -07:00
parent 8d97ba6b22
commit 76308e7fd2
5 changed files with 65 additions and 9 deletions

View File

@@ -18,6 +18,13 @@ class ThreadPoolInterface {
public:
virtual void Schedule(std::function<void()> fn) = 0;
// Returns the number of threads in the pool.
virtual size_t NumThreads() const = 0;
// Returns a logical thread index between 0 and NumThreads() - 1 if called
// from one of the threads in the pool. Returns NumThreads() otherwise.
virtual size_t CurrentThreadId() const = 0;
virtual ~ThreadPoolInterface() {}
};