Proposal for reworking host memory

src/axom/core/ArrayBase.hpp

@@ -894,6 +894,7 @@ struct DeviceStagingBuffer
                       T* data,
                       IndexType begin,
                       IndexType nelems,
+                      HostAllocator host_allocator,
                       bool read_from_data = false)
     : m_data(data)
     , m_begin(begin)
@@ -911,8 +912,7 @@ struct DeviceStagingBuffer
 #if defined(AXOM_USE_CUDA) && defined(AXOM_USE_UMPIRE)
     if(m_deviceStage)
     {
-      int allocator_id = axom::detail::getAllocatorID<axom::MemorySpace::Host>();
-      m_staging_buf = axom::allocate<T>(nelems, allocator_id);
+      m_staging_buf = axom::allocate<T>(nelems, host_allocator.getID());
       if(read_from_data)
       {
         axom::copy(m_staging_buf, m_data + begin, sizeof(T) * nelems);
@@ -921,6 +921,7 @@ struct DeviceStagingBuffer
 #else
     AXOM_UNUSED_VAR(space);
     AXOM_UNUSED_VAR(read_from_data);
+    AXOM_UNUSED_VAR(host_allocator);
 #endif
   }
 
@@ -986,20 +987,34 @@ struct ArrayOps
   using StagingBuffer = DeviceStagingBuffer<T>;
 
 public:
-  ArrayOps(int allocId, bool preferDevice)
+  AXOM_HOST_DEVICE ArrayOps() : m_host_allocator(axom::MALLOC_ALLOCATOR_ID) { }
+
+  AXOM_HOST_DEVICE ArrayOps(int allocId,
+                            bool preferDevice,
+                            HostAllocator hostAllocator = HostAllocator {axom::MALLOC_ALLOCATOR_ID})
+    : m_host_allocator(hostAllocator)
   {
 #if defined(AXOM_USE_GPU) && defined(AXOM_USE_UMPIRE)
+  #if defined(AXOM_DEVICE_CODE)
+    AXOM_UNUSED_VAR(allocId);
+    AXOM_UNUSED_VAR(preferDevice);
+  #else
     space = getAllocatorSpace(allocId);
+    if(space == MemorySpace::Malloc)
+    {
+      space = MemorySpace::Host;
+    }
 
     bool isUnifiedSpace = false;
     isUnifiedSpace = (space == MemorySpace::Unified || space == MemorySpace::Pinned);
-  #if defined(AXOM_USE_HIP)
+    #if defined(AXOM_USE_HIP)
     isUnifiedSpace = (isUnifiedSpace || space == MemorySpace::Device);
-  #endif
+    #endif
     if(!preferDevice && isUnifiedSpace)
     {
       space = MemorySpace::Host;
     }
+  #endif
 #else
     AXOM_UNUSED_VAR(allocId);
     AXOM_UNUSED_VAR(preferDevice);
@@ -1042,7 +1057,7 @@ struct ArrayOps
 #endif
       // Object is neither trivially default-constructible nor trivially-
       // copyable. Construct instances on the host.
-      StagingBuffer tmp_buf(space, data, begin, nelems);
+      StagingBuffer tmp_buf(space, data, begin, nelems, m_host_allocator);
       T* data_host = tmp_buf.getStagingBuffer();
       for(IndexType i = 0; i < nelems; ++i)
       {
@@ -1074,7 +1089,7 @@ struct ArrayOps
 #endif
     // Object is not trivially-copyable, so ensure copy constructors are
     // called on the host.
-    StagingBuffer tmp_buf(space, array, begin, nelems);
+    StagingBuffer tmp_buf(space, array, begin, nelems, m_host_allocator);
     std::uninitialized_fill_n(tmp_buf.getStagingBuffer(), nelems, value);
   }
 
@@ -1097,8 +1112,8 @@ struct ArrayOps
     {
       // HostOp::fill_range will handle the copy to our "staging" host buffer,
       // regardless of the source memory space.
-      StagingBuffer dst_buf(space, array, begin, nelems);
-      DeviceStagingBuffer<T> src_buf(valueSpace, const_cast<T*>(values), 0, nelems, true);
+      StagingBuffer dst_buf(space, array, begin, nelems, m_host_allocator);
+      DeviceStagingBuffer<T> src_buf(valueSpace, const_cast<T*>(values), 0, nelems, m_host_allocator, true);
       std::uninitialized_copy(src_buf.getStagingBuffer(),
                               src_buf.getStagingBuffer() + nelems,
                               dst_buf.getStagingBuffer());
@@ -1136,8 +1151,13 @@ struct ArrayOps
       else
       {
         // Strided case - element-by-element copy
-        StagingBuffer dst_buf(space, array, begin, nelems);
-        DeviceStagingBuffer<T> src_buf(valueSpace, const_cast<T*>(values), 0, nelems * src_stride, true);
+        StagingBuffer dst_buf(space, array, begin, nelems, m_host_allocator);
+        DeviceStagingBuffer<T> src_buf(valueSpace,
+                                       const_cast<T*>(values),
+                                       0,
+                                       nelems * src_stride,
+                                       m_host_allocator,
+                                       true);
 
         T* dst = dst_buf.getStagingBuffer();
         const T* src = src_buf.getStagingBuffer();
@@ -1152,8 +1172,13 @@ struct ArrayOps
     else
     {
       // Non-trivially copyable - use placement new with stride
-      StagingBuffer dst_buf(space, array, begin, nelems);
-      DeviceStagingBuffer<T> src_buf(valueSpace, const_cast<T*>(values), 0, nelems * src_stride, true);
+      StagingBuffer dst_buf(space, array, begin, nelems, m_host_allocator);
+      DeviceStagingBuffer<T> src_buf(valueSpace,
+                                     const_cast<T*>(values),
+                                     0,
+                                     nelems * src_stride,
+                                     m_host_allocator,
+                                     true);
 
       T* dst = dst_buf.getStagingBuffer();
       const T* src = src_buf.getStagingBuffer();
@@ -1206,7 +1231,7 @@ struct ArrayOps
   {
     if constexpr(!std::is_trivially_destructible_v<T>)
     {
-      StagingBuffer tmp_buf(space, array, begin, nelems, true);
+      StagingBuffer tmp_buf(space, array, begin, nelems, m_host_allocator, true);
       T* array_host = tmp_buf.getStagingBuffer();
       for(IndexType i = 0; i < nelems; i++)
       {
@@ -1225,18 +1250,15 @@ struct ArrayOps
    */
   void move(T* array, IndexType src_begin, IndexType src_end, IndexType dst)
   {
-#if defined(AXOM_USE_GPU) && defined(AXOM_USE_UMPIRE)
-  #ifdef AXOM_USE_CUDA
+#if defined(AXOM_USE_GPU) && defined(AXOM_USE_UMPIRE) && defined(AXOM_USE_CUDA)
     // CUDA-only: we require non-trivial types to be trivially-relocatable.
     // This enables us to do simple memcpys for move operations.
     bool presume_trivially_relocatable = (space == MemorySpace::Device);
-  #else
-    constexpr bool presume_trivially_relocatable = false;
-  #endif
-    if(std::is_trivially_copyable_v<T> || presume_trivially_relocatable)
+    if(space == MemorySpace::Device &&
+       (std::is_trivially_copyable_v<T> || presume_trivially_relocatable))
     {
-      // Since this memory is on the device-side, we copy it to a temporary buffer
-      // first.
+      // Device-only CUDA memory cannot be shifted from the host in-place.
+      // Copy it through a temporary device buffer instead.
       IndexType nelems = src_end - src_begin;
       T* tmp_buf = axom::allocate<T>(nelems, axom::execution_space<ExecSpace>::allocatorID());
       axom::copy(tmp_buf, array + src_begin, nelems * sizeof(T));
@@ -1299,6 +1321,9 @@ struct ArrayOps
       destroy(values, 0, nelems);
     }
   }
+
+private:
+  HostAllocator m_host_allocator;
 };
 
 template <typename T, int SliceDim, typename BaseArray>

src/axom/core/ArrayView.hpp

@@ -379,7 +379,7 @@ AXOM_HOST_DEVICE ArrayView<T, DIM, SPACE>::ArrayView(ArrayBase<T, DIM, OtherArra
 #if !defined(AXOM_DEVICE_CODE) && defined(AXOM_DEBUG)
   // If it's not dynamic, the allocator ID from the argument array has to match the template param.
   // If that's not the case then things have gone horribly wrong somewhere.
-  if(SPACE != MemorySpace::Dynamic && SPACE != axom::detail::getAllocatorSpace(m_allocator_id))
+  if(!axom::isAllocatorCompatibleWithMemorySpace(m_allocator_id, SPACE))
   {
     std::cerr << "Input argument allocator does not match the explicitly "
                  "provided memory space\n";
@@ -403,7 +403,7 @@ AXOM_HOST_DEVICE ArrayView<T, DIM, SPACE>::ArrayView(
 #if !defined(AXOM_DEVICE_CODE) && defined(AXOM_DEBUG)
   // If it's not dynamic, the allocator ID from the argument array has to match the template param.
   // If that's not the case then things have gone horribly wrong somewhere.
-  if(SPACE != MemorySpace::Dynamic && SPACE != axom::detail::getAllocatorSpace(m_allocator_id))
+  if(!axom::isAllocatorCompatibleWithMemorySpace(m_allocator_id, SPACE))
   {
     std::cerr << "Input argument allocator does not match the explicitly "
                  "provided memory space\n";

src/axom/core/CMakeLists.txt

@@ -30,6 +30,7 @@ set(core_headers
     utilities/CommandLineUtilities.hpp
     utilities/ConstexprAssert.hpp
     utilities/FileUtilities.hpp
+    utilities/MemoryTesting.hpp
     utilities/RAII.hpp
     utilities/Sorting.hpp
     utilities/StringUtilities.hpp

src/axom/core/FlatMap.hpp

@@ -160,9 +160,11 @@ class FlatMap : detail::flat_map::SequentialLookupPolicy<typename Hash::result_t
                   "Cannot copy an axom::FlatMap when value type is not "
                   "copy-constructible.");
     // Copy all elements.
-    detail::flat_map::copyBuckets<KeyValuePair, LookupPolicy>(m_metadata.view(),
-                                                              other.m_buckets.view(),
-                                                              m_buckets.view());
+    detail::flat_map::copyBuckets<KeyValuePair, LookupPolicy>(
+      m_metadata.view(),
+      other.m_buckets.view(),
+      m_buckets.view(),
+      HostAllocator {m_buckets.getHostAllocatorID()});
   }
 
   /*!
@@ -207,9 +209,11 @@ class FlatMap : detail::flat_map::SequentialLookupPolicy<typename Hash::result_t
     , m_loadCount(other.m_loadCount)
   {
     // Copy all elements.
-    detail::flat_map::copyBuckets<KeyValuePair, LookupPolicy>(m_metadata.view(),
-                                                              other.m_buckets.view(),
-                                                              m_buckets.view());
+    detail::flat_map::copyBuckets<KeyValuePair, LookupPolicy>(
+      m_metadata.view(),
+      other.m_buckets.view(),
+      m_buckets.view(),
+      HostAllocator {m_buckets.getHostAllocatorID()});
   }
 
   /// \brief Destructor for a FlatMap instance.
@@ -218,8 +222,10 @@ class FlatMap : detail::flat_map::SequentialLookupPolicy<typename Hash::result_t
     // Destroy all elements.
     if(m_size > 0)
     {
-      detail::flat_map::destroyBuckets<KeyValuePair, LookupPolicy>(m_metadata.view(),
-                                                                   m_buckets.view());
+      detail::flat_map::destroyBuckets<KeyValuePair, LookupPolicy>(
+        m_metadata.view(),
+        m_buckets.view(),
+        HostAllocator {m_buckets.getHostAllocatorID()});
     }
 
     // Unlike in clear() we don't need to reset metadata here.
@@ -384,7 +390,10 @@ class FlatMap : detail::flat_map::SequentialLookupPolicy<typename Hash::result_t
   void clear()
   {
     // Destroy all elements.
-    detail::flat_map::destroyBuckets<KeyValuePair, LookupPolicy>(m_metadata.view(), m_buckets.view());
+    detail::flat_map::destroyBuckets<KeyValuePair, LookupPolicy>(
+      m_metadata.view(),
+      m_buckets.view(),
+      HostAllocator {m_buckets.getHostAllocatorID()});
 
     // Also reset metadata.
     IndexType numGroupsRounded = IndexType {1} << m_numGroups2;

src/axom/core/detail/FlatMapOps.hpp

@@ -34,17 +34,23 @@ inline void setSentinel(axom::ArrayView<GroupBucket> metadata)
 }
 
 template <typename KVPair, typename LookupPolicy, typename StoragePair = TypeErasedStorage<KVPair>>
-inline void destroyBuckets(axom::ArrayView<GroupBucket> metadata, axom::ArrayView<StoragePair> buckets)
+inline void destroyBuckets(axom::ArrayView<GroupBucket> metadata,
+                           axom::ArrayView<StoragePair> buckets,
+                           HostAllocator host_allocator)
 {
+#if !defined(AXOM_USE_UMPIRE) || !defined(AXOM_USE_CUDA)
+  // Note: HIP can access device memory from the host and does not need special
+  // handling - we just defer to the host path in all cases.
+  AXOM_UNUSED_VAR(host_allocator);
+#endif
+
   if(std::is_trivially_destructible<KVPair>::value)
   {
     // Nothing to do.
     return;
   }
 
 #if defined(AXOM_USE_UMPIRE) && defined(AXOM_USE_CUDA)
-  // Note: HIP can access device memory from the host and does not need special
-  // handling - we just defer to the host path in all cases.
   MemorySpace space = getAllocatorSpace(metadata.getAllocatorID());
   // CUDA-only: buckets located in device-only memory and non-trivially
   // destructible. We'll need to "relocate" the objects to the host to
@@ -56,9 +62,8 @@ inline void destroyBuckets(axom::ArrayView<GroupBucket> metadata, axom::ArrayVie
   axom::Array<GroupBucket> metadata_host;
   if(space == MemorySpace::Device)
   {
-    int host_allocator_id = axom::execution_space<axom::SEQ_EXEC>::allocatorID();
-    metadata_host = axom::Array<GroupBucket>(metadata, host_allocator_id);
-    buckets_host = axom::Array<StoragePair>(buckets, host_allocator_id);
+    metadata_host = axom::Array<GroupBucket>(metadata, host_allocator.getID(), host_allocator);
+    buckets_host = axom::Array<StoragePair>(buckets, host_allocator.getID(), host_allocator);
     metadata = metadata_host.view();
     buckets = buckets_host.view();
   }
@@ -74,8 +79,15 @@ inline void destroyBuckets(axom::ArrayView<GroupBucket> metadata, axom::ArrayVie
 template <typename KVPair, typename LookupPolicy, typename StoragePair = TypeErasedStorage<KVPair>>
 inline void copyBuckets(axom::ArrayView<const GroupBucket> metadata,
                         axom::ArrayView<const StoragePair> from_buckets,
-                        axom::ArrayView<StoragePair> to_buckets)
+                        axom::ArrayView<StoragePair> to_buckets,
+                        HostAllocator host_allocator)
 {
+#if !defined(AXOM_USE_UMPIRE) || !defined(AXOM_USE_CUDA)
+  // Note: HIP can access device memory from the host and does not need special
+  // handling - we just defer to the host path in all cases.
+  AXOM_UNUSED_VAR(host_allocator);
+#endif
+
   if(std::is_trivially_copyable<KVPair>::value)
   {
     axom::copy(to_buckets.data(), from_buckets.data(), sizeof(StoragePair) * from_buckets.size());
@@ -84,34 +96,29 @@ inline void copyBuckets(axom::ArrayView<const GroupBucket> metadata,
 
   axom::ArrayView<StoragePair> to_buckets_stage = to_buckets;
 #if defined(AXOM_USE_UMPIRE) && defined(AXOM_USE_CUDA)
-  // Note: HIP can access device memory from the host and does not need special
-  // handling - we just defer to the host path in all cases.
-
   // Non-trivially copyable:
   // "Relocate" to the host to call host-based copy constructor.
   MemorySpace meta_space = getAllocatorSpace(metadata.getAllocatorID());
   axom::Array<GroupBucket> metadata_host;
   if(meta_space == MemorySpace::Device)
   {
-    int host_allocator_id = axom::execution_space<axom::SEQ_EXEC>::allocatorID();
-    metadata_host = axom::Array<GroupBucket>(metadata, host_allocator_id);
+    metadata_host = axom::Array<GroupBucket>(metadata, host_allocator.getID(), host_allocator);
     metadata = metadata_host.view();
   }
   MemorySpace from_space = getAllocatorSpace(from_buckets.getAllocatorID());
   axom::Array<StoragePair> from_buckets_host;
   if(from_space == MemorySpace::Device)
   {
-    int host_allocator_id = axom::execution_space<axom::SEQ_EXEC>::allocatorID();
-    from_buckets_host = axom::Array<StoragePair>(from_buckets, host_allocator_id);
+    from_buckets_host =
+      axom::Array<StoragePair>(from_buckets, host_allocator.getID(), host_allocator);
     from_buckets = from_buckets_host.view();
   }
 
   MemorySpace to_space = getAllocatorSpace(to_buckets.getAllocatorID());
   axom::Array<StoragePair> to_buckets_host;
   if(to_space == MemorySpace::Device)
   {
-    int host_allocator_id = axom::execution_space<axom::SEQ_EXEC>::allocatorID();
-    to_buckets_host = axom::Array<StoragePair>(to_buckets, host_allocator_id);
+    to_buckets_host = axom::Array<StoragePair>(to_buckets, host_allocator.getID(), host_allocator);
     to_buckets_stage = to_buckets_host.view();
   }
 #endif