disable pinned-mem from RMM pool (#126)

NWChemEx · Jun 13, 2024 · 8963362 · 8963362
1 parent fa88ad7
commit 8963362
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 10 deletions.
diff --git a/src/tamm/gpu_streams.hpp b/src/tamm/gpu_streams.hpp
@@ -273,6 +273,28 @@ static inline bool gpuEventQuery(gpuEvent_t& event) {
 #endif
 }
 
+static inline void* getPinnedMem(size_t bytes) {
+  void* ptr = nullptr;
+#if defined(USE_CUDA)
+  cudaMallocHost((void**) &ptr, bytes);
+#elif defined(USE_HIP)
+  hipMallocHost((void**) &ptr, bytes);
+#elif defined(USE_DPCPP)
+  ptr = (void*) sycl::malloc_host(bytes, tamm::GPUStreamPool::getInstance().getStream().first);
+#endif
+  return ptr;
+}
+
+static inline void freePinnedMem(void* ptr) {
+#if defined(USE_CUDA)
+  cudaFreeHost(ptr);
+#elif defined(USE_HIP)
+  hipFreeHost(ptr);
+#elif defined(USE_DPCPP)
+  sycl::free(ptr, tamm::GPUStreamPool::getInstance().getStream().first);
+#endif
+}
+
 class GPUStreamPool {
 protected:
   int                      default_deviceID{0};

diff --git a/src/tamm/rmm_memory_pool.hpp b/src/tamm/rmm_memory_pool.hpp
@@ -6,7 +6,7 @@
 
 #include "tamm/mr/gpu_memory_resource.hpp"
 #include "tamm/mr/per_device_resource.hpp"
-#include "tamm/mr/pinned_memory_resource.hpp"
+//#include "tamm/mr/pinned_memory_resource.hpp"
 #endif
 
 #include <tamm/errors.hpp>
@@ -56,8 +56,8 @@ class RMMMemoryManager {
 #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP)
   using device_pool_mr = rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>;
   std::unique_ptr<device_pool_mr> deviceMR;
-  using pinned_pool_mr = rmm::mr::pool_memory_resource<rmm::mr::pinned_memory_resource>;
-  std::unique_ptr<pinned_pool_mr> pinnedHostMR;
+  // using pinned_pool_mr = rmm::mr::pool_memory_resource<rmm::mr::pinned_memory_resource>;
+  // std::unique_ptr<pinned_pool_mr> pinnedHostMR;
 #endif
 
 private:
@@ -71,8 +71,8 @@ class RMMMemoryManager {
 #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP)
   /// Returns a RMM device pool handle
   device_pool_mr& getDeviceMemoryPool() { return *(deviceMR.get()); }
-  /// Returns a RMM pinnedHost pool handle
-  pinned_pool_mr& getPinnedMemoryPool() { return *(pinnedHostMR.get()); }
+  // /// Returns a RMM pinnedHost pool handle
+  // pinned_pool_mr& getPinnedMemoryPool() { return *(pinnedHostMR.get()); }
 #endif
 
   /// Returns a RMM host pool handle
@@ -88,7 +88,7 @@ class RMMMemoryManager {
     hostMR.reset();
 #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP)
     deviceMR.reset();
-    pinnedHostMR.reset();
+    // pinnedHostMR.reset();
 #endif
 
     this->invalid_state = true;
@@ -248,10 +248,11 @@ class RMMMemoryManager {
 
     deviceMR = std::make_unique<device_pool_mr>(new rmm::mr::gpu_memory_resource, max_device_bytes);
 
-    size_t max_pinned_host_bytes{0};
-    max_pinned_host_bytes = 0.18 * free;
-    pinnedHostMR =
-      std::make_unique<pinned_pool_mr>(new rmm::mr::pinned_memory_resource, max_pinned_host_bytes);
+    // size_t max_pinned_host_bytes{0};
+    // max_pinned_host_bytes = 0.18 * free;
+    // pinnedHostMR =
+    //   std::make_unique<pinned_pool_mr>(new rmm::mr::pinned_memory_resource,
+    //   max_pinned_host_bytes);
 #endif
     hostMR = std::make_unique<host_pool_mr>(new rmm::mr::new_delete_resource, max_host_bytes);