ManagedSharedPtr (#298)

* Add ManagedSharedPtr as an alternative to managed_ptr. It adds basic reference counting on the host and experimental automatic migration

* Improves object lifetime management and extends it to work with additional types

* Some reorganization and clean up

---------

Co-authored-by: Alan Dayton <[email protected]>

Author: mdavis36

cmake/SetupChaiOptions.cmake

@@ -4,6 +4,7 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 ############################################################################
+option(CHAI_ENABLE_EXPERIMENTAL "Enable experimental chai features." On)
 option(CHAI_ENABLE_GPU_SIMULATION_MODE "Enable GPU Simulation Mode" Off)
 option(CHAI_ENABLE_OPENMP "Enable OpenMP" Off)
 option(CHAI_ENABLE_MPI "Enable MPI (for umpire replay only)" Off)

src/chai/ArrayManager.hpp

@@ -7,12 +7,12 @@
 #ifndef CHAI_ArrayManager_HPP
 #define CHAI_ArrayManager_HPP
 
-#include "chai/config.hpp"
 #include "chai/ChaiMacros.hpp"
 #include "chai/ExecutionSpaces.hpp"
-#include "chai/PointerRecord.hpp"
 #include "chai/Types.hpp"
 
+#include "chai/PointerRecord.hpp"
+
 #if defined(CHAI_ENABLE_RAJA_PLUGIN)
 #include "chai/pluginLinker.hpp"
 #endif
@@ -22,107 +22,12 @@
 #include "umpire/Allocator.hpp"
 #include "umpire/util/MemoryMap.hpp"
 
-#if defined(CHAI_ENABLE_CUDA)
-#include <cuda_runtime_api.h>
-#endif
-#if defined(CHAI_ENABLE_HIP)
-#include "hip/hip_runtime_api.h"
-#endif
 
-namespace chai
-{
-// CHAI_GPU_ERROR_CHECK macro
-#if defined(CHAI_ENABLE_CUDA) || defined(CHAI_ENABLE_HIP)
+#include "chai/DeviceHelpers.hpp"
 
-#ifdef CHAI_ENABLE_GPU_ERROR_CHECKING
 
-#ifdef CHAI_ENABLE_CUDA
-inline void gpuErrorCheck(cudaError_t code, const char *file, int line, bool abort=true)
-{
-   if (code != cudaSuccess) {
-      fprintf(stderr, "[CHAI] GPU Error: %s %s %d\n", cudaGetErrorString(code), file, line);
-      if (abort) {
-         exit(code);
-      }
-   }
-}
-#elif defined(CHAI_ENABLE_HIP)
-inline void gpuErrorCheck(hipError_t code, const char *file, int line, bool abort=true)
+namespace chai
 {
-   if (code != hipSuccess) {
-      fprintf(stderr, "[CHAI] GPU Error: %s %s %d\n", hipGetErrorString(code), file, line);
-      if (abort) {
-         exit(code);
-      }
-   }
-}
-#endif
-
-
-#define CHAI_GPU_ERROR_CHECK(code) { gpuErrorCheck((code), __FILE__, __LINE__); }
-#else // CHAI_ENABLE_GPU_ERROR_CHECKING
-#define CHAI_GPU_ERROR_CHECK(code) code
-#endif // CHAI_ENABLE_GPU_ERROR_CHECKING
-
-#endif
-
-// wrapper for hip/cuda synchronize
-inline void synchronize() {
-#if defined (CHAI_ENABLE_HIP) &&!defined(__HIP_DEVICE_COMPILE__)
-   CHAI_GPU_ERROR_CHECK(hipDeviceSynchronize());
-#elif defined (CHAI_ENABLE_CUDA) &&!defined(__CUDA_ARCH__)
-   CHAI_GPU_ERROR_CHECK(cudaDeviceSynchronize());
-#endif
-}
-
-#if defined(CHAI_GPUCC) || defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
-
-// wrapper for hip/cuda free
-CHAI_HOST inline void gpuFree(void* buffer) {
-#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
-   free(buffer);
-#elif defined (CHAI_ENABLE_HIP)
-   CHAI_GPU_ERROR_CHECK(hipFree(buffer));
-#elif defined (CHAI_ENABLE_CUDA)
-   CHAI_GPU_ERROR_CHECK(cudaFree(buffer));
-#endif
-}
-
-// wrapper for hip/cuda malloc
-CHAI_HOST inline void gpuMalloc(void** devPtr, size_t size) {
-#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
-   *devPtr = (void*)malloc(size);
-#elif defined (CHAI_ENABLE_HIP)
-   CHAI_GPU_ERROR_CHECK(hipMalloc(devPtr, size));
-#elif defined (CHAI_ENABLE_CUDA)
-   CHAI_GPU_ERROR_CHECK(cudaMalloc(devPtr, size));
-#endif
-}
-
-// wrapper for hip/cuda managed malloc
-CHAI_HOST inline void gpuMallocManaged(void** devPtr, size_t size) {
-#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
-   *devPtr = (void*)malloc(size);
-#elif defined (CHAI_ENABLE_HIP)
-   CHAI_GPU_ERROR_CHECK(hipMallocManaged(devPtr, size));
-#elif defined (CHAI_ENABLE_CUDA)
-   CHAI_GPU_ERROR_CHECK(cudaMallocManaged(devPtr, size));
-#endif
-}
-
-// wrapper for hip/cuda mem copy
-CHAI_HOST inline void  gpuMemcpy(void* dst, const void* src, size_t count, gpuMemcpyKind kind) {
-#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
-   memcpy(dst, src, count);
-#elif defined (CHAI_ENABLE_HIP)
-   CHAI_GPU_ERROR_CHECK(hipMemcpy(dst, src, count, kind));
-#elif defined (CHAI_ENABLE_CUDA)
-   CHAI_GPU_ERROR_CHECK(cudaMemcpy(dst, src, count, kind));
-#endif
-}
-
-#endif //#if defined(CHAI_GPUCC)
-
 /*!
  * \brief Singleton that manages caching and movement of ManagedArray objects.
  *

src/chai/CMakeLists.txt

@@ -15,13 +15,24 @@ set (chai_headers
   ArrayManager.hpp
   ArrayManager.inl
   ChaiMacros.hpp
+  DeviceHelpers.hpp
   ExecutionSpaces.hpp
   ManagedArray.hpp
   ManagedArray.inl
   managed_ptr.hpp
   PointerRecord.hpp
   Types.hpp)
 
+if(CHAI_ENABLE_EXPERIMENTAL)
+  set(chai_headers
+    ${chai_headers}
+    ManagedSharedPtr.hpp
+    SharedPtrCounter.hpp
+    SharedPtrManager.hpp
+    SharedPtrManager.inl
+    SharedPointerRecord.hpp)
+endif()
+
 if(CHAI_DISABLE_RM)
   set(chai_headers
     ${chai_headers}
@@ -31,6 +42,12 @@ endif ()
 set (chai_sources
   ArrayManager.cpp)
 
+if(CHAI_ENABLE_EXPERIMENTAL)
+  set (chai_sources
+    ${chai_sources}
+    SharedPtrManager.cpp)
+endif ()
+
 set (chai_depends
   umpire)
 

src/chai/ChaiMacros.hpp

@@ -31,6 +31,13 @@
 #define gpuMemcpyDeviceToDevice cudaMemcpyDeviceToDevice
 #define gpuMemcpyDefault cudaMemcpyDefault
 
+#define gpuSuccess cudaSuccess
+#define gpuError_t cudaError_t
+#define gpuGetErrorString cudaGetErrorString
+#define gpuPeekAtLastError cudaPeekAtLastError
+#define gpuDeviceSynchronize cudaDeviceSynchronize
+
+
 // NOTE: Cannot have if defined(__HIPCC__) in the condition below, since __HIPCC__ comes from the included header hip_runtime below.
 #elif defined(CHAI_ENABLE_HIP)
 
@@ -48,6 +55,12 @@
 #define gpuMemcpyDeviceToDevice hipMemcpyDeviceToDevice
 #define gpuMemcpyDefault hipMemcpyDefault
 
+#define gpuSuccess hipSuccess
+#define gpuError_t hipError_t
+#define gpuGetErrorString hipGetErrorString
+#define gpuPeekAtLastError hipPeekAtLastError
+#define gpuDeviceSynchronize hipDeviceSynchronize
+
 #else
 
 #define CHAI_HOST

src/chai/DeviceHelpers.hpp

@@ -0,0 +1,88 @@
+#ifndef CHAI_DEVICE_HELPERS_HPP
+#define CHAI_DEVICE_HELPERS_HPP
+
+#include "chai/config.hpp"
+#include "chai/ChaiMacros.hpp"
+
+namespace chai
+{
+// CHAI_GPU_ERROR_CHECK macro
+#ifdef CHAI_ENABLE_DEVICE
+
+#ifdef CHAI_ENABLE_GPU_ERROR_CHECKING
+
+inline void gpuErrorCheck(gpuError_t code, const char *file, int line, bool abort=true)
+{
+   if (code != gpuSuccess) {
+      fprintf(stderr, "[CHAI] GPU Error: %s %s %d\n", gpuGetErrorString(code), file, line);
+      if (abort) {
+         exit(code);
+      }
+   }
+}
+
+#define CHAI_GPU_ERROR_CHECK(code) { ::chai::gpuErrorCheck((code), __FILE__, __LINE__); }
+#else // CHAI_ENABLE_GPU_ERROR_CHECKING
+#define CHAI_GPU_ERROR_CHECK(code) code
+#endif // CHAI_ENABLE_GPU_ERROR_CHECKING
+
+#endif
+
+// wrapper for hip/cuda synchronize
+inline void synchronize() {
+#if defined(CHAI_ENABLE_DEVICE) && !defined(CHAI_DEVICE_COMPILE)
+   CHAI_GPU_ERROR_CHECK(gpuDeviceSynchronize());
+#endif
+}
+
+#if defined(CHAI_GPUCC) || defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
+
+// wrapper for hip/cuda free
+CHAI_HOST inline void gpuFree(void* buffer) {
+#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
+   free(buffer);
+#elif defined (CHAI_ENABLE_HIP)
+   CHAI_GPU_ERROR_CHECK(hipFree(buffer));
+#elif defined (CHAI_ENABLE_CUDA)
+   CHAI_GPU_ERROR_CHECK(cudaFree(buffer));
+#endif
+}
+
+// wrapper for hip/cuda malloc
+CHAI_HOST inline void gpuMalloc(void** devPtr, size_t size) {
+#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
+   *devPtr = (void*)malloc(size);
+#elif defined (CHAI_ENABLE_HIP)
+   CHAI_GPU_ERROR_CHECK(hipMalloc(devPtr, size));
+#elif defined (CHAI_ENABLE_CUDA)
+   CHAI_GPU_ERROR_CHECK(cudaMalloc(devPtr, size));
+#endif
+}
+
+// wrapper for hip/cuda managed malloc
+CHAI_HOST inline void gpuMallocManaged(void** devPtr, size_t size) {
+#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
+   *devPtr = (void*)malloc(size);
+#elif defined (CHAI_ENABLE_HIP)
+   CHAI_GPU_ERROR_CHECK(hipMallocManaged(devPtr, size));
+#elif defined (CHAI_ENABLE_CUDA)
+   CHAI_GPU_ERROR_CHECK(cudaMallocManaged(devPtr, size));
+#endif
+}
+
+// wrapper for hip/cuda mem copy
+CHAI_HOST inline void  gpuMemcpy(void* dst, const void* src, size_t count, gpuMemcpyKind kind) {
+#if defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
+   memcpy(dst, src, count);
+#elif defined (CHAI_ENABLE_HIP)
+   CHAI_GPU_ERROR_CHECK(hipMemcpy(dst, src, count, kind));
+#elif defined (CHAI_ENABLE_CUDA)
+   CHAI_GPU_ERROR_CHECK(cudaMemcpy(dst, src, count, kind));
+#endif
+}
+
+#endif //#if defined(CHAI_GPUCC)
+
+}  // namespace chai
+
+#endif // CHAI_DEVICE_HELPERS_HPP

src/chai/ManagedArray.hpp

@@ -373,14 +373,14 @@ class ManagedArray : public CHAICopyable
   // shenanigan reasons need to be defined here.
 #if !defined(CHAI_DISABLE_RM)
   // if T is a CHAICopyable, then it is important to initialize all the
-  // ManagedArrays to nullptr at allocation, since it is extremely easy to
+  // elements with default constructors, since it is extremely easy to
   // trigger a moveInnerImpl, which expects inner values to be initialized.
   template <bool B = std::is_base_of<CHAICopyable, T>::value,
             typename std::enable_if<B, int>::type = 0>
   CHAI_HOST bool initInner(size_t start = 0)
   {
     for (size_t i = start; i < m_size/sizeof(T); ++i) {
-      m_active_base_pointer[i] = nullptr;
+      new (&m_active_base_pointer[i]) T();
     }
     return true;
   }
@@ -392,6 +392,26 @@ class ManagedArray : public CHAICopyable
   {
     return false;
   }
+
+  // if T is a CHAICopyable, then it is important to free all the
+  // CHAICopyable containers, which expect inner values to be initialized.
+  template <bool B = std::is_base_of<CHAICopyable, T>::value,
+            typename std::enable_if<B, int>::type = 0>
+  CHAI_HOST bool freeInner(size_t start = 0)
+  {
+    for (size_t i = start; i < m_size/sizeof(T); ++i) {
+      m_active_base_pointer[i].~T();
+    }
+    return true;
+  }
+
+  // Do not deep initialize if T is not a CHAICopyable.
+  template <bool B = std::is_base_of<CHAICopyable, T>::value,
+            typename std::enable_if<!B, int>::type = 0>
+  CHAI_HOST bool freeInner(size_t = 0)
+  {
+    return false;
+  }
 #endif
 protected:
   /*!

src/chai/ManagedArray.inl

@@ -249,12 +249,14 @@ CHAI_HOST void ManagedArray<T>::reallocate(size_t elems)
       // trigger a moveInnerImpl, which expects inner values to be initialized.
       if (initInner(old_size/sizeof(T))) {
         // if we are active on the  GPU, we need to send any newly initialized inner members to the device
+#if defined(CHAI_ENABLE_CUDA) || defined(CHAI_ENABLE_HIP) || defined(CHAI_ENABLE_GPU_SIMULATION_MODE)
         if (m_pointer_record->m_last_space == GPU && old_size < m_size) {
           umpire::ResourceManager & umpire_rm = umpire::ResourceManager::getInstance();
           void *src = (void *)(((char *)(m_pointer_record->m_pointers[CPU])) + old_size);
           void *dst = (void *)(((char *)(m_pointer_record->m_pointers[GPU])) + old_size);
           umpire_rm.copy(dst,src,m_size-old_size);
         }
+#endif
       }
 
       CHAI_LOG(Debug, "m_active_ptr reallocated at address: " << m_active_pointer);
@@ -276,6 +278,8 @@ CHAI_HOST void ManagedArray<T>::free(ExecutionSpace space)
     if (m_pointer_record == &ArrayManager::s_null_record) {
        m_pointer_record = m_resource_manager->makeManaged((void *)m_active_base_pointer,m_size,space,true);
     }
+    freeInner();
+
     m_resource_manager->free(m_pointer_record, space);
     m_active_pointer = nullptr;
     m_active_base_pointer = nullptr;