Minor fixes

NWChemEx · Aug 21, 2024 · 42a7b81 · 42a7b81
1 parent f55c1e6
commit 42a7b81
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 6 deletions.
diff --git a/src/tamm/execution_context.cpp b/src/tamm/execution_context.cpp
@@ -39,7 +39,7 @@ ExecutionContext::ExecutionContext(ProcGroup pg, DistributionKind default_dist_k
   ranks_pn_ = GA_Cluster_nprocs(GA_Cluster_proc_nodeid(pg.rank().value()));
 #endif
   nnodes_  = pg.size().value() / ranks_pn_;
-  gpus_pn_ = ranks_pn_ / ranks_per_gpu_pool();
+  tamm::getHardwareGPUCount(&gpus_pn_);
 
 #if defined(__APPLE__)
   {

diff --git a/src/tamm/gpu_streams.hpp b/src/tamm/gpu_streams.hpp
@@ -15,7 +15,7 @@
 #include <cuda_runtime_api.h>
 #include <nvml.h>
 #elif defined(USE_HIP)
-#include "rocm_smi/rocm_smi.h"
+#include <rocm_smi/rocm_smi.h>
 #include <hip/hip_runtime.h>
 #include <rocblas/rocblas.h>
 #elif defined(USE_DPCPP)

diff --git a/src/tamm/rmm_memory_pool.hpp b/src/tamm/rmm_memory_pool.hpp
@@ -66,8 +66,6 @@ class RMMMemoryManager {
   uint32_t tamm_rpg;
 
 public:
-  uint32_t get_rpg() { return tamm_rpg; }
-
 #if defined(USE_CUDA) || defined(USE_HIP) || defined(USE_DPCPP)
   /// Returns a RMM device pool handle
   device_pool_mr& getDeviceMemoryPool() { return *(deviceMR.get()); }
@@ -252,8 +250,6 @@ class RMMMemoryManager {
   RMMMemoryManager& operator=(RMMMemoryManager&&)      = delete;
 };
 
-static inline uint32_t ranks_per_gpu_pool() { return RMMMemoryManager::getInstance().get_rpg(); }
-
 // The reset pool & reinitialize only is being used for the (T) segement of cannonical
 static inline void reset_rmm_pool() { RMMMemoryManager::getInstance().reset(); }