diff options
author | Benoit Steiner <benoit.steiner.goog@gmail.com> | 2018-08-01 19:14:34 +0000 |
---|---|---|
committer | Benoit Steiner <benoit.steiner.goog@gmail.com> | 2018-08-01 19:14:34 +0000 |
commit | 93b9e36e10a813cd92ffd397ffb0c1fcf60ca158 (patch) | |
tree | b1f7b47f815e20102a92a999dd7acce6e9d1a283 /unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h | |
parent | 17221115c9f7e382c84c5d053f885470e904f4a4 (diff) | |
parent | 385f7b8d0ca926d00c71987ab308202511e5c753 (diff) |
Merged in paultucker/eigen (pull request PR-431)
Optional ThreadPoolDevice allocator
Approved-by: Benoit Steiner <benoit.steiner.goog@gmail.com>
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h')
-rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h | 23 |
1 files changed, 20 insertions, 3 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h index 5a16ebe50..cc134228a 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h @@ -91,18 +91,31 @@ static EIGEN_STRONG_INLINE void wait_until_ready(SyncType* n) { } } +// An abstract interface to a device specific memory allocator. +class Allocator { + public: + virtual ~Allocator() {} + EIGEN_DEVICE_FUNC virtual void* allocate(size_t num_bytes) const = 0; + EIGEN_DEVICE_FUNC virtual void deallocate(void* buffer) const = 0; +}; // Build a thread pool device on top the an existing pool of threads. struct ThreadPoolDevice { // The ownership of the thread pool remains with the caller. - ThreadPoolDevice(ThreadPoolInterface* pool, int num_cores) : pool_(pool), num_threads_(num_cores) { } + ThreadPoolDevice(ThreadPoolInterface* pool, int num_cores, Allocator* allocator = nullptr) + : pool_(pool), num_threads_(num_cores), allocator_(allocator) { } EIGEN_STRONG_INLINE void* allocate(size_t num_bytes) const { - return internal::aligned_malloc(num_bytes); + return allocator_ ? allocator_->allocate(num_bytes) + : internal::aligned_malloc(num_bytes); } EIGEN_STRONG_INLINE void deallocate(void* buffer) const { - internal::aligned_free(buffer); + if (allocator_) { + allocator_->deallocate(buffer); + } else { + internal::aligned_free(buffer); + } } EIGEN_STRONG_INLINE void* allocate_temp(size_t num_bytes) const { @@ -275,9 +288,13 @@ struct ThreadPoolDevice { // Thread pool accessor. ThreadPoolInterface* getPool() const { return pool_; } + // Allocator accessor. + Allocator* allocator() const { return allocator_; } + private: ThreadPoolInterface* pool_; int num_threads_; + Allocator* allocator_; }; |