Arm backend: Track target memory usage

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

Signed-off-by: Zingo Andersen <zingo.andersen@arm.com>
Change-Id: I73d3531f2430edaa91608a1b06dee4aedf79c5a3

examples/arm/executor_runner/arm_executor_runner.cpp

@@ -313,6 +313,8 @@ int main(int argc, const char* argv[]) {
       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());
@@ -324,6 +326,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;
 
   torch::executor::HierarchicalAllocator planned_memory(
       {planned_spans.data(), planned_spans.size()});
@@ -334,6 +338,7 @@ int main(int argc, const char* argv[]) {
   torch::executor::MemoryManager memory_manager(
       &method_allocator, &planned_memory, &temp_allocator);
 
+  size_t method_loaded_membase = method_allocator.used_size();
   Result<torch::executor::Method> method =
       program->load_method(method_name, &memory_manager);
   if (!method.ok()) {
@@ -343,10 +348,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);
 
@@ -357,12 +364,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(
@@ -413,6 +423,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() {