Arm backend: Track target memory usage (#5341)

Summary:
This adds memory stats API to the Memory Allocator and make it possible for the backends to track and report runtime memory usage.

Pull Request resolved: #5341

Reviewed By: manuelcandales

Differential Revision: D62874650

Pulled By: digantdesai

fbshipit-source-id: 9b737d973466544c2df491236ed5e6a14c57eb40

Author: zingo

examples/arm/executor_runner/arm_executor_runner.cpp

@@ -28,7 +28,7 @@
 // needs to be large enough to take an entire model. On the FVP,
 // network_model_sec is linked to the DDR, which is large (256MB on
 // Corstone-300).
-const size_t input_allocation_pool_size = 100 * 1024 * 1024;
+const size_t input_allocation_pool_size = 60 * 1024 * 1024;
 unsigned char __attribute__((
     section("network_model_sec"),
     aligned(16))) input_allocation_pool[input_allocation_pool_size];
@@ -63,7 +63,7 @@ using executorch::runtime::Span;
 using executorch::runtime::Tag;
 using executorch::runtime::TensorInfo;
 
-#define METHOD_ALLOCATOR_POOL_SIZE (70 * 1024 * 1024)
+#define METHOD_ALLOCATOR_POOL_SIZE (60 * 1024 * 1024)
 unsigned char __attribute__((
     section("network_model_sec"),
     aligned(16))) method_allocation_pool[METHOD_ALLOCATOR_POOL_SIZE];
@@ -326,6 +326,8 @@ int main(int argc, const char* argv[]) {
   std::vector<Span<uint8_t>> planned_spans; // Passed to the allocator
   size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
 
+  size_t planned_buffer_membase = method_allocator.used_size();
+
   for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
     size_t buffer_size =
         static_cast<size_t>(method_meta->memory_planned_buffer_size(id).get());
@@ -337,6 +339,8 @@ int main(int argc, const char* argv[]) {
     planned_buffers.push_back(buffer);
     planned_spans.push_back({planned_buffers.back(), buffer_size});
   }
+  size_t planned_buffer_memsize =
+      method_allocator.used_size() - planned_buffer_membase;
 
   HierarchicalAllocator planned_memory(
       {planned_spans.data(), planned_spans.size()});
@@ -347,6 +351,7 @@ int main(int argc, const char* argv[]) {
   MemoryManager memory_manager(
       &method_allocator, &planned_memory, &temp_allocator);
 
+  size_t method_loaded_membase = method_allocator.used_size();
   Result<Method> method = program->load_method(method_name, &memory_manager);
   if (!method.ok()) {
     ET_LOG(
@@ -355,10 +360,12 @@ int main(int argc, const char* argv[]) {
         method_name,
         method.error());
   }
+  size_t method_loaded_memsize =
+      method_allocator.used_size() - method_loaded_membase;
   ET_LOG(Info, "Method loaded.");
 
   ET_LOG(Info, "Preparing inputs...");
-
+  size_t input_membase = method_allocator.used_size();
   auto inputs =
       ::prepare_input_tensors(*method, method_allocator, input_buffers);
 
@@ -369,12 +376,15 @@ int main(int argc, const char* argv[]) {
         method_name,
         inputs.error());
   }
+  size_t input_memsize = method_allocator.used_size() - input_membase;
   ET_LOG(Info, "Input prepared.");
 
   ET_LOG(Info, "Starting the model execution...");
+  size_t executor_membase = method_allocator.used_size();
   StartMeasurements();
   Error status = method->execute();
   StopMeasurements();
+  size_t executor_memsize = method_allocator.used_size() - executor_membase;
 
   if (status != Error::Ok) {
     ET_LOG(
@@ -425,6 +435,25 @@ int main(int argc, const char* argv[]) {
   }
 out:
   ET_LOG(Info, "Program complete, exiting.");
+  if (method_allocator.size() != 0) {
+    size_t method_allocator_used = method_allocator.used_size();
+    ET_LOG(
+        Info,
+        "Method allocator area ( method_allocator_planned: %zu method_allocator_loaded: %zu method_allocator_input: %zu method_allocator_executor: %zu ) total: %zu",
+        planned_buffer_memsize,
+        method_loaded_memsize,
+        input_memsize,
+        executor_memsize,
+        method_allocator_used);
+    ET_LOG(
+        Info,
+        "Method allocator area method_allocator_used: %d / method_allocator_size: %d  method_allocator_free: %d ( used: %d %% ) ",
+        method_allocator_used,
+        method_allocator.size(),
+        method_allocator.free_size(),
+        100 * method_allocator_used / method_allocator.size());
+  }
+
 #ifdef SEMIHOSTING
   _exit(0);
 #endif

runtime/core/memory_allocator.h

@@ -1,6 +1,7 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
+ * Copyright 2024 Arm Limited and/or its affiliates.
  *
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
@@ -150,6 +151,16 @@ class MemoryAllocator {
     return size_;
   }
 
+  // Returns the used size of the allocator's memory buffer.
+  virtual uint32_t used_size() const {
+    return cur_ - begin_;
+  }
+
+  // Returns the free size of the allocator's memory buffer.
+  virtual uint32_t free_size() const {
+    return end_ - cur_;
+  }
+
   // Resets the current pointer to the base address. It does nothing to
   // the contents.
   virtual void reset() {