diff options
author | Paul Tucker <tucker@google.com> | 2018-07-16 17:26:05 -0700 |
---|---|---|
committer | Paul Tucker <tucker@google.com> | 2018-07-16 17:26:05 -0700 |
commit | b3e7c9132d41da75c0a6af783300cb11101010db (patch) | |
tree | 66336060cd3f70438bd3933ee54a691f6b23bf98 /unsupported/Eigen/CXX11/src/Tensor | |
parent | 40797dbea334405fb49ea1ba2cd972a23a7dbe02 (diff) |
Add optional Allocator argument to ThreadPoolDevice constructor.
When supplied, this allocator will be used in place of
internal::aligned_malloc. This permits e.g. use of a NUMA-node specific
allocator where the thread-pool is also restricted a single NUMA-node.
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor')
-rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h | 16 |
1 files changed, 13 insertions, 3 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h index 90fd99027..be397e1b6 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h @@ -95,14 +95,20 @@ static EIGEN_STRONG_INLINE void wait_until_ready(SyncType* n) { // 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) + : pool_(pool), num_threads_(num_cores), allocator_(nullptr) { } 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 memcpy(void* dst, const void* src, size_t n) const { @@ -267,9 +273,13 @@ struct ThreadPoolDevice { // Thread pool accessor. ThreadPoolInterface* getPool() const { return pool_; } + // Allocator accessor. + Allocator* getAllocator() const { return allocator_; } + private: ThreadPoolInterface* pool_; int num_threads_; + Allocator* allocator_; }; |