From 1d5af0693c4d54cf15aa9a787d5765ddfaf706dc Mon Sep 17 00:00:00 2001
From: Rasmus Munk Larsen <rmlarsen@google.com>
Date: Thu, 19 Sep 2019 13:54:49 -0700
Subject: Add support for asynchronous evaluation of tensor casting
 expressions.

---
 .../Eigen/CXX11/src/Tensor/TensorConversion.h      | 31 ++++++++++++++++++++++
 1 file changed, 31 insertions(+)

(limited to 'unsupported/Eigen/CXX11/src/Tensor/TensorConversion.h')

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorConversion.h b/unsupported/Eigen/CXX11/src/Tensor/TensorConversion.h
index e96f31537..fa329bfe6 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorConversion.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorConversion.h
@@ -178,6 +178,27 @@ template <typename Eval, typename EvalPointerType> struct ConversionSubExprEval<
   }
 };
 
+#ifdef EIGEN_USE_THREADS
+template <bool SameType, typename Eval, typename EvalPointerType,
+          typename EvalSubExprsCallback>
+struct ConversionSubExprEvalAsync {
+  static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void run(
+      Eval& impl, EvalPointerType, EvalSubExprsCallback done) {
+    impl.evalSubExprsIfNeededAsync(nullptr, std::move(done));
+  }
+};
+
+template <typename Eval, typename EvalPointerType,
+          typename EvalSubExprsCallback>
+struct ConversionSubExprEvalAsync<true, Eval, EvalPointerType,
+                                  EvalSubExprsCallback> {
+  static EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void run(
+      Eval& impl, EvalPointerType data, EvalSubExprsCallback done) {
+    impl.evalSubExprsIfNeededAsync(data, std::move(done));
+  }
+};
+#endif
+
 namespace internal {
 
 template <typename SrcType, typename TargetType, bool IsSameT>
@@ -299,6 +320,16 @@ struct TensorEvaluator<const TensorConversionOp<TargetType, ArgType>, Device>
     return ConversionSubExprEval<IsSameType, TensorEvaluator<ArgType, Device>, EvaluatorPointerType>::run(m_impl, data);
   }
 
+#ifdef EIGEN_USE_THREADS
+  template <typename EvalSubExprsCallback>
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void evalSubExprsIfNeededAsync(
+      EvaluatorPointerType data, EvalSubExprsCallback done) {
+    ConversionSubExprEvalAsync<IsSameType, TensorEvaluator<ArgType, Device>,
+                               EvaluatorPointerType,
+        EvalSubExprsCallback>::run(m_impl, data, std::move(done));
+  }
+#endif
+
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void cleanup()
   {
     m_impl.cleanup();
-- 
cgit v1.2.3