aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/common_runtime/eigen_thread_pool.h
blob: 2554f3521b80f5c26c7f27649c5205e565f9e244 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#ifndef TENSORFLOW_COMMON_RUNTIME_EIGEN_THREAD_POOL_H_
#define TENSORFLOW_COMMON_RUNTIME_EIGEN_THREAD_POOL_H_

#include "tensorflow/core/lib/core/threadpool.h"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"

namespace tensorflow {

class EigenThreadPoolWrapper : public Eigen::ThreadPoolInterface {
 public:
  explicit EigenThreadPoolWrapper(thread::ThreadPool* pool) : pool_(pool) {}
  ~EigenThreadPoolWrapper() override {}

  void Schedule(std::function<void()> fn) override { pool_->Schedule(fn); }

 private:
  thread::ThreadPool* pool_ = nullptr;
};

}  // namespace tensorflow

#endif  // TENSORFLOW_COMMON_RUNTIME_EIGEN_THREAD_POOL_H_