Add AllocatorWithStats

This is a layer under Allocator that can provide stats. The idea is to allow users to expose stats for their allocators or not.
LLNL · Nov 24, 2021 · e5d6aca · e5d6aca
1 parent 0287c8e
commit e5d6aca
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 8 deletions.
diff --git a/include/RAJA/policy/cuda/MemUtils_CUDA.hpp b/include/RAJA/policy/cuda/MemUtils_CUDA.hpp
@@ -107,6 +107,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_device_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<DeviceBaseAllocator>(
         std::string("RAJA::cuda::default_device_allocator"),
+        Platform::cuda,
         DeviceBaseAllocator()));
 }
 ///
@@ -115,6 +116,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_device_zeroed_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<DeviceZeroedBaseAllocator>(
         std::string("RAJA::cuda::default_device_zeroed_allocator"),
+        Platform::cuda,
         DeviceZeroedBaseAllocator()));
 }
 ///
@@ -123,6 +125,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_pinned_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<PinnedBaseAllocator>(
         std::string("RAJA::cuda::default_pinned_allocator"),
+        Platform::cuda,
         PinnedBaseAllocator()));
 }
 

diff --git a/include/RAJA/policy/hip/MemUtils_HIP.hpp b/include/RAJA/policy/hip/MemUtils_HIP.hpp
@@ -110,6 +110,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_device_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<DeviceBaseAllocator>(
         std::string("RAJA::hip::default_device_allocator"),
+        Platform::hip,
         DeviceBaseAllocator()));
 }
 ///
@@ -118,6 +119,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_device_zeroed_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<DeviceZeroedBaseAllocator>(
         std::string("RAJA::hip::default_device_zeroed_allocator"),
+        Platform::hip,
         DeviceZeroedBaseAllocator()));
 }
 ///
@@ -126,6 +128,7 @@ inline std::unique_ptr<RAJA::Allocator> make_default_pinned_allocator()
   return std::unique_ptr<RAJA::Allocator>(
       new AllocatorPool<PinnedBaseAllocator>(
         std::string("RAJA::hip::default_pinned_allocator"),
+        Platform::hip,
         PinnedBaseAllocator()));
 }
 

diff --git a/include/RAJA/util/Allocator.hpp b/include/RAJA/util/Allocator.hpp
@@ -85,6 +85,30 @@ struct Allocator
 
   virtual void release() = 0;
 
+  virtual const std::string& getName() const noexcept = 0;
+};
+
+
+/*! \class AllocatorWithStats
+ ******************************************************************************
+ *
+ * \brief  AllocatorWithStats Provides a generic interface for allocation and
+ *         getting allocation statistics in RAJA
+ *
+ ******************************************************************************
+ */
+struct AllocatorWithStats : Allocator
+{
+  AllocatorWithStats() = default;
+
+  // not copyable or movable
+  AllocatorWithStats(AllocatorWithStats const&) = delete;
+  AllocatorWithStats(AllocatorWithStats &&) = delete;
+  AllocatorWithStats& operator=(AllocatorWithStats const&) = delete;
+  AllocatorWithStats& operator=(AllocatorWithStats &&) = delete;
+
+  virtual ~AllocatorWithStats() = default;
+
   virtual size_t getHighWatermark() const noexcept = 0;
 
   virtual size_t getCurrentSize() const noexcept = 0;
@@ -93,9 +117,7 @@ struct Allocator
 
   virtual size_t getAllocationCount() const noexcept = 0;
 
-  virtual const std::string& getName() const noexcept = 0;
-
-  // virtual Platform getPlatform() const noexcept = 0;
+  virtual Platform getPlatform() const noexcept = 0;
 };
 
 namespace detail

diff --git a/include/RAJA/util/AllocatorPool.hpp b/include/RAJA/util/AllocatorPool.hpp
@@ -325,18 +325,20 @@ struct MemoryArena
  ******************************************************************************
  */
 template <typename allocator_t>
-struct AllocatorPool : Allocator
+struct AllocatorPool : AllocatorWithStats
 {
   using allocator_type = allocator_t;
 
   static const size_t default_default_arena_size = 32ull * 1024ull * 1024ull;
 
   AllocatorPool(std::string const& name,
+                Platform platform,
                 allocator_type const& aloc = allocator_type{},
                 size_t default_arena_size = default_default_arena_size)
     : m_default_arena_size(default_arena_size)
     , m_alloc(aloc)
     , m_name(name) // std::string("RAJA::AllocatorPool<")+m_alloc.getName()+">")
+    , m_platform(platform)
   {
   }
 
@@ -489,10 +491,10 @@ struct AllocatorPool : Allocator
     return m_name;
   }
 
-  // Platform getPlatform() const noexcept final
-  // {
-  //   return m_alloc.getPlatform();
-  // }
+  Platform getPlatform() const noexcept final
+  {
+    return m_platform;
+  }
 
 private:
   using arena_container_type = std::list<detail::MemoryArena>;
@@ -501,6 +503,7 @@ struct AllocatorPool : Allocator
   size_t m_default_arena_size;
   allocator_t m_alloc;
   std::string m_name;
+  Platform m_platform;
 
   size_t m_highWatermark = 0;
   size_t m_currentSize = 0;