Umpire/CHAI changes

CMakeLists.txt

@@ -106,6 +106,7 @@ endif()
 # Configure CHAI/Umpire for memory management
 # 
 option(ENABLE_CHAI "Enable CHAI/Umpire memory management" Off)
+option(ENABLE_CHAI_SINGLE_MEMORY "Enable single memory configuration in CHAI" Off)
 if(ENABLE_CHAI)
   # Find camp
   if (NOT TARGET camp)
@@ -140,14 +141,18 @@ if(ENABLE_CHAI)
   endif ()
 
   if (ENABLE_HIP)
-    # Set CHAI-related HIP variables for single memory address space.
-    # Use the THIN_GPU_ALLOCATE mode of CHAI on MI250X.
-    # This is to mirror the expected behavior of El Capitan.
-    # Requires env variables HSA_XNACK set to 1 and MPICH_GPU_SUPPORT_ENABLED set to 1 to work on the MI250X.
-    set (CHAI_ENABLE_UM OFF CACHE BOOL "")
-    set (CHAI_ENABLE_PINNED ON CACHE BOOL "")
-    set(CHAI_DISABLE_RM ON CACHE BOOL "")
-    set(CHAI_THIN_GPU_ALLOCATE ON CACHE BOOL "")
+    if(ENABLE_CHAI_SINGLE_MEMORY)
+        # Set CHAI-related HIP variables for single memory address space.
+        # Use the THIN_GPU_ALLOCATE mode of CHAI on MI250X.
+        # This is to mirror the expected behavior of El Capitan.
+        # Requires env variables HSA_XNACK set to 1 and MPICH_GPU_SUPPORT_ENABLED set to 1 to work on the MI250X.
+        set (CHAI_ENABLE_UM OFF CACHE BOOL "")
+        set (CHAI_ENABLE_PINNED ON CACHE BOOL "")
+        set(CHAI_DISABLE_RM ON CACHE BOOL "")
+        set(CHAI_THIN_GPU_ALLOCATE ON CACHE BOOL "")
+
+        set(KRIPKE_USE_CHAI_SINGLE_MEMORY 1)
+    endif()
 
     message(STATUS "Kripke: Setting CAMP_HAVE_HIP")
     blt_add_target_definitions(TO camp

src/Kripke/Core/DataStore.cpp

@@ -45,6 +45,16 @@ DataStore::~DataStore(){
 
 }
 
+double DataStore::getUmpireDeviceHighWatermark() {
+#ifdef KRIPKE_USE_CHAI
+  auto chai_resource_manager = chai::ArrayManager::getInstance();
+  auto device_allocator = chai_resource_manager->getAllocator(chai::GPU);
+  return ((double) device_allocator.getHighWatermark()) / (1024 * 1024 * 1024);
+#else
+  return 0.0;
+#endif
+}
+
 void DataStore::addVariable(std::string const &name,
   Kripke::Core::BaseVar *var)
 {

src/Kripke/Core/DataStore.h

@@ -27,6 +27,8 @@ class DataStore {
     DataStore(DataStore const &) = delete;
     DataStore &operator=(DataStore const &) = delete;
 
+    double getUmpireDeviceHighWatermark();
+
     void addVariable(std::string const &name, Kripke::Core::BaseVar *);
 
     template<typename T, typename ... CTOR_ARGS>

src/Kripke/Core/Field.h

@@ -67,8 +67,12 @@ namespace Core {
           m_chunk_to_size[chunk_id] = sdom_size;
 #ifndef KRIPKE_USE_CHAI
           m_chunk_to_data[chunk_id] = new ElementType[sdom_size];
+#else
+#ifdef KRIPKE_USE_CHAI_SINGLE_MEMORY
+          m_chunk_to_data[chunk_id].allocate(sdom_size, chai::GPU,
 #else
           m_chunk_to_data[chunk_id].allocate(sdom_size, chai::CPU,
+#endif
               [=](const chai::PointerRecord* record, chai::Action action, chai::ExecutionSpace space){
                 /*printf("CHAI[%s, %d]: ", BaseVar::getName().c_str(), (int)chunk_id);
                 switch(action){

src/KripkeConfig.h.in

@@ -11,6 +11,7 @@
 #cmakedefine KRIPKE_USE_OPENMP
 #cmakedefine KRIPKE_USE_CUDA
 #cmakedefine KRIPKE_USE_CHAI
+#cmakedefine KRIPKE_USE_CHAI_SINGLE_MEMORY
 #cmakedefine KRIPKE_USE_HIP
 
 #cmakedefine KRIPKE_USE_CALIPER

src/kripke.cpp

@@ -536,6 +536,21 @@ int main(int argc, char **argv) {
     printf("  Grind time :        %e [(seconds/iteration)/unknowns]\n", grind_time);
     printf("  Sweep efficiency :  %4.5lf [100.0 * SweepSubdomain time / SweepSolver time]\n", sweep_eff);
     printf("  Number of unknowns: %lu\n", (unsigned long) num_unknowns);
+
+#ifdef KRIPKE_USE_CHAI
+    double umpire_device_high_watermark = data_store.getUmpireDeviceHighWatermark();
+#ifdef KRIPKE_USE_CALIPER
+    adiak::value("umpire_device_pool_size", vars.dev_pool_size);
+    adiak::value("umpire_device_high_watermark", umpire_device_high_watermark);
+#endif
+    printf("\n");
+    printf("Memory Usage\n");
+    printf("============\n");
+    printf("\n");
+    printf("  Umpire device pool size:         %4.2lf GB\n", (double) vars.dev_pool_size);
+    printf("  Umpire device high water mark:   %4.2lf GB\n", umpire_device_high_watermark);
+#endif
+
   }
 
 #ifdef KRIPKE_USE_CALIPER