Fix Dynamic memory space when configured with Umpire

src/axom/core/memory_management.hpp

@@ -9,6 +9,7 @@
 // Axom includes
 #include "axom/config.hpp"  // for AXOM compile-time definitions
 #include "axom/core/Macros.hpp"
+#include "axom/core/utilities/Utilities.hpp"
 
 // Umpire includes
 #ifdef AXOM_USE_UMPIRE
@@ -209,16 +210,22 @@ inline T* allocate(std::size_t n, int allocID) noexcept
 {
   const std::size_t numbytes = n * sizeof(T);
 
+  if(allocID == DYNAMIC_ALLOCATOR_ID)
+  {
+    return static_cast<T*>(std::malloc(numbytes));
+  }
+
 #ifdef AXOM_USE_UMPIRE
 
   umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
   umpire::Allocator allocator = rm.getAllocator(allocID);
   return static_cast<T*>(allocator.allocate(numbytes));
 
-#else
-  AXOM_UNUSED_VAR(allocID);
-  return static_cast<T*>(std::malloc(numbytes));
 #endif
+  // Without Umpire, only DYNAMIC_ALLOCATOR_ID is valid.
+  axom::utilities::processAbort();
+
+  return nullptr;  // Silence warning.
 }
 //------------------------------------------------------------------------------
 template <typename T>
@@ -232,14 +239,16 @@ inline void deallocate(T*& pointer) noexcept
 #ifdef AXOM_USE_UMPIRE
 
   umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
-  rm.deallocate(pointer);
-
-#else
-
-  std::free(pointer);
+  if(rm.hasAllocator(pointer))
+  {
+    rm.deallocate(pointer);
+    pointer = nullptr;
+    return;
+  }
 
 #endif
 
+  std::free(pointer);
   pointer = nullptr;
 }
 
@@ -258,8 +267,26 @@ inline T* reallocate(T* pointer, std::size_t n, int allocID) noexcept
   }
   else
   {
-    pointer = static_cast<T*>(rm.reallocate(pointer, numbytes));
+    if(rm.hasAllocator(pointer))
+    {
+      pointer = static_cast<T*>(rm.reallocate(pointer, numbytes));
+    }
+    else
+    {
+      /*
+        Reallocate from non-Umpire to Umpire, manually, using
+        allocate, copy and deallocate.  Because we don't know the
+        current size, we first do a (extra) reallocate within the
+        current space just so we have the size for the copy.
+        Is there a better way?
+      */
+      auto tmpPointer = std::realloc(pointer, numbytes);
+      pointer = axom::allocate<T>(n, allocID);
+      copy(pointer, tmpPointer, numbytes);
+      deallocate(tmpPointer);
+    }
   }
+  return pointer;
 
 #else
 
@@ -354,6 +381,8 @@ inline MemorySpace getAllocatorSpace(int allocatorId)
     return MemorySpace::Dynamic;
   }
 #endif
+
+  return MemorySpace::Dynamic;  // Silence warning.
 }
 
 #ifdef AXOM_USE_UMPIRE

src/axom/core/tests/core_memory_management.hpp

@@ -506,3 +506,72 @@ TEST(core_memory_management, allocator_id_from_pointer)
   delete[] buf;
 }
 #endif
+
+//------------------------------------------------------------------------------
+TEST(core_memory_management, interspace_reallocation)
+{
+  // Allocator ids to test.
+  std::vector<int> allocIds(1, axom::DYNAMIC_ALLOCATOR_ID);
+#ifdef AXOM_USE_UMPIRE
+  allocIds.push_back(axom::detail::getAllocatorID<axom::MemorySpace::Host>());
+  #ifdef AXOM_USE_GPU
+  allocIds.push_back(axom::detail::getAllocatorID<axom::MemorySpace::Device>());
+  allocIds.push_back(axom::detail::getAllocatorID<axom::MemorySpace::Unified>());
+    // Does it make sense to check Pinned and Constant memory spaces?
+  #endif
+#endif
+
+  // We'll allocate N items, reallocate to K items, reallocate back to N.
+  constexpr std::size_t N = 5;
+  constexpr std::size_t K = 8;
+  constexpr std::size_t maxNK = std::max(N, K);
+  constexpr std::size_t minNK = std::min(N, K);
+
+  // origOnHost and tempOnHost are for initialization and results-checking on host.
+  int* origOnHost = axom::allocate<int>(maxNK, axom::DYNAMIC_ALLOCATOR_ID);
+  for(std::size_t i = 0; i < maxNK; ++i)
+  {
+    origOnHost[i] = 100 + i;
+  }
+  int* tempOnHost = axom::allocate<int>(maxNK, axom::DYNAMIC_ALLOCATOR_ID);
+
+  // Count differences between origOnHost and tempOnHost.
+  auto countDiffs = [=]() {
+    std::size_t diffCount = 0;
+    for(std::size_t j = 0; j < minNK; ++j)
+    {
+      diffCount += tempOnHost[j] != origOnHost[j];
+    }
+    return diffCount;
+  };
+
+  std::size_t diffCount = 0;
+  for(auto srcAllocId : allocIds)
+  {
+    for(auto dstAllocId : allocIds)
+    {
+      std::cout << "Testing allocator ids " << srcAllocId << " and "
+                << dstAllocId << std::endl;
+      // For each combination of srcAllocId and dstAllocId,
+      // allocate src, reallocate to dst, reallocate back to src.
+
+      int* src = axom::allocate<int>(N, srcAllocId);
+      axom::copy(src, origOnHost, N * sizeof(int));
+
+      int* dst = axom::reallocate(src, K, dstAllocId);
+      axom::copy(tempOnHost, dst, N * sizeof(int));
+      diffCount = countDiffs();
+      EXPECT_EQ(diffCount, 0);
+
+      src = axom::reallocate(dst, N, srcAllocId);
+      axom::copy(tempOnHost, src, N * sizeof(int));
+      diffCount = countDiffs();
+      EXPECT_EQ(diffCount, 0);
+
+      axom::deallocate(src);
+    }
+  }
+
+  axom::deallocate(origOnHost);
+  axom::deallocate(tempOnHost);
+}

src/thirdparty/tests/CMakeLists.txt

@@ -39,7 +39,7 @@ if ( RAJA_FOUND )
   axom_add_executable(NAME raja_smoke_test
                       SOURCES raja_smoke.cpp
                       OUTPUT_DIR ${TEST_OUTPUT_DIRECTORY}
-                      DEPENDS_ON ${raja_smoke_dependencies}
+                      DEPENDS_ON ${raja_smoke_dependencies} core
                       FOLDER axom/thirdparty/tests )
 
   target_include_directories(raja_smoke_test PUBLIC