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,23 @@ 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.
+ * A stack of counters is maintained, use `.push_counters()` and `.pop_counters()`
+ * to track statistics at different nesting levels.
+ *
+ * `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 +87,23 @@ 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
+ *
+ * @note When updating the peak value, we assume that `val` is the inner counter of
+ * `this` on the counter stack so its peak value becomes `this->value + val.peak`.
+ *
+ * @param val Value to add
+ * @return Reference to this object
+ */
+ counter& operator+=(const counter& val)
+ {
+ peak = std::max(value + val.peak, peak);
+ value += val.value;
+ total += val.total;
+ return *this;
+ }
  };
 
  /**
@@ -96,6 +117,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 +154,7 @@ class statistics_resource_adaptor final : public device_memory_resource {
  {
  read_lock_t lock(mtx_);
 
- return bytes_;
+ return counter_stack_.top().first;
  }
 
  /**
@@ -145,7 +168,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 pair of counters <bytes, allocations> from the stack _before_ the push
+ */
+ std::pair<counter, counter> push_counters()
+ {
+ write_lock_t lock(mtx_);
+ // auto [bytes, allocations] = counter_stack_.top();
+ // bytes.
+ 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 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();
+ // The new top inherits the statistics
+ counter_stack_.top().first += ret.first;
+ counter_stack_.top().second += ret.second;
+ return ret;
  }
 
  private:
@@ -171,8 +227,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 +249,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 +269,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
 };
 
 /**

python/rmm/rmm/_lib/memory_resource.pyx

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2022, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -177,20 +177,20 @@ cdef extern from "rmm/mr/device/logging_resource_adaptor.hpp" \
 
 cdef extern from "rmm/mr/device/statistics_resource_adaptor.hpp" \
  namespace "rmm::mr" nogil:
- cdef cppclass statistics_resource_adaptor[Upstream](
- device_memory_resource):
+ cdef cppclass statistics_resource_adaptor[Upstream](device_memory_resource):
  struct counter:
  counter()
 
  int64_t value
  int64_t peak
  int64_t total
 
- statistics_resource_adaptor(
- Upstream* upstream_mr) except +
+ statistics_resource_adaptor(Upstream* upstream_mr) except +
 
  counter get_bytes_counter() except +
  counter get_allocations_counter() except +
+ pair[counter, counter] pop_counters() except +
+ pair[counter, counter] push_counters() except +
 
 cdef extern from "rmm/mr/device/tracking_resource_adaptor.hpp" \
  namespace "rmm::mr" nogil:
@@ -793,6 +793,9 @@ cdef class StatisticsResourceAdaptor(UpstreamResourceAdaptor):
  allocations/deallocations performed by an upstream memory resource.
  Includes the ability to query these statistics at any time.
 
+ A stack of counters is maintained, use `.push_counters()` and
+ `.pop_counters()` to track statistics at different nesting levels.
+
  Parameters
  ----------
  upstream : DeviceMemoryResource
@@ -812,12 +815,11 @@ cdef class StatisticsResourceAdaptor(UpstreamResourceAdaptor):
  Returns:
  dict: Dictionary containing allocation counts and bytes.
  """
+ cdef statistics_resource_adaptor[device_memory_resource]* mr = \
+ <statistics_resource_adaptor[device_memory_resource]*> self.c_obj.get()
 
- counts = (<statistics_resource_adaptor[device_memory_resource]*>(
- self.c_obj.get()))[0].get_allocations_counter()
- byte_counts = (<statistics_resource_adaptor[device_memory_resource]*>(
- self.c_obj.get()))[0].get_bytes_counter()
-
+ counts = deref(mr).get_allocations_counter()
+ byte_counts = deref(mr).get_bytes_counter()
  return {
  "current_bytes": byte_counts.value,
  "current_count": counts.value,
@@ -827,6 +829,41 @@ cdef class StatisticsResourceAdaptor(UpstreamResourceAdaptor):
  "total_count": counts.total,
  }
 
+ def pop_counters(self) -> dict:
+ """
+ Pop a counter pair (bytes and allocations) from the stack
+ """
+ cdef statistics_resource_adaptor[device_memory_resource]* mr = \
+ <statistics_resource_adaptor[device_memory_resource]*> self.c_obj.get()
+
+ bytes_and_allocs = deref(mr).pop_counters()
+ return {
+ "current_bytes": bytes_and_allocs.first.value,
+ "current_count": bytes_and_allocs.second.value,
+ "peak_bytes": bytes_and_allocs.first.peak,
+ "peak_count": bytes_and_allocs.second.peak,
+ "total_bytes": bytes_and_allocs.first.total,
+ "total_count": bytes_and_allocs.second.total,
+ }
+
+ def push_counters(self) -> dict:
+ """
+ Push a new counter pair (bytes and allocations) on the stack
+ """
+
+ cdef statistics_resource_adaptor[device_memory_resource]* mr = \
+ <statistics_resource_adaptor[device_memory_resource]*> self.c_obj.get()
+
+ bytes_and_allocs = deref(mr).push_counters()
+ return {
+ "current_bytes": bytes_and_allocs.first.value,
+ "current_count": bytes_and_allocs.second.value,
+ "peak_bytes": bytes_and_allocs.first.peak,
+ "peak_count": bytes_and_allocs.second.peak,
+ "total_bytes": bytes_and_allocs.first.total,
+ "total_count": bytes_and_allocs.second.total,
+ }
+
 cdef class TrackingResourceAdaptor(UpstreamResourceAdaptor):
 
  def __cinit__(

python/rmm/rmm/statistics.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from contextlib import contextmanager
+from typing import Dict, Optional
+
+import rmm.mr
+
+
+def enable_statistics() -> None:
+ """Enable allocation statistics
+
+ This function is idempotent, if statistics has been enabled for the
+ current RMM resource stack, this is a no-op.
+
+ Warning
+ -------
+ This modifies the current RMM memory resource. StatisticsResourceAdaptor
+ is pushed onto the current RMM memory resource stack and must remain the
+ the top must resource throughout the statistics gathering.
+ """
+
+ mr = rmm.mr.get_current_device_resource()
+ if not isinstance(mr, rmm.mr.StatisticsResourceAdaptor):
+ rmm.mr.set_current_device_resource(
+ rmm.mr.StatisticsResourceAdaptor(mr)
+ )
+
+
+def get_statistics() -> Optional[Dict[str, int]]:
+ """Get the current allocation statistics
+
+ Return
+ ------
+ If enabled, returns the current tracked statistics.
+ If disabled, returns None.
+ """
+ mr = rmm.mr.get_current_device_resource()
+ if isinstance(mr, rmm.mr.StatisticsResourceAdaptor):
+ return mr.allocation_counts
+ return None
+
+
+def push_statistics() -> Optional[Dict[str, int]]:
+ """Push new counters on the current allocation statistics stack
+
+ This returns the current tracked statistics and push a new set
+ of zero counters on the stack of statistics.
+
+ If statistics are disabled (the current memory resource is not an
+ instance of StatisticsResourceAdaptor), this function is a no-op.
+
+ Return
+ ------
+ If enabled, returns the current tracked statistics _before_ the pop.
+ If disabled, returns None.
+ """
+ mr = rmm.mr.get_current_device_resource()
+ if isinstance(mr, rmm.mr.StatisticsResourceAdaptor):
+ return mr.push_counters()
+ return None
+
+
+def pop_statistics() -> Optional[Dict[str, int]]:
+ """Pop the counters of the current allocation statistics stack
+
+ This returns the counters of current tracked statistics and pops
+ them from the stack.
+
+ If statistics are disabled (the current memory resource is not an
+ instance of StatisticsResourceAdaptor), this function is a no-op.
+
+ Return
+ ------
+ If enabled, returns the popped counters.
+ If disabled, returns None.
+ """
+ mr = rmm.mr.get_current_device_resource()
+ if isinstance(mr, rmm.mr.StatisticsResourceAdaptor):
+ return mr.pop_counters()
+ return None
+
+
+@contextmanager
+def statistics():
+ """Context to enable allocation statistics.
+
+ If statistics has been enabled already (the current memory resource is an
+ instance of StatisticsResourceAdaptor), new counters are pushed on the
+ current allocation statistics stack when entering the context and popped
+ again when exiting using `push_statistics()` and `push_statistics()`.
+
+ If statistics has not been enabled, StatisticsResourceAdaptor is set as
+ the current RMM memory resource when entering the context and removed
+ again when exiting.
+
+ Raises
+ ------
+ ValueError
+ If the current RMM memory source was changed while in the context.
+ """
+
+ if push_statistics() is None:
+ # Save the current non-statistics memory resource for later cleanup
+ prior_non_stats_mr = rmm.mr.get_current_device_resource()
+ enable_statistics()
+ else:
+ prior_non_stats_mr = None
+
+ try:
+ current_mr = rmm.mr.get_current_device_resource()
+ yield
+ finally:
+ if current_mr is not rmm.mr.get_current_device_resource():
+ raise ValueError(
+ "RMM memory source stack was changed "
+ "while in the statistics context"
+ )
+ if prior_non_stats_mr is None:
+ pop_statistics()
+ else:
+ rmm.mr.set_current_device_resource(prior_non_stats_mr)