Add a stack to the statistics resource

include/rmm/mr/device/statistics_resource_adaptor.hpp

@@ -21,6 +21,7 @@
 #include <cstddef>
 #include <mutex>
 #include <shared_mutex>
+#include <stack>
 
 namespace rmm::mr {
 /**
@@ -36,20 +37,27 @@ namespace rmm::mr {
  * resource in order to satisfy allocation requests, but any existing
  * allocations will be untracked. Tracking statistics stores the current, peak
  * and total memory allocations for both the number of bytes and number of calls
- * to the memory resource. `statistics_resource_adaptor` is intended as a debug
- * adaptor and shouldn't be used in performance-sensitive code.
+ * to the memory resource.
+ *
+ * This resource supports nested statistics, which makes it possible to track statistics
+ * of a code block. Use `.push_counters()` to start tracking statistics on a code block
+ * and use `.pop_counters()` to stop the tracking. The nested statistics are cascading
+ * such that the statistics tracked by a code block include the statistics tracked in
+ * all its tracked sub code blocks.
+ *
+ * `statistics_resource_adaptor` is intended as a debug adaptor and shouldn't be
+ * used in performance-sensitive code.
  *
  * @tparam Upstream Type of the upstream resource used for
  * allocation/deallocation.
  */
 template <typename Upstream>
 class statistics_resource_adaptor final : public device_memory_resource {
  public:
- // can be a std::shared_mutex once C++17 is adopted
  using read_lock_t =
- std::shared_lock<std::shared_timed_mutex>; ///< Type of lock used to synchronize read access
+ std::shared_lock<std::shared_mutex>; ///< Type of lock used to synchronize read access
  using write_lock_t =
- std::unique_lock<std::shared_timed_mutex>; ///< Type of lock used to synchronize write access
+ std::unique_lock<std::shared_mutex>; ///< Type of lock used to synchronize write access
  /**
  * @brief Utility struct for counting the current, peak, and total value of a number
  */
@@ -83,6 +91,24 @@ class statistics_resource_adaptor final : public device_memory_resource {
  value -= val;
  return *this;
  }
+
+ /**
+ * @brief Add `val` to the current value and update the peak value if necessary
+ *
+ * When updating the peak value, we assume that `val` is tracking a code block inside the
+ * code block tracked by `this`. Because nested statistics are cascading, we have to convert
+ * `val.peak` to the peak it would have been if it was part of the statistics tracked by `this`.
+ * We do this by adding the current value that was active when `val` started tracking such that
+ * we get `std::max(value + val.peak, peak)`.
+ *
+ * @param val Value to add
+ */
+ void add_counters_from_tracked_sub_block(const counter& val)
+ {
+ peak = std::max(value + val.peak, peak);
+ value += val.value;
+ total += val.total;
+ }
  };
 
  /**
@@ -96,6 +122,8 @@ class statistics_resource_adaptor final : public device_memory_resource {
  statistics_resource_adaptor(Upstream* upstream) : upstream_{upstream}
  {
  RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
+ // Initially, we push a single counter pair on the stack
+ push_counters();
  }
 
  statistics_resource_adaptor() = delete;
@@ -131,7 +159,7 @@ class statistics_resource_adaptor final : public device_memory_resource {
  {
  read_lock_t lock(mtx_);
 
- return bytes_;
+ return counter_stack_.top().first;
  }
 
  /**
@@ -145,7 +173,40 @@ class statistics_resource_adaptor final : public device_memory_resource {
  {
  read_lock_t lock(mtx_);
 
- return allocations_;
+ return counter_stack_.top().second;
+ }
+
+ /**
+ * @brief Push a pair of zero counters on the stack, which becomes the new
+ * counters returned by `get_bytes_counter()` and `get_allocations_counter()`
+ *
+ * @return top pair of counters <bytes, allocations> from the stack _before_
+ * the push
+ */
+ std::pair<counter, counter> push_counters()
+ {
+ write_lock_t lock(mtx_);
+ auto ret = counter_stack_.top();
+ counter_stack_.push(std::make_pair(counter{}, counter{}));
+ return ret;
+ }
+
+ /**
+ * @brief Pop a pair of counters from the stack
+ *
+ * @return top pair of counters <bytes, allocations> from the stack _before_
+ * the pop
+ */
+ std::pair<counter, counter> pop_counters()
+ {
+ write_lock_t lock(mtx_);
+ if (counter_stack_.size() < 2) { throw std::out_of_range("cannot pop the last counter pair"); }
+ auto ret = counter_stack_.top();
+ counter_stack_.pop();
+ // Update the new top pair of counters
+ counter_stack_.top().first.add_counters_from_tracked_sub_block(ret.first);
+ counter_stack_.top().second.add_counters_from_tracked_sub_block(ret.second);
+ return ret;
  }
 
  private:
@@ -171,8 +232,8 @@ class statistics_resource_adaptor final : public device_memory_resource {
  write_lock_t lock(mtx_);
 
  // Increment the allocation_count_ while we have the lock
- bytes_ += bytes;
- allocations_ += 1;
+ counter_stack_.top().first += bytes;
+ counter_stack_.top().second += 1;
  }
 
  return ptr;
@@ -193,8 +254,8 @@ class statistics_resource_adaptor final : public device_memory_resource {
  write_lock_t lock(mtx_);
 
  // Decrement the current allocated counts.
- bytes_ -= bytes;
- allocations_ -= 1;
+ counter_stack_.top().first -= bytes;
+ counter_stack_.top().second -= 1;
  }
  }
 
@@ -213,10 +274,10 @@ class statistics_resource_adaptor final : public device_memory_resource {
  return get_upstream_resource() == cast->get_upstream_resource();
  }
 
- counter bytes_; // peak, current and total allocated bytes
- counter allocations_; // 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
+ // Stack of counter pairs <bytes, allocations>
+ std::stack<std::pair<counter, counter>> counter_stack_;
+ std::shared_mutex mutable mtx_; // mutex for thread safe access to allocations_
+ Upstream* upstream_;  // the upstream resource used for satisfying allocation requests
 };
 
 /**

include/rmm/mr/device/tracking_resource_adaptor.hpp

@@ -53,11 +53,10 @@ namespace rmm::mr {
 template <typename Upstream>
 class tracking_resource_adaptor final : public device_memory_resource {
  public:
- // can be a std::shared_mutex once C++17 is adopted
  using read_lock_t =
- std::shared_lock<std::shared_timed_mutex>; ///< Type of lock used to synchronize read access
+ std::shared_lock<std::shared_mutex>; ///< Type of lock used to synchronize read access
  using write_lock_t =
- std::unique_lock<std::shared_timed_mutex>; ///< Type of lock used to synchronize write access
+ std::unique_lock<std::shared_mutex>; ///< Type of lock used to synchronize write access
  /**
  * @brief Information stored about an allocation. Includes the size
  * and a stack trace if the `tracking_resource_adaptor` was initialized
@@ -271,7 +270,7 @@ 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
- std::shared_timed_mutex mutable mtx_; // mutex for thread safe access to allocations_
+ std::shared_mutex mutable mtx_;  // mutex for thread safe access to allocations_
  Upstream* upstream_; // the upstream resource used for satisfying allocation requests
 };
 

python/rmm/docs/guide.md

@@ -187,3 +187,106 @@ allocator.
 
 >>> torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
 ```
+
+
+
+## Memory statistics and profiling
+
+RMM can profile memory usage and track memory statistics by using either of the following:
+ - Use the context manager `rmm.statistics.statistics()` to enable statistics tracking for a specific code block.
+ - Call `rmm.statistics.enable_statistics()` to enable statistics tracking globally.
+
+Common to both usages is that they modify the currently active RMM memory resource. The current device resource is wrapped with a `StatisticsResourceAdaptor` which must remain the topmost resource throughout the statistics tracking:
+```python
+>>> import rmm
+>>> import rmm.statistics
+
+>>> # We start with the default cuda memory resource
+>>> rmm.mr.get_current_device_resource()
+<rmm._lib.memory_resource.CudaMemoryResource at 0x7f7e6c0a1ce0>
+
+>>> # When using statistics, we get a StatisticsResourceAdaptor with the context
+>>> with rmm.statistics.statistics():
+... rmm.mr.get_current_device_resource()
+<rmm._lib.memory_resource.StatisticsResourceAdaptor at 0x7f7e6c524900>
+
+>>> # We can also enable statistics globally
+>>> rmm.statistics.enable_statistics()
+>>> print(rmm.mr.get_current_device_resource())
+<rmm._lib.memory_resource.StatisticsResourceAdaptor at 0x7f662c2bb3c0>
+```
+
+With statistics enabled, you can query statistics of the current and peak bytes and number of allocations performed by the current RMM memory resource:
+```python
+>>> buf = rmm.DeviceBuffer(size=10)
+>>> rmm.statistics.get_statistics()
+Statistics(current_bytes=16, current_count=1, peak_bytes=16, peak_count=1, total_bytes=16, total_count=1)
+```
+
+### Memory Profiler
+To profile a specific block of code, first enable memory statistics by calling `rmm.statistics.enable_statistics()`. To profile a function, use `profiler` as a function decorator:
+```python
+>>> @rmm.statistics.profiler()
+... def f(size):
+... rmm.DeviceBuffer(size=size)
+>>> f(1000)
+
+>>> # By default, the profiler write to rmm.statistics.default_profiler_records
+>>> print(rmm.statistics.default_profiler_records.report())
+Memory Profiling
+================
+
+Legends:
+ ncalls - number of times the function or code block was called
+ memory_peak - peak memory allocated in function or code block (in bytes)
+ memory_total - total memory allocated in function or code block (in bytes)
+
+Ordered by: memory_peak
+
+ncalls memory_peak memory_total filename:lineno(function)
+ 1 1,008 1,008 <ipython-input-11-5fc63161ac29>:1(f)
+```
+
+To profile a code block, use `profiler` as a context manager:
+```python
+>>> with rmm.statistics.profiler(name="my code block"):
+... rmm.DeviceBuffer(size=20)
+>>> print(rmm.statistics.default_profiler_records.report())
+Memory Profiling
+================
+
+Legends:
+ ncalls - number of times the function or code block was called
+ memory_peak - peak memory allocated in function or code block (in bytes)
+ memory_total - total memory allocated in function or code block (in bytes)
+
+Ordered by: memory_peak
+
+ncalls memory_peak memory_total filename:lineno(function)
+ 1 1,008 1,008 <ipython-input-11-5fc63161ac29>:1(f)
+ 1 32 32 my code block
+```
+
+The `profiler` supports nesting:
+```python
+>>> with rmm.statistics.profiler(name="outer"):
+... buf1 = rmm.DeviceBuffer(size=10)
+... with rmm.statistics.profiler(name="inner"):
+... buf2 = rmm.DeviceBuffer(size=10)
+>>> print(rmm.statistics.default_profiler_records.report())
+Memory Profiling
+================
+
+Legends:
+ ncalls - number of times the function or code block was called
+ memory_peak - peak memory allocated in function or code block (in bytes)
+ memory_total - total memory allocated in function or code block (in bytes)
+
+Ordered by: memory_peak
+
+ncalls memory_peak memory_total filename:lineno(function)
+ 1 1,008 1,008 <ipython-input-4-865fbe04e29f>:1(f)
+ 1 32 32 my code block
+ 1 32 32 outer
+ 1 16 16 inner
+```

python/rmm/docs/python_api.rst

@@ -37,3 +37,12 @@ Memory Allocators
  :members:
  :undoc-members:
  :show-inheritance:
+
+Memory Statistics
+-----------------
+
+.. automodule:: rmm.statistics
+ :members:
+ :inherited-members:
+ :undoc-members:
+ :show-inheritance: