Add Statistics Resource Adaptor and cython bindings to tracking_resource_adaptor and statistics_resource_adaptor

CMakeLists.txt

@@ -81,6 +81,7 @@ endif(CUDA_STATIC_RUNTIME)
 
 target_link_libraries(rmm INTERFACE rmm::Thrust)
 target_link_libraries(rmm INTERFACE spdlog::spdlog_header_only)
+target_link_libraries(rmm INTERFACE dl)
 target_compile_features(rmm INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
 
 # Set logging level. Must go before including gtests and benchmarks.

include/rmm/detail/stack_trace.hpp

@@ -25,6 +25,8 @@
 #include <sstream>
 
 #if defined(RMM_ENABLE_STACK_TRACES)
+#include <cxxabi.h>
+#include <dlfcn.h>
 #include <execinfo.h>
 #include <memory>
 #include <vector>
@@ -60,12 +62,32 @@ class stack_trace {
 #if defined(RMM_ENABLE_STACK_TRACES)
     std::unique_ptr<char*, decltype(&::free)> strings(
       backtrace_symbols(st.stack_ptrs.data(), st.stack_ptrs.size()), &::free);
+
     if (strings.get() == nullptr) {
       os << "But no stack trace could be found!" << std::endl;
     } else {
-      ///@todo: support for demangling of C++ symbol names
+      // Iterate over the stack pointers converting to a string
       for (std::size_t i = 0; i < st.stack_ptrs.size(); ++i) {
-        os << "#" << i << " in " << strings.get()[i] << std::endl;
+        // Leading index
+        os << "#" << i << " in ";
+
+        auto const str = [&] {
+          Dl_info info;
+          if (dladdr(st.stack_ptrs[i], &info)) {
+            int status = -1;  // Demangle the name. This can occasionally fail
+
+            std::unique_ptr<char, decltype(&::free)> demangled(
+              abi::__cxa_demangle(info.dli_sname, nullptr, 0, &status), &::free);
+            // If it fails, fallback to the dli_name.
+            if (status == 0 or info.dli_sname) {
+              auto name = status == 0 ? demangled.get() : info.dli_sname;
+              return name + std::string(" from ") + info.dli_fname;
+            }
+          }
+          return std::string(strings.get()[i]);
+        }();
+
+        os << str << std::endl;
       }
     }
 #else

include/rmm/mr/device/tracking_resource_adaptor.hpp

@@ -65,6 +65,29 @@ class tracking_resource_adaptor final : public device_memory_resource {
         allocation_size{size} {};
   };
 
+  /**
+   * @brief Utility struct for counting the current, peak, and total value of a number
+   */
+  struct counter {
+    int64_t value{0};  // Current value
+    int64_t peak{0};   // Max value of `value`
+    int64_t total{0};  // Sum of all added values
+
+    counter& operator+=(int64_t x)
+    {
+      value += x;
+      total += x;
+      peak = std::max(value, peak);
+      return *this;
+    }
+
+    counter& operator-=(int64_t x)
+    {
+      value -= x;
+      return *this;
+    }
+  };
+
   /**
    * @brief Construct a new tracking resource adaptor using `upstream` to satisfy
    * allocation requests.
@@ -75,13 +98,13 @@ class tracking_resource_adaptor final : public device_memory_resource {
    * @param capture_stacks If true, capture stacks for allocation calls
    */
   tracking_resource_adaptor(Upstream* upstream, bool capture_stacks = false)
-    : capture_stacks_{capture_stacks}, allocated_bytes_{0}, upstream_{upstream}
+    : capture_stacks_{capture_stacks}, upstream_{upstream}
   {
     RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
   }
 
   tracking_resource_adaptor()                                 = delete;
-  ~tracking_resource_adaptor()                                = default;
+  virtual ~tracking_resource_adaptor()                        = default;
   tracking_resource_adaptor(tracking_resource_adaptor const&) = delete;
   tracking_resource_adaptor(tracking_resource_adaptor&&)      = default;
   tracking_resource_adaptor& operator=(tracking_resource_adaptor const&) = delete;
@@ -133,27 +156,56 @@ class tracking_resource_adaptor final : public device_memory_resource {
    * @return std::size_t number of bytes that have been allocated through this
    * allocator.
    */
-  std::size_t get_allocated_bytes() const noexcept { return allocated_bytes_; }
+  std::size_t get_allocated_bytes() const noexcept { return allocation_bytes_.value; }
 
   /**
-   * @brief Log any outstanding allocations via RMM_LOG_DEBUG
+   * @brief Returns a `counter` struct for this adaptor containing the current,
+   * peak, and total number of allocated bytes or allocation counts for this
+   * adaptor since it was created.
    *
+   * @param return_bytes true to return bytes counter, false to re
+   * @return counter
    */
-  void log_outstanding_allocations() const
+  counter get_counter(bool return_bytes = true) const noexcept
   {
-#if SPDLOG_ACTIVE_LEVEL <= SPDLOG_LEVEL_DEBUG
     read_lock_t lock(mtx_);
-    if (not allocations_.empty()) {
-      std::ostringstream oss;
+
+    return return_bytes ? allocation_bytes_ : allocation_count_;
+  }
+
+  /**
+   * @brief Gets a string containing the outstanding allocation pointers, their
+   * size, and optionally the stack trace for when each pointer was allocated.
+   *
+   * @return std::string Containing the outstanding allocation pointers.
+   */
+  std::string get_outstanding_allocations_str() const
+  {
+    read_lock_t lock(mtx_);
+
+    std::ostringstream oss;
+
+    if (!allocations_.empty()) {
       for (auto const& al : allocations_) {
         oss << al.first << ": " << al.second.allocation_size << " B";
         if (al.second.strace != nullptr) {
           oss << " : callstack:" << std::endl << *al.second.strace;
         }
         oss << std::endl;
       }
-      RMM_LOG_DEBUG("Outstanding Allocations: {}", oss.str());
     }
+
+    return oss.str();
+  }
+
+  /**
+   * @brief Log any outstanding allocations via RMM_LOG_DEBUG
+   *
+   */
+  void log_outstanding_allocations() const
+  {
+#if SPDLOG_ACTIVE_LEVEL <= SPDLOG_LEVEL_DEBUG
+    RMM_LOG_DEBUG("Outstanding Allocations: {}", get_outstanding_allocations_str());
 #endif  // SPDLOG_ACTIVE_LEVEL <= SPDLOG_LEVEL_DEBUG
   }
 
@@ -179,8 +231,11 @@ class tracking_resource_adaptor final : public device_memory_resource {
     {
       write_lock_t lock(mtx_);
       allocations_.emplace(p, allocation_info{bytes, capture_stacks_});
+
+      // Increment the allocation_count_ while we have the lock
+      allocation_bytes_ += bytes;
+      allocation_count_ += 1;
     }
-    allocated_bytes_ += bytes;
 
     return p;
   }
@@ -197,11 +252,41 @@ class tracking_resource_adaptor final : public device_memory_resource {
   void do_deallocate(void* p, std::size_t bytes, cuda_stream_view stream) override
   {
     upstream_->deallocate(p, bytes, stream);
+
     {
       write_lock_t lock(mtx_);
-      allocations_.erase(p);
+
+      const auto found = allocations_.find(p);
+
+      // Ensure the allocation is found and the number of bytes match
+      if (found == allocations_.end()) {
+        // Don't throw but log an error. Throwing in a descructor (or any noexcept) will call
+        // std::terminate
+        RMM_LOG_ERROR(
+          "Deallocating a pointer that was not tracked. Ptr: {:p} [{}B], Current Num. Allocations: "
+          "{}",
+          fmt::ptr(p),
+          bytes,
+          this->allocations_.size());
+      } else {
+        allocations_.erase(found);
+
+        auto allocated_bytes = found->second.allocation_size;
+
+        if (allocated_bytes != bytes) {
+          // Don't throw but log an error. Throwing in a descructor (or any noexcept) will call
+          // std::terminate
+          RMM_LOG_ERROR(
+            "Alloc bytes ({}) and Dealloc bytes ({}) do not match", allocated_bytes, bytes);
+
+          bytes = allocated_bytes;
+        }
+      }
+
+      // Decrement the current allocated counts.
+      allocation_bytes_ -= bytes;
+      allocation_count_ -= 1;
     }
-    allocated_bytes_ -= bytes;
   }
 
   /**
@@ -239,7 +324,8 @@ class tracking_resource_adaptor final : public device_memory_resource {
 
   bool capture_stacks_;                           // whether or not to capture call stacks
   std::map<void*, allocation_info> allocations_;  // map of active allocations
-  std::atomic<std::size_t> allocated_bytes_;      // number of bytes currently allocated
+  counter allocation_bytes_;                      // peak, current and total allocated bytes
+  counter allocation_count_;                      // peak, current and total allocation count
   std::shared_timed_mutex mutable mtx_;           // mutex for thread safe access to allocations_
   Upstream* upstream_;  // the upstream resource used for satisfying allocation requests
 };

python/rmm/_lib/memory_resource.pxd

@@ -50,4 +50,7 @@ cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor):
     cpdef get_file_name(self)
     cpdef flush(self)
 
+cdef class TrackingResourceAdaptor(UpstreamResourceAdaptor):
+    pass
+
 cpdef DeviceMemoryResource get_current_device_resource()

python/rmm/_lib/memory_resource.pyx

@@ -4,7 +4,7 @@ import warnings
 from collections import defaultdict
 
 from cython.operator cimport dereference as deref
-from libc.stdint cimport int8_t
+from libc.stdint cimport int8_t, int64_t
 from libcpp cimport bool
 from libcpp.cast cimport dynamic_cast
 from libcpp.memory cimport make_shared, make_unique, shared_ptr, unique_ptr
@@ -83,6 +83,25 @@ cdef extern from "rmm/mr/device/logging_resource_adaptor.hpp" \
 
         void flush() except +
 
+cdef extern from "rmm/mr/device/tracking_resource_adaptor.hpp" \
+        namespace "rmm::mr" nogil:
+    cdef cppclass tracking_resource_adaptor[Upstream](device_memory_resource):
+        struct counter:
+            counter()
+
+            int64_t value
+            int64_t peak
+            int64_t total
+
+        tracking_resource_adaptor(
+            Upstream* upstream_mr,
+            bool capture_stacks) except +
+
+        counter get_counter(bool return_bytes) except +
+
+        string get_outstanding_allocations_str() except +
+        void log_outstanding_allocations() except +
+
 cdef extern from "rmm/mr/device/per_device_resource.hpp" namespace "rmm" nogil:
 
     cdef cppclass cuda_device_id:
@@ -457,6 +476,76 @@ cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor):
         self.c_obj.reset()
 
 
+cdef class TrackingResourceAdaptor(UpstreamResourceAdaptor):
+
+    def __cinit__(
+        self,
+        DeviceMemoryResource upstream_mr,
+        bool capture_stacks=False
+    ):
+        self.c_obj.reset(
+            new tracking_resource_adaptor[device_memory_resource](
+                upstream_mr.get_mr(),
+                capture_stacks
+            )
+        )
+
+    def __init__(
+        self,
+        DeviceMemoryResource upstream_mr,
+        bool capture_stacks=False
+    ):
+        """
+        Memory resource that logs tracks allocations/deallocations
+        performed by an upstream memory resource. Includes the ability to
+        query all outstanding allocations with the stack trace, if desired.
+
+        Parameters
+        ----------
+        upstream : DeviceMemoryResource
+            The upstream memory resource.
+        capture_stacks : bool
+            Whether or not to capture the stack trace with each allocation.
+        """
+        pass
+
+    @property
+    def allocation_counts(self) -> dict:
+        counts = (<tracking_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get()))[0].get_counter(False)
+        byte_counts = (<tracking_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get()))[0].get_counter(True)
+
+        return {
+            "current_bytes": byte_counts.value,
+            "current_count": counts.value,
+            "peak_bytes": byte_counts.peak,
+            "peak_count": counts.peak,
+            "total_bytes": byte_counts.total,
+            "total_count": counts.total,
+        }
+
+    def get_outstanding_allocations_str(self) -> str:
+        """
+        Returns a string containing information about the current outstanding
+        allocations. For each allocation, the address, size and optional
+        stack trace are shown.
+        """
+
+        return (<tracking_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get())
+        )[0].get_outstanding_allocations_str().decode('UTF-8')
+
+    def log_outstanding_allocations(self):
+        """
+        Logs the output of `get_outstanding_allocations_str` to the current
+        RMM log file if enabled.
+        """
+
+        (<tracking_resource_adaptor[device_memory_resource]*>(
+            self.c_obj.get()))[0].log_outstanding_allocations()
+
+
 # Global per-device memory resources; dict of int:DeviceMemoryResource
 cdef _per_device_mrs = defaultdict(CudaMemoryResource)
 

python/rmm/mr.py

@@ -8,6 +8,7 @@
     LoggingResourceAdaptor,
     ManagedMemoryResource,
     PoolMemoryResource,
+    TrackingResourceAdaptor,
     _flush_logs,
     _initialize,
     disable_logging,
@@ -31,6 +32,7 @@
     "LoggingResourceAdaptor",
     "ManagedMemoryResource",
     "PoolMemoryResource",
+    "TrackingResourceAdaptor",
     "_flush_logs",
     "_initialize",
     "set_per_device_resource",