Address review comments

Author: tbennun

src/callbacks/gpu_memory_usage.cpp

@@ -30,6 +30,7 @@
 #include "lbann/models/model.hpp"
 #include "lbann/utils/gpu/helpers.hpp"
 #include "lbann/utils/serialize.hpp"
+#include <h2/gpu/memory_utils.hpp>
 #include <iomanip>
 #include <sstream>
 
@@ -79,13 +80,7 @@ void gpu_memory_usage::write_specific_proto(lbann_data::Callback& proto) const
 void gpu_memory_usage::on_epoch_begin(model* m)
 {
 #ifdef LBANN_HAS_GPU
-  size_t available;
-  size_t total;
-#ifdef LBANN_HAS_CUDA
-  FORCE_CHECK_CUDA(cudaMemGetInfo(&available, &total));
-#elif defined(LBANN_HAS_ROCM)
-  FORCE_CHECK_ROCM(hipMemGetInfo(&available, &total));
-#endif
+  auto const [available, total] = h2::gpu::mem_info();
   size_t used = total - available;
   auto comm = m->get_comm();
   if (comm->am_trainer_master()) {

src/callbacks/memory_profiler.cpp

@@ -37,6 +37,7 @@
 #include "lbann/proto/callbacks.pb.h"
 
 #include <algorithm>
+#include <h2/gpu/memory_utils.hpp>
 #include <string>
 
 namespace lbann {
@@ -45,17 +46,12 @@ namespace callback {
 /**
  * @brief Returns the currently used memory, or 0 if LBANN was not compiled with
  * GPU support.
+ * TODO(later): Gather across all ranks?
  */
 static inline size_t get_used_gpu_memory()
 {
 #ifdef LBANN_HAS_GPU
-  size_t available;
-  size_t total;
-#ifdef LBANN_HAS_CUDA
-  FORCE_CHECK_CUDA(cudaMemGetInfo(&available, &total));
-#elif defined(LBANN_HAS_ROCM)
-  FORCE_CHECK_ROCM(hipMemGetInfo(&available, &total));
-#endif
+  auto const [available, total] = h2::gpu::mem_info();
   // TODO(later): Might be nicer to return a struct with gathered information
   // (min, max, median across ranks)
   return total - available;
@@ -71,14 +67,7 @@ static inline size_t get_used_gpu_memory()
 static inline size_t get_total_gpu_memory()
 {
 #ifdef LBANN_HAS_GPU
-  size_t available;
-  size_t total;
-#ifdef LBANN_HAS_CUDA
-  FORCE_CHECK_CUDA(cudaMemGetInfo(&available, &total));
-#elif defined(LBANN_HAS_ROCM)
-  FORCE_CHECK_ROCM(hipMemGetInfo(&available, &total));
-#endif
-  return total;
+  return h2::gpu::mem_info().total;
 #else
   return 0;
 #endif
@@ -103,8 +92,9 @@ memory_profiler::memory_profiler(bool detailed_first_step)
   : callback_base(), m_detailed_first_step(detailed_first_step)
 {
 #ifndef LBANN_HAS_GPU
-  LBANN_WARNING("Memory profiler callback only provides unaccounted memory "
-                "information with GPU support.");
+  LBANN_WARNING(
+    "Without GPU support, the memory profiler callback does not ",
+    "provide raw memory usage information, only expected allocation size.");
 #endif
 }