Extend OrtAllocator API to get Allocator statistics

include/onnxruntime/core/framework/allocator.h

@@ -101,7 +101,9 @@ class IAllocator {
   const OrtMemoryInfo& Info() const { return memory_info_; };
 
   // Each implementation of IAllocator can override and provide their own implementation
-  virtual void GetStats(AllocatorStats* /*stats*/) { return; }
+  virtual void GetStats(AllocatorStats* stats) {
+    stats->Clear();
+  }
 
   static bool CalcMemSizeForArray(size_t nmemb, size_t size, size_t* out) noexcept {
     return CalcMemSizeForArrayWithAlignment(nmemb, size, 0, out);

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -340,6 +340,25 @@ typedef struct OrtAllocator {
    * those made during session initialization. This allows for separate memory management strategies for these allocations.
    */
   void*(ORT_API_CALL* Reserve)(struct OrtAllocator* this_, size_t size);  ///< Returns a pointer to an allocated block of `size` bytes
+
+  /**
+   * @brief Function used to get the statistics of the allocator.
+   *
+   * Return a pointer to the OrtKeyValuePairs structure that contains the statistics of the allocator.
+   * Supported keys are:
+   * - Limit: Bytes limit of the allocator. -1 if no limit is set.
+   * - InUse: Number of bytes in use.
+   * - TotalAllocated: The total number of allocated bytes by the allocator.
+   * - MaxInUse: The maximum bytes in use.
+   * - NumAllocs: Number of allocations.
+   * - NumReserves: Number of reserves. (Number of calls to Reserve() in arena-based allocators)
+   * - NumArenaExtensions: Number of arena extensions (Relevant only for arena based allocators)
+   * - NumArenaShrinkages: Number of arena shrinkages (Relevant only for arena based allocators)
+   * - MaxAllocSize: The max single allocation seen.
+   *
+   * NOTE: If the allocator does not implement this function, the OrtKeyValuePairs instance will be empty.
+   */
+  ORT_API2_STATUS(GetStats, _In_ const struct OrtAllocator* this_, _Outptr_ OrtKeyValuePairs** out);
 } OrtAllocator;
 
 typedef void(ORT_API_CALL* OrtLoggingFunction)(
@@ -5266,6 +5285,21 @@ struct OrtApi {
    * \since Version 1.23
    */
   ORT_API2_STATUS(GetTensorSizeInBytes, _In_ const OrtValue* ort_value, _Out_ size_t* size);
+
+  /** \brief Calls OrtAllocator::GetStats function
+   *
+   * Return a pointer to the OrtKeyValuePairs structure that contains the statistics of the allocator.
+   *
+   * NOTE: If the allocator does not implement this function, the OrtKeyValuePairs instance will be empty.
+   *
+   * \param[in] ort_allocator The allocator to get stats from
+   * \param[out] out A pointer to the OrtKeyValuePairs instance that contains the stats
+   *
+   * \snippet{doc} snippets.dox OrtStatus Return Value
+   *
+   * \since Version 1.23.
+   */
+  ORT_API2_STATUS(AllocatorGetStats, _In_ const OrtAllocator* ort_allocator, _Outptr_ OrtKeyValuePairs** out);
 };
 
 /*

include/onnxruntime/core/session/onnxruntime_cxx_api.h

@@ -2154,6 +2154,13 @@ struct AllocatorImpl : Base<T> {
   MemoryAllocation GetAllocation(size_t size);
   void Free(void* p);
   ConstMemoryInfo GetInfo() const;
+
+  /** \brief Function that returns the statistics of the allocator.
+   *
+   * \param[out] stats: A pointer to a KeyValuePairs object that will be filled with the allocator statistics.
+   * \return A Status indicating success or failure.
+   */
+  Status GetStats(KeyValuePairs* stats) const;
 };
 
 }  // namespace detail

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -243,6 +243,13 @@ inline ConstMemoryInfo AllocatorImpl<T>::GetInfo() const {
   return ConstMemoryInfo{out};
 }
 
+template <typename T>
+inline Status AllocatorImpl<T>::GetStats(KeyValuePairs* stats) const {
+  OrtKeyValuePairs* out;
+  ORT_CXX_RETURN_ON_API_FAIL(GetApi().AllocatorGetStats(this->p_, &out));
+  *stats = KeyValuePairs(out);
+  return Ort::Status();
+}
 }  // namespace detail
 
 inline AllocatorWithDefaultOptions::AllocatorWithDefaultOptions() {

onnxruntime/core/framework/allocator_stats.h

@@ -51,4 +51,4 @@ struct AllocatorStats {
     return ss.str();
   }
 };
-}  // namespace onnxruntime
+}  // namespace onnxruntime

onnxruntime/core/session/allocator_adapters.cc

@@ -3,6 +3,7 @@
 
 #include "allocator_adapters.h"
 #include "core/framework/error_code_helper.h"
+#include "core/session/abi_key_value_pairs.h"
 #include "core/session/inference_session.h"
 #include "core/session/ort_env.h"
 #include "core/session/ort_apis.h"
@@ -11,6 +12,7 @@ namespace onnxruntime {
 
 namespace {
 constexpr uint32_t kOrtAllocatorReserveMinVersion = 18;
+constexpr uint32_t kOrtAllocatorStatsMinVersion = 23;
 }  // namespace
 
 OrtAllocatorImplWrappingIAllocator::OrtAllocatorImplWrappingIAllocator(onnxruntime::AllocatorPtr&& i_allocator)
@@ -26,6 +28,18 @@ OrtAllocatorImplWrappingIAllocator::OrtAllocatorImplWrappingIAllocator(onnxrunti
     OrtAllocator::Reserve =
         [](OrtAllocator* this_, size_t size) { return static_cast<OrtAllocatorImplWrappingIAllocator*>(this_)->Reserve(size); };
   }
+  if (OrtAllocator::version >= kOrtAllocatorStatsMinVersion) {
+    OrtAllocator::GetStats =
+        [](const OrtAllocator* this_, OrtKeyValuePairs** stats) noexcept -> OrtStatusPtr {
+      API_IMPL_BEGIN
+      auto kvp = std::make_unique<OrtKeyValuePairs>();
+      auto stats_map = static_cast<const OrtAllocatorImplWrappingIAllocator*>(this_)->Stats();
+      kvp->Copy(stats_map);
+      *stats = reinterpret_cast<OrtKeyValuePairs*>(kvp.release());
+      return nullptr;
+      API_IMPL_END
+    };
+  }
 }
 
 void* OrtAllocatorImplWrappingIAllocator::Alloc(size_t size) {
@@ -44,6 +58,26 @@ const OrtMemoryInfo* OrtAllocatorImplWrappingIAllocator::Info() const {
   return &i_allocator_->Info();
 }
 
+std::unordered_map<std::string, std::string> OrtAllocatorImplWrappingIAllocator::Stats() const {
+  AllocatorStats stats{};
+  i_allocator_->GetStats(&stats);
+
+  // Allocators which does not implement GetStats() will return empty stats
+  std::unordered_map<std::string, std::string> entries;
+  if (stats.num_allocs > 0 || stats.bytes_limit != 0) {
+    entries.insert_or_assign("Limit", std::to_string(stats.bytes_limit));
+    entries.insert_or_assign("InUse", std::to_string(stats.bytes_in_use));
+    entries.insert_or_assign("TotalAllocated", std::to_string(stats.total_allocated_bytes));
+    entries.insert_or_assign("MaxInUse", std::to_string(stats.max_bytes_in_use));
+    entries.insert_or_assign("NumAllocs", std::to_string(stats.num_allocs));
+    entries.insert_or_assign("NumReserves", std::to_string(stats.num_reserves));
+    entries.insert_or_assign("NumArenaExtensions", std::to_string(stats.num_arena_extensions));
+    entries.insert_or_assign("NumArenaShrinkages", std::to_string(stats.num_arena_shrinkages));
+    entries.insert_or_assign("MaxAllocSize", std::to_string(stats.max_alloc_size));
+  }
+  return entries;
+}
+
 onnxruntime::AllocatorPtr OrtAllocatorImplWrappingIAllocator::GetWrappedIAllocator() {
   return i_allocator_;
 }

onnxruntime/core/session/allocator_adapters.h

@@ -6,6 +6,8 @@
 #include "core/framework/allocator.h"
 #include "core/session/onnxruntime_cxx_api.h"
 
+#include <string>
+
 namespace onnxruntime {
 
 // Since all allocators are of type 'OrtAllocator' and there is a single
@@ -29,6 +31,8 @@
   const OrtMemoryInfo* Info() const;
   void* Reserve(size_t size);
 
+  std::unordered_map<std::string, std::string> Stats() const;
+
   ORT_DISALLOW_COPY_AND_ASSIGNMENT(OrtAllocatorImplWrappingIAllocator);
 
   onnxruntime::AllocatorPtr GetWrappedIAllocator();

onnxruntime/core/session/default_cpu_allocator_c_api.cc

@@ -3,6 +3,7 @@
 
 #include "core/framework/utils.h"
 #include "core/framework/allocator.h"
+#include "core/session/abi_key_value_pairs.h"
 #include "core/session/allocator_adapters.h"
 #include "core/session/onnxruntime_cxx_api.h"
 #include "core/session/ort_apis.h"
@@ -17,6 +18,15 @@
         [](OrtAllocator* this_, void* p) { static_cast<OrtDefaultCpuAllocator*>(this_)->Free(p); };
     OrtAllocator::Info =
         [](const OrtAllocator* this_) { return static_cast<const OrtDefaultCpuAllocator*>(this_)->Info(); };
+    OrtAllocator::Reserve =
+        [](OrtAllocator* this_, size_t size) { return static_cast<OrtDefaultCpuAllocator*>(this_)->Alloc(size); };
+    OrtAllocator::GetStats =
+        [](const OrtAllocator* /*this_*/, OrtKeyValuePairs** stats) noexcept -> OrtStatusPtr {
+      // Default allocator does not support stats, return an empty OrtKeyValuePairs.
+      auto kvp = std::make_unique<OrtKeyValuePairs>();
+      *stats = reinterpret_cast<OrtKeyValuePairs*>(kvp.release());
+      return nullptr;
+    };
     Ort::ThrowOnError(OrtApis::CreateCpuMemoryInfo(OrtDeviceAllocator, OrtMemTypeDefault, &cpu_memory_info));
   }
 

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -1567,6 +1567,12 @@ ORT_API_STATUS_IMPL(OrtApis::AllocatorGetInfo, _In_ const OrtAllocator* ptr, _Ou
   API_IMPL_END
 }
 
+ORT_API_STATUS_IMPL(OrtApis::AllocatorGetStats, _In_ const OrtAllocator* ptr, _Outptr_ OrtKeyValuePairs** out) {
+  API_IMPL_BEGIN
+  return ptr->GetStats(ptr, out);
+  API_IMPL_END
+}
+
 template <typename T>
 ORT_STATUS_PTR OrtGetNumSequenceElements(const OrtValue* p_ml_value, size_t* out) {
   auto& data = p_ml_value->Get<T>();
@@ -3024,6 +3030,7 @@ static constexpr OrtApi ort_api_1_to_23 = {
     &OrtApis::GetEpApi,
     // End of Version 22 - DO NOT MODIFY ABOVE (see above text for more information)
     &OrtApis::GetTensorSizeInBytes,
+    &OrtApis::AllocatorGetStats,
 };
 
 // OrtApiBase can never change as there is no way to know what version of OrtApiBase is returned by OrtGetApiBase.

onnxruntime/core/session/ort_apis.h

@@ -601,4 +601,5 @@ ORT_API(const OrtEpApi*, GetEpApi);
 
 ORT_API_STATUS_IMPL(GetTensorSizeInBytes, _In_ const OrtValue* ort_value, _Out_ size_t* size);
 
+ORT_API_STATUS_IMPL(AllocatorGetStats, _In_ const OrtAllocator* ptr, _Outptr_ OrtKeyValuePairs** out);
 }  // namespace OrtApis

onnxruntime/test/shared_lib/test_allocator.cc

@@ -31,4 +31,10 @@ TEST(CApiTest, DefaultAllocator) {
   ASSERT_EQ(allocation.size(), 100U);
   ASSERT_NE(allocation.get(), nullptr);
   memset(allocation.get(), 0, 100U);
+
+  // Default Allocator does not implement GetStats, we expect the stats to be empty.
+  Ort::KeyValuePairs stats{};
+  auto status = default_allocator.GetStats(&stats);
+  ASSERT_TRUE(status.IsOK());
+  ASSERT_EQ(0, stats.GetKeyValuePairs().size());
 }

onnxruntime/test/shared_lib/test_inference.cc

@@ -1992,6 +1992,31 @@ TEST(CApiTest, get_allocator_cpu) {
   auto mem_allocation = cpu_allocator.GetAllocation(1024);
   ASSERT_NE(nullptr, mem_allocation.get());
   ASSERT_EQ(1024U, mem_allocation.size());
+
+  Ort::KeyValuePairs stats;
+  auto status = cpu_allocator.GetStats(&stats);
+
+  // CPU allocator may not support arena usage.
+  // See func DoesCpuAllocatorSupportArenaUsage() in allocator_utils.cc.
+  if (allocator_info.GetAllocatorType() == OrtAllocatorType::OrtArenaAllocator) {
+    ASSERT_TRUE(status.IsOK());
+
+    ASSERT_EQ("-1", std::string(stats.GetValue("Limit")));
+    ASSERT_EQ("1024", std::string(stats.GetValue("InUse")));
+    ASSERT_EQ("1024", std::string(stats.GetValue("MaxInUse")));
+    ASSERT_EQ("1024", std::string(stats.GetValue("MaxAllocSize")));
+    ASSERT_EQ("2", std::string(stats.GetValue("NumAllocs")));
+    ASSERT_EQ("0", std::string(stats.GetValue("NumReserves")));
+
+    // We don't check values of the following stats
+    ASSERT_NE(nullptr, stats.GetValue("TotalAllocated"));
+    ASSERT_NE(nullptr, stats.GetValue("NumArenaExtensions"));
+    ASSERT_NE(nullptr, stats.GetValue("NumArenaShrinkages"));
+  } else {
+    // If the allocator is not an arena allocator, we expect the stats to be empty.
+    ASSERT_TRUE(status.IsOK());
+    ASSERT_EQ(0, stats.GetKeyValuePairs().size());
+  }
 }
 
 #ifdef USE_CUDA
@@ -2014,6 +2039,22 @@ TEST(CApiTest, get_allocator_cuda) {
   auto mem_allocation = cuda_allocator.GetAllocation(1024);
   ASSERT_NE(nullptr, mem_allocation.get());
   ASSERT_EQ(1024U, mem_allocation.size());
+
+  Ort::KeyValuePairs stats;
+  auto status = cuda_allocator.GetStats(&stats);
+  ASSERT_TRUE(status.IsOK());
+
+  ASSERT_EQ("-1", std::string(stats.GetValue("Limit")));
+  ASSERT_EQ("1024", std::string(stats.GetValue("InUse")));
+  ASSERT_EQ("1024", std::string(stats.GetValue("MaxInUse")));
+  ASSERT_EQ("1024", std::string(stats.GetValue("MaxAllocSize")));
+  ASSERT_EQ("2", std::string(stats.GetValue("NumAllocs")));
+  ASSERT_EQ("0", std::string(stats.GetValue("NumReserves")));
+
+  // We don't check values of the following stats
+  ASSERT_NE(nullptr, stats.GetValue("TotalAllocated"));
+  ASSERT_NE(nullptr, stats.GetValue("NumArenaExtensions"));
+  ASSERT_NE(nullptr, stats.GetValue("NumArenaShrinkages"));
 }
 #endif