Only use TiledArray allocator if available

Signed-off-by: Joseph Schuchart <[email protected]>
devreal · Oct 28, 2024 · d405b7e · d405b7e
1 parent ce80a38
commit d405b7e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/examples/matrixtile.h b/examples/matrixtile.h
@@ -12,7 +12,8 @@
 
 #include <TiledArray/external/device.h>
 #if defined(TILEDARRAY_HAS_DEVICE)
-#define ALLOCATOR TiledArray::device_pinned_allocator<T>
+template<typename T>
+using Allocator = TiledArray::device_pinned_allocator<T>;
 
 inline void allocator_init(int argc, char **argv) {
   // initialize MADNESS so that TA allocators can be created
@@ -26,15 +27,16 @@ inline void allocator_fini() {
   madness::finalize();
 }
 #else  // TILEDARRAY_HAS_DEVICE
-#define ALLOCATOR std::allocator<T>
+template<typename T>
+using Allocator = std::allocator<T>;
 
 inline void allocator_init(int argc, char **argv) { }
 
 inline void allocator_fini() { }
 
 #endif // TILEDARRAY_HAS_DEVICE
 
-template <typename T, class Allocator = ALLOCATOR>
+template <typename T, class Allocator = Allocator<T>>
 class MatrixTile : public ttg::TTValue<MatrixTile<T, Allocator>> {
  public:
   using metadata_t = typename std::tuple<std::size_t, std::size_t, std::size_t>;

diff --git a/examples/potrf/potrf.h b/examples/potrf/potrf.h
@@ -95,7 +95,7 @@ namespace potrf {
                   ttg::Edge<Key2, MatrixTile<typename MatrixT::element_type>>& output_result) {
     using T = typename MatrixT::element_type;
 #if defined(ENABLE_DEVICE_KERNEL)
-    auto iallocator = std::make_shared<TiledArray::device_pinned_allocator<int>>();
+    auto iallocator = std::make_shared<Allocator<int>>();
     //std::cout << "Creating CUDA POTRF task " << std::endl;
     auto f_dev = [=, iallocator = std::move(iallocator)]
                     (const Key1& key, MatrixTile<T>&& tile_kk,