GitHub-Proxy/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2025-05-03 01:04:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add optional Allocator argument to ThreadPoolDevice constructor.

Browse Source 
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.
This commit is contained in:
Paul Tucker 2018-07-16 17:26:05 -07:00
parent 40797dbea3
commit b3e7c9132d
1 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
unsupported/Eigen/CXX11/src/Tensor/TensorDeviceThreadPool.h
View File

@ -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_;
};