Merge pull request #470 from ValeevGroup/evaleev/fixup/misc-frontier

misc mixups for Frontier
ValeevGroup · Aug 30, 2024 · 236af70 · 236af70
2 parents c10f8d5 + 9bda8ea
commit 236af70
Show file tree

Hide file tree

Showing 6 changed files with 39 additions and 5 deletions.
diff --git a/src/TiledArray/conversions/vector_of_arrays.h b/src/TiledArray/conversions/vector_of_arrays.h
@@ -5,6 +5,8 @@
 #ifndef TILEDARRAY_CONVERSIONS_VECTOR_OF_ARRAYS_H_
 #define TILEDARRAY_CONVERSIONS_VECTOR_OF_ARRAYS_H_
 
+#include <TiledArray/dist_array.h>
+
 namespace TiledArray {
 
 namespace detail {

diff --git a/src/TiledArray/device/device_task_fn.h b/src/TiledArray/device/device_task_fn.h
@@ -121,7 +121,8 @@ struct deviceTaskFn : public TaskInterface {
       } else {
         // TODO should we use device callback or device events??
         // insert device callback
-        TiledArray::device::launchHostFunc(*stream_, device_callback, task_);
+        DeviceSafeCall(TiledArray::device::launchHostFunc(
+            *stream_, device_callback, task_));
         // processed sync, clear state
         stream_ = {};
       }

diff --git a/src/TiledArray/device/thrust.h b/src/TiledArray/device/thrust.h
@@ -32,6 +32,15 @@
 #include <cuda_runtime_api.h>
 #endif
 
+// rocthrust headers rely on THRUST_DEVICE_SYSTEM being defined, which is only
+// defined by the HIP-specific compilers to be usable with host compiler define
+// it here explicitly
+#ifdef TILEDARRAY_HAS_HIP
+#ifndef THRUST_DEVICE_SYSTEM
+#define THRUST_DEVICE_SYSTEM THRUST_DEVICE_SYSTEM_HIP
+#endif
+#endif
+
 #include <thrust/device_vector.h>
 #include <thrust/host_vector.h>
 

diff --git a/src/TiledArray/tile.h b/src/TiledArray/tile.h
@@ -215,11 +215,23 @@ class Tile {
 
   // Dimension information accessors -----------------------------------------
 
-  /// Size accessors
+  /// Size accessor
 
   /// \return The number of elements in the tensor
   decltype(auto) size() const { return tensor().size(); }
 
+  /// Total size accessor
+
+  /// \return The number of elements in the tensor, tallied across batches (if
+  /// any)
+  decltype(auto) total_size() const {
+    if constexpr (detail::has_member_function_total_size_anyreturn_v<
+                      tensor_type>) {
+      return tensor().total_size();
+    } else
+      return size();
+  }
+
   /// Range accessor
 
   /// \return An object describes the upper and lower bounds of  the tensor data

diff --git a/src/TiledArray/type_traits.h b/src/TiledArray/type_traits.h
@@ -322,6 +322,8 @@ GENERATE_HAS_MEMBER_TYPE(mapped_type)
 
 GENERATE_HAS_MEMBER_FUNCTION_ANYRETURN(size)
 GENERATE_HAS_MEMBER_FUNCTION(size)
+GENERATE_HAS_MEMBER_FUNCTION_ANYRETURN(total_size)
+GENERATE_HAS_MEMBER_FUNCTION(total_size)
 GENERATE_HAS_MEMBER_FUNCTION_ANYRETURN(data)
 GENERATE_HAS_MEMBER_FUNCTION(data)
 GENERATE_HAS_MEMBER_FUNCTION_ANYRETURN(empty)

diff --git a/tests/conversions.cpp b/tests/conversions.cpp
@@ -23,13 +23,21 @@
  *
  */
 
-#include "range_fixture.h"
-#include "tiledarray.h"
 #include "unit_test_config.h"
 
-#include "TiledArray/conversions/concat.h"
 #include "TiledArray/conversions/vector_of_arrays.h"
 
+#include "TiledArray/conversions/concat.h"
+
+#include "TiledArray/conversions/dense_to_sparse.h"
+#include "TiledArray/conversions/make_array.h"
+#include "TiledArray/conversions/sparse_to_dense.h"
+#include "TiledArray/conversions/to_new_tile_type.h"
+
+#include "TiledArray/expressions/tsr_expr.h"
+
+#include "range_fixture.h"
+
 using namespace TiledArray;
 
 struct ConversionsFixture : public TiledRangeFixture {