amends 577fda2 to always use the old implementation of ta_tensor_to_u…

…m_tensor when element conversion is needed

src/TiledArray/device/btas_um_tensor.h

@@ -67,7 +67,7 @@ struct ArchiveLoadImpl<Archive, TiledArray::btasUMTensorVarray<T>> {
                           TiledArray::btasUMTensorVarray<T> &t) {
     TiledArray::Range range{};
     TiledArray::device_um_btas_varray<T> store{};
-    ar &range &store;
+    ar & range & store;
     t = TiledArray::btasUMTensorVarray<T>(std::move(range), std::move(store));
     // device::setDevice(TiledArray::deviceEnv::instance()->default_device_id());
     // auto &stream = device::stream_for(range);
@@ -83,7 +83,7 @@ struct ArchiveStoreImpl<Archive, TiledArray::btasUMTensorVarray<T>> {
     auto stream = TiledArray::device::stream_for(t.range());
     TiledArray::to_execution_space<TiledArray::ExecutionSpace::Host>(
         t.storage(), stream);
-    ar &t.range() & t.storage();
+    ar & t.range() & t.storage();
   }
 };
 
@@ -674,25 +674,12 @@ template <typename UMTensor, typename TATensor, typename Policy>
 typename std::enable_if<!std::is_same<UMTensor, TATensor>::value,
                         TiledArray::DistArray<UMTensor, Policy>>::type
 ta_tensor_to_um_tensor(const TiledArray::DistArray<TATensor, Policy> &array) {
-  auto convert_tile_memcpy = [](const TATensor &tile) {
-    /// UMTensor must be wrapped into TA::Tile
-
-    using Tensor = typename UMTensor::tensor_type;
-
-    auto stream = device::stream_for(tile.range());
-    typename Tensor::storage_type storage;
-    make_device_storage(storage, tile.range().area(), stream);
-    Tensor result(tile.range(), std::move(storage));
-
-    DeviceSafeCall(
-        device::memcpyAsync(result.data(), tile.data(),
-                            tile.size() * sizeof(typename Tensor::value_type),
-                            device::MemcpyDefault, stream));
-
-    device::sync_madness_task_with(stream);
-    return TiledArray::Tile<Tensor>(std::move(result));
-  };
+  using inT = typename TATensor::value_type;
+  using outT = typename UMTensor::value_type;
+  // check if element conversion is necessary
+  constexpr bool T_conversion = !std::is_same_v<inT, outT>;
 
+  // this is safe even when need to convert element types, but less efficient
   auto convert_tile_um = [](const TATensor &tile) {
     /// UMTensor must be wrapped into TA::Tile
 
@@ -711,14 +698,46 @@ ta_tensor_to_um_tensor(const TiledArray::DistArray<TATensor, Policy> &array) {
     TiledArray::to_execution_space<TiledArray::ExecutionSpace::Device>(
         result.storage(), stream);
 
+    // N.B. move! without it have D-to-H transfer due to calling UM
+    // allocator construct() on the host
     return TiledArray::Tile<Tensor>(std::move(result));
   };
 
-  const char *use_legacy_conversion =
-      std::getenv("TA_DEVICE_LEGACY_UM_CONVERSION");
-  auto um_array = use_legacy_conversion
-                      ? to_new_tile_type(array, convert_tile_um)
-                      : to_new_tile_type(array, convert_tile_memcpy);
+  if constexpr (T_conversion) {
+    auto um_array = to_new_tile_type(array, convert_tile_um);
+  } else {
+    // this is more efficient for copying:
+    // - avoids copy on host followed by UM transfer, instead uses direct copy
+    // - replaced unneeded copy (which also caused D-to-H transfer due to
+    // calling UM allocator construct() on the host) by move
+    // This eliminates all spurious UM traffic in (T) W3 contractions
+    auto convert_tile_memcpy = [](const TATensor &tile) {
+      /// UMTensor must be wrapped into TA::Tile
+
+      using Tensor = typename UMTensor::tensor_type;
+
+      auto stream = device::stream_for(tile.range());
+      typename Tensor::storage_type storage;
+      make_device_storage(storage, tile.range().area(), stream);
+      Tensor result(tile.range(), std::move(storage));
+
+      DeviceSafeCall(
+          device::memcpyAsync(result.data(), tile.data(),
+                              tile.size() * sizeof(typename Tensor::value_type),
+                              device::MemcpyDefault, stream));
+
+      device::sync_madness_task_with(stream);
+      // N.B. move! without it have D-to-H transfer due to calling UM
+      // allocator construct() on the host
+      return TiledArray::Tile<Tensor>(std::move(result));
+    };
+
+    const char *use_legacy_conversion =
+        std::getenv("TA_DEVICE_LEGACY_UM_CONVERSION");
+    auto um_array = use_legacy_conversion
+                        ? to_new_tile_type(array, convert_tile_um)
+                        : to_new_tile_type(array, convert_tile_memcpy);
+  }
 
   array.world().gop.fence();
   return um_array;