Extend OrtAllocator API to get Allocator statistics (microsoft#24785)

### Description
<!-- Describe your changes. -->
Extend IAllocator to get Allocator statistics:
- Add `OrtAllocator::GetStats` and `AllocatorGetStats` C-API.
- Add `Ort::Allocator::GetStats` Cxx API to parse the string and return
as map.
- Add UT.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Our system integrates multiple models for inference, each with varying
memory demands. Providing a mechanism to retrieve detailed memory
statistics would be useful for analyzing memory usage across models and
devices more effectively.

Author: toothache

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 = {};
+  }
 
   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,26 @@ 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
+   * and the user should call OrtApi::ReleaseKeyValuePairs.
+   * 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)(
@@ -5277,6 +5297,22 @@ 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
+   * and the user should call OrtApi::ReleaseKeyValuePairs.
+   *
+   * 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

@@ -2155,6 +2155,12 @@ 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.
+   *
+   * \return A pointer to a KeyValuePairs object that will be filled with the allocator statistics.
+   */
+  KeyValuePairs GetStats() const;
 };
 
 }  // namespace detail

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

@@ -243,6 +243,12 @@ inline ConstMemoryInfo AllocatorImpl<T>::GetInfo() const {
   return ConstMemoryInfo{out};
 }
 
+template <typename T>
+inline KeyValuePairs AllocatorImpl<T>::GetStats() const {
+  OrtKeyValuePairs* out;
+  ThrowOnError(GetApi().AllocatorGetStats(this->p_, &out));
+  return KeyValuePairs(out);
+}
 }  // 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"
@@ -26,6 +27,16 @@ OrtAllocatorImplWrappingIAllocator::OrtAllocatorImplWrappingIAllocator(onnxrunti
     OrtAllocator::Reserve =
         [](OrtAllocator* this_, size_t size) { return static_cast<OrtAllocatorImplWrappingIAllocator*>(this_)->Reserve(size); };
   }
+  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 +55,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 @@ struct OrtAllocatorImplWrappingIAllocator final : public OrtAllocatorImpl {
   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 @@ struct OrtDefaultCpuAllocator : onnxruntime::OrtAllocatorImpl {
         [](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