diff options
author | Eugene Zhulenev <ezhulenev@google.com> | 2019-08-30 14:49:40 -0700 |
---|---|---|
committer | Eugene Zhulenev <ezhulenev@google.com> | 2019-08-30 14:49:40 -0700 |
commit | 66665e7e76d2ad5aa37775b3777e9a53c6d1c18c (patch) | |
tree | cb62a23e970d9125475abd95e4c9e68a02a04461 /unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h | |
parent | f6c51d9209ccc04d28c39f4c8059e7d3e74d6e07 (diff) |
Asynchronous expression evaluation with TensorAsyncDevice
Diffstat (limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h')
-rw-r--r-- | unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h | 41 |
1 files changed, 41 insertions, 0 deletions
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h b/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h index 29e50a3b2..5122b3623 100644 --- a/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h +++ b/unsupported/Eigen/CXX11/src/Tensor/TensorDevice.h @@ -63,6 +63,47 @@ template <typename ExpressionType, typename DeviceType> class TensorDevice { ExpressionType& m_expression; }; +#ifdef EIGEN_USE_THREADS + +/** \class TensorAsyncDevice + * \ingroup CXX11_Tensor_Module + * + * \brief Pseudo expression providing an operator = that will evaluate its + * argument asynchronously on the specified device (currently supports only + * ThreadPoolDevice). + * + * Example: + * std::function<void()> done = []() {}; + * C.device(EIGEN_THREAD_POOL, std::move(done)) = A + B; + */ + +template <typename ExpressionType, typename DeviceType> +class TensorAsyncDevice { + public: + TensorAsyncDevice(const DeviceType& device, ExpressionType& expression, + std::function<void()> done) + : m_device(device), m_expression(expression), m_done(std::move(done)) {} + + template <typename OtherDerived> + EIGEN_STRONG_INLINE TensorAsyncDevice& operator=(const OtherDerived& other) { + typedef TensorAssignOp<ExpressionType, const OtherDerived> Assign; + typedef internal::TensorAsyncExecutor<const Assign, DeviceType> Executor; + + // WARNING: After assignment 'm_done' callback will be in undefined state. + Assign assign(m_expression, other); + Executor::runAsync(assign, m_device, std::move(m_done)); + + return *this; + } + + protected: + const DeviceType& m_device; + ExpressionType& m_expression; + std::function<void()> m_done; +}; + +#endif // EIGEN_USE_THREADS + } // end namespace Eigen #endif // EIGEN_CXX11_TENSOR_TENSOR_DEVICE_H |