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
@@ -355,6 +360,7 @@ configure_file(${PROJECT_SOURCE_DIR}/src/KripkeConfig.h.in
 blt_add_library(
   NAME        kripke
   SOURCES     "src/Kripke/Core/BaseVar.cpp"
+              "src/Kripke/Core/MemoryManager.cpp"
               "src/Kripke/Core/DataStore.cpp"
               "src/Kripke/Core/DomainVar.cpp"
               "src/Kripke/Generate.cpp"

src/Kripke/Core/DataStore.cpp

@@ -19,21 +19,7 @@
 using namespace Kripke;
 using namespace Kripke::Core;
 
-DataStore::DataStore(int dev_pool_size){
-#ifdef KRIPKE_USE_CHAI
-  auto &rm = umpire::ResourceManager::getInstance();
-  const char * allocator_name = "KRIPKE_DEVICE_POOL";
-  size_t umpire_dev_pool_size = ((size_t) dev_pool_size) * 1024 * 1024 * 1024;
-  size_t umpire_dev_block_size = 512;
-  auto dev_pool_allocator = rm.makeAllocator<umpire::strategy::QuickPool>(allocator_name, rm.getAllocator("DEVICE"), umpire_dev_pool_size, umpire_dev_block_size);
-  auto chai_resource_manager = chai::ArrayManager::getInstance();
-  chai_resource_manager->setAllocator(chai::GPU, dev_pool_allocator);
-  // force allocation of GPU memory pool
-  auto tmp = new chai::ManagedArray<int>(100, chai::GPU);
-  tmp->free(chai::GPU);
-  delete tmp;
-#endif // KRIPKE_USE_CHAI
-}
+DataStore::DataStore(){}
 
 DataStore::~DataStore(){
 

src/Kripke/Core/DataStore.h

@@ -22,7 +22,7 @@ namespace Core {
  */
 class DataStore {
   public:
-    DataStore(int dev_pool_size = 4);
+    DataStore();
     ~DataStore();
     DataStore(DataStore const &) = delete;
     DataStore &operator=(DataStore const &) = delete;

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/Kripke/Core/MemoryManager.cpp

@@ -0,0 +1,54 @@
+//
+// Copyright (c) 2014-25, Lawrence Livermore National Security, LLC
+// and Kripke project contributors. See the Kripke/COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//
+
+#include <Kripke.h>
+#include <Kripke/Core/MemoryManager.h>
+
+#ifdef KRIPKE_USE_CHAI
+#define DEBUG
+#include <umpire/Umpire.hpp>
+#include <umpire/strategy/QuickPool.hpp>
+#include <chai/ManagedArray.hpp>
+#undef DEBUG
+#endif
+
+using namespace Kripke;
+using namespace Kripke::Core;
+
+MemoryManager::MemoryManager(int device_pool_size) : device_pool_size(device_pool_size) {
+#ifdef KRIPKE_USE_CHAI
+  auto &rm = umpire::ResourceManager::getInstance();
+  const char * allocator_name = "KRIPKE_DEVICE_POOL";
+  size_t umpire_device_pool_size = ((size_t) device_pool_size) * 1024 * 1024 * 1024;
+  size_t umpire_dev_block_size = 512;
+  auto device_pool_allocator = rm.makeAllocator<umpire::strategy::QuickPool>(allocator_name, rm.getAllocator("DEVICE"), umpire_device_pool_size, umpire_dev_block_size);
+  auto chai_resource_manager = chai::ArrayManager::getInstance();
+  chai_resource_manager->setAllocator(chai::GPU, device_pool_allocator);
+  // force allocation of GPU memory pool
+  auto tmp = new chai::ManagedArray<int>(100, chai::GPU);
+  tmp->free(chai::GPU);
+  delete tmp;
+#endif // KRIPKE_USE_CHAI
+}
+
+double MemoryManager::getDeviceMemoryPoolSize() {
+#ifdef KRIPKE_USE_CHAI
+  return (double) device_pool_size;
+#else
+      return 0.0;
+#endif
+}
+
+double MemoryManager::getDeviceMemoryHighWatermark() {
+#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
+}

src/Kripke/Core/MemoryManager.h

@@ -0,0 +1,28 @@
+//
+// Copyright (c) 2014-25, Lawrence Livermore National Security, LLC
+// and Kripke project contributors. See the Kripke/COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+//
+
+#ifndef KRIPKE_CORE_MEMORYMANAGER_H__
+#define KRIPKE_CORE_MEMORYMANAGER_H__
+
+#include <Kripke.h>
+
+namespace Kripke {
+namespace Core {
+
+class MemoryManager {
+  protected:
+    int device_pool_size;
+
+  public:
+    MemoryManager(int device_pool_size);
+    double getDeviceMemoryPoolSize();
+    double getDeviceMemoryHighWatermark();
+};
+
+} } // namespace
+
+#endif

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

@@ -7,6 +7,7 @@
 
 #include <Kripke.h>
 #include <Kripke/Core/Comm.h>
+#include <Kripke/Core/MemoryManager.h>
 #include <Kripke/Core/DataStore.h>
 #include <Kripke/Core/Set.h>
 #include <Kripke/ArchLayout.h>
@@ -500,7 +501,8 @@ int main(int argc, char **argv) {
 
   // Allocate problem
 
-  Kripke::Core::DataStore data_store(vars.dev_pool_size);
+  Kripke::Core::MemoryManager memory_manager(vars.dev_pool_size);
+  Kripke::Core::DataStore data_store;
   Kripke::generateProblem(data_store, vars);
 
   // Run the solver
@@ -536,6 +538,22 @@ 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 device_memory_pool_size = memory_manager.getDeviceMemoryPoolSize();
+    double device_memory_high_watermark = memory_manager.getDeviceMemoryHighWatermark();
+#ifdef KRIPKE_USE_CALIPER
+    adiak::value("umpire_device_pool_size", device_memory_pool_size);
+    adiak::value("umpire_device_high_watermark", device_memory_high_watermark);
+#endif
+    printf("\n");
+    printf("Memory Usage\n");
+    printf("============\n");
+    printf("\n");
+    printf("  Device pool size:         %4.2lf GB\n", device_memory_pool_size);
+    printf("  Device high water mark:   %4.2lf GB\n", device_memory_high_watermark);
+#endif
+
   }
 
 #ifdef KRIPKE_USE_CALIPER