Asynchronous expression evaluation with TensorAsyncDevice

1 files changed, 11 insertions, 0 deletions

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
index dbacf494e..095c85dc4 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorBase.h
@@ -1063,6 +1063,17 @@ class TensorBase : public TensorBase<Derived, ReadOnlyAccessors> {
       return TensorDevice<Derived, DeviceType>(dev, derived());
     }
 
+#ifdef EIGEN_USE_THREADS
+    // Select the async device on which to evaluate the expression.
+    template <typename DeviceType>
+    typename internal::enable_if<
+        internal::is_same<DeviceType, ThreadPoolDevice>::value,
+        TensorAsyncDevice<Derived, DeviceType>>::type
+    device(const DeviceType& dev, std::function<void()> done) {
+      return TensorAsyncDevice<Derived, DeviceType>(dev, derived(), std::move(done));
+    }
+#endif  // EIGEN_USE_THREADS
+
  protected:
     EIGEN_DEVICE_FUNC
     EIGEN_STRONG_INLINE Derived& derived() { return *static_cast<Derived*>(this); }