diff options
| -rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorDeviceType.h | 22 |
1 files changed, 19 insertions, 3 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceType.h b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceType.h index bb05e4177..efd207507 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceType.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorDeviceType.h @@ -43,11 +43,14 @@ typedef std::promise<void> Promise; static EIGEN_STRONG_INLINE void wait_until_ready(const Future* f) { f->wait(); - // eigen_assert(f->ready()); } +static EIGEN_STRONG_INLINE void get_when_ready(Future* f) { + f->get(); +} + struct ThreadPoolDevice { - ThreadPoolDevice(/*ThreadPool* pool, */size_t num_cores) : num_threads_(num_cores) { } + ThreadPoolDevice(size_t num_cores) : num_threads_(num_cores) { } EIGEN_STRONG_INLINE void* allocate(size_t num_bytes) const { return internal::aligned_malloc(num_bytes); @@ -79,9 +82,9 @@ struct ThreadPoolDevice { } private: - // todo: NUMA, ... size_t num_threads_; }; + #endif @@ -114,6 +117,10 @@ static inline int sharedMemPerBlock() { return m_deviceProperties.sharedMemPerBlock; } +static inline void setCudaSharedMemConfig(cudaSharedMemConfig config) { + cudaError_t status = cudaDeviceSetSharedMemConfig(config); + assert(status == cudaSuccess); +} struct GpuDevice { // The cudastream is not owned: the caller is responsible for its initialization and eventual destruction. @@ -163,10 +170,19 @@ struct GpuDevice { return 32; } + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void synchronize() const { + cudaStreamSynchronize(*stream_); + } + private: // TODO: multigpu. const cudaStream_t* stream_; }; + +#define LAUNCH_CUDA_KERNEL(kernel, gridsize, blocksize, sharedmem, device, ...) \ + (kernel) <<< (gridsize), (blocksize), (sharedmem), (device).stream() >>> (__VA_ARGS__); \ + assert(cudaGetLastError() == cudaSuccess); + #endif } // end namespace Eigen |