clone(DistArray) supports device-based arrays

ValeevGroup · Sep 28, 2023 · 4b79b5a · 4b79b5a
1 parent 089dcc6
commit 4b79b5a
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 10 deletions.
diff --git a/src/TiledArray/conversions/clone.h b/src/TiledArray/conversions/clone.h
@@ -26,6 +26,10 @@
 #ifndef TILEDARRAY_CONVERSIONS_CLONE_H__INCLUDED
 #define TILEDARRAY_CONVERSIONS_CLONE_H__INCLUDED
 
+#ifdef TILEDARRAY_HAS_DEVICE
+#include "TiledArray/device/device_task_fn.h"
+#endif
+
 namespace TiledArray {
 
 /// Forward declarations
@@ -53,12 +57,28 @@ inline DistArray<Tile, Policy> clone(const DistArray<Tile, Policy>& arg) {
     if (arg.is_zero(index)) continue;
 
     // Spawn a task to clone the tiles
-    Future<value_type> tile = world.taskq.add(
-        [](const value_type& tile) -> value_type {
-          using TiledArray::clone;
-          return clone(tile);
-        },
-        arg.find(index));
+
+    Future<value_type> tile;
+    if constexpr (!detail::is_device_tile_v<value_type>) {
+      tile = world.taskq.add(
+          [](const value_type& tile) -> value_type {
+            using TiledArray::clone;
+            return clone(tile);
+          },
+          arg.find(index));
+    } else {
+#ifdef TILEDARRAY_HAS_DEVICE
+      tile = madness::add_device_task(
+          world,
+          [](const value_type& tile) -> value_type {
+            using TiledArray::clone;
+            return clone(tile);
+          },
+          arg.find(index));
+#else
+      abort();  // unreachable
+#endif
+    }
 
     // Store result tile
     result.set(index, tile);

diff --git a/src/TiledArray/external/device.h b/src/TiledArray/external/device.h
@@ -807,6 +807,7 @@ inline std::optional<Stream>*& madness_task_stream_opt_ptr_accessor() {
 
 inline std::optional<Stream>& tls_stream_opt_accessor() {
   static thread_local std::optional<Stream> stream_opt =
+
       device::Env::instance()->stream(0);
   return stream_opt;
 }

diff --git a/src/TiledArray/tensor/type_traits.h b/src/TiledArray/tensor/type_traits.h
@@ -314,8 +314,7 @@ template <typename ReduceOp, typename Result, typename... Args>
 constexpr const bool is_reduce_op_v =
     is_reduce_op_<void, ReduceOp, Result, Args...>::value;
 
-/// detect cuda tile
-#ifdef TILEDARRAY_HAS_DEVICE
+/// detect device tile types
 template <typename T>
 struct is_device_tile : public std::false_type {};
 
@@ -329,8 +328,6 @@ struct is_device_tile<LazyArrayTile<T, Op>>
 template <typename T>
 static constexpr const auto is_device_tile_v = is_device_tile<T>::value;
 
-#endif
-
 template <typename Tensor, typename Enabler = void>
 struct default_permutation;