Migrate runner-like targets to use the new HierarchicalAllocator span ctor (#388)

Summary:
Pull Request resolved: #388

Stop using the deprecated HierarchicalAllocator ctor.
ghstack-source-id: 201128788
exported-using-ghexport

Reviewed By: tarun292

Differential Revision: D49344926

fbshipit-source-id: 860f67a679ab7cf2bc725defe1a3455f89835947

Author: dbort

examples/bundled_executor_runner/bundled_executor_runner.cpp

@@ -154,26 +154,27 @@ int main(int argc, char** argv) {
       MemoryAllocator(kRuntimeMemorySize, runtime_pool)};
   runtime_allocator.enable_profiling("runtime allocator");
 
-  // The non-const allocator is used to provide the memory-planned buffers that
-  // back mutable tensors. Since it was planned ahead of time, the Program knows
-  // how big each of the allocators needs to be.
+  // The non-const buffers will back the mutable tensors used by the method. The
+  // sizes of these buffers were determined ahead of time during the
+  // memory-planning pasees.
   //
-  // These buffers correspond to different hardware memory banks. Most mobile
-  // environments will only have a single buffer. Some embedded environments may
-  // have more than one for, e.g., slow/large DRAM and fast/small SRAM.
+  // Each buffer typically corresponds to a different hardware memory bank. Most
+  // mobile environments will only have a single buffer. Some embedded
+  // environments may have more than one for, e.g., slow/large DRAM and
+  // fast/small SRAM, or for memory associated with particular cores.
   std::vector<std::unique_ptr<uint8_t[]>> non_const_buffers;
-  std::vector<MemoryAllocator> non_const_allocators;
+  std::vector<Span<uint8_t>> non_const_spans;
   size_t num_non_const_buffers = method_meta->num_non_const_buffers();
   for (size_t id = 0; id < num_non_const_buffers; ++id) {
-    size_t buffer_size = method_meta->non_const_buffer_size(id).get();
+    // .get() will always succeed because id < num_non_const_buffers.
+    size_t buffer_size =
+        static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
     ET_LOG(Info, "Setting up non-const buffer %zu, size %zu.", id, buffer_size);
     non_const_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
-    non_const_allocators.push_back(
-        MemoryAllocator(buffer_size, non_const_buffers.back().get()));
-    non_const_allocators.back().enable_profiling("non_const_allocators");
+    non_const_spans.push_back({non_const_buffers.back().get(), buffer_size});
   }
   HierarchicalAllocator non_const_allocator(
-      non_const_allocators.size(), non_const_allocators.data());
+      {non_const_spans.data(), non_const_spans.size()});
 
   // Allocator for bundled input.
   MemoryAllocator bundled_input_allocator{

examples/executor_runner/executor_runner.cpp

@@ -114,33 +114,27 @@ int main(int argc, char** argv) {
       MemoryAllocator(kRuntimeMemorySize, runtime_pool)};
   runtime_allocator.enable_profiling("runtime allocator");
 
-  // The non-const allocator is used to provide the memory-planned buffers that
-  // back mutable tensors. Since it was planned ahead of time, the Program knows
-  // how big each of the allocators needs to be.
+  // The non-const buffers will back the mutable tensors used by the method. The
+  // sizes of these buffers were determined ahead of time during the
+  // memory-planning pasees.
   //
-  // These buffers correspond to different hardware memory banks. Most mobile
-  // environments will only have a single buffer. Some embedded environments may
-  // have more than one for, e.g., slow/large DRAM and fast/small SRAM.
+  // Each buffer typically corresponds to a different hardware memory bank. Most
+  // mobile environments will only have a single buffer. Some embedded
+  // environments may have more than one for, e.g., slow/large DRAM and
+  // fast/small SRAM, or for memory associated with particular cores.
   std::vector<std::unique_ptr<uint8_t[]>> non_const_buffers;
-  std::vector<MemoryAllocator> non_const_allocators;
+  std::vector<Span<uint8_t>> non_const_spans;
   size_t num_non_const_buffers = method_meta->num_non_const_buffers();
   for (size_t id = 0; id < num_non_const_buffers; ++id) {
-    auto buffer_size = method_meta->non_const_buffer_size(id);
-    ET_CHECK_MSG(
-        buffer_size.ok(),
-        "Failed to get size of non-const buffer %zu for method %s: 0x%x",
-        id,
-        method_name,
-        (unsigned int)buffer_size.error());
-    ET_LOG(
-        Info, "Setting up non-const buffer %zu, size %zu.", id, *buffer_size);
-    non_const_buffers.push_back(std::make_unique<uint8_t[]>(*buffer_size));
-    non_const_allocators.push_back(
-        MemoryAllocator(*buffer_size, non_const_buffers.back().get()));
-    non_const_allocators.back().enable_profiling("non_const_allocators");
+    // .get() will always succeed because id < num_non_const_buffers.
+    size_t buffer_size =
+        static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
+    ET_LOG(Info, "Setting up non-const buffer %zu, size %zu.", id, buffer_size);
+    non_const_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
+    non_const_spans.push_back({non_const_buffers.back().get(), buffer_size});
   }
   HierarchicalAllocator non_const_allocator(
-      non_const_allocators.size(), non_const_allocators.data());
+      {non_const_spans.data(), non_const_spans.size()});
 
   // The constant allocator is not currently used. Please initialize with a
   // zero-sized allocator.

exir/backend/test/demos/rpc/ExecutorBackend.cpp

@@ -79,24 +79,21 @@ class ExecutorBackend final : public PyTorchBackendInterface {
 
     auto num_non_const_buffers = method_meta->num_non_const_buffers();
 
-    uint8_t** non_const_buffers = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
-        runtime_allocator, uint8_t*, num_non_const_buffers);
-    MemoryAllocator* non_const_allocators = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
-        runtime_allocator, MemoryAllocator, num_non_const_buffers);
+    Span<uint8_t>* non_const_buffers = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
+        runtime_allocator, Span<uint8_t>, num_non_const_buffers);
 
     for (size_t id = 0; id < num_non_const_buffers; ++id) {
-      auto buffer_size = method_meta->non_const_buffer_size(id);
+      size_t buffer_size =
+          static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
       uint8_t* buffer_i = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
-          runtime_allocator, uint8_t, buffer_size.get());
-      non_const_buffers[id] = buffer_i;
-      new (&non_const_allocators[id])
-          MemoryAllocator(static_cast<uint32_t>(buffer_size.get()), buffer_i);
+          runtime_allocator, uint8_t, buffer_size);
+      non_const_buffers[id] = {buffer_i, buffer_size};
     }
 
     auto client_non_const_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
         runtime_allocator, HierarchicalAllocator);
     new (client_non_const_allocator)
-        HierarchicalAllocator(num_non_const_buffers, non_const_allocators);
+        HierarchicalAllocator({non_const_buffers, num_non_const_buffers});
 
     // Allocate some memory from runtime allocator for the client executor, in
     // real case, like if it's an executor in dsp, it should allocate memory

sdk/runners/executor_runner.cpp

@@ -264,26 +264,27 @@ int main(int argc, char** argv) {
       MemoryAllocator(kRuntimeMemorySize, runtime_pool)};
   runtime_allocator.enable_profiling("runtime allocator");
 
-  // The non-const allocator is used to provide the memory-planned buffers that
-  // back mutable tensors. Since it was planned ahead of time, the Program knows
-  // how big each of the allocators needs to be.
+  // The non-const buffers will back the mutable tensors used by the method. The
+  // sizes of these buffers were determined ahead of time during the
+  // memory-planning pasees.
   //
-  // These buffers correspond to different hardware memory banks. Most mobile
-  // environments will only have a single buffer. Some embedded environments may
-  // have more than one for, e.g., slow/large DRAM and fast/small SRAM.
+  // Each buffer typically corresponds to a different hardware memory bank. Most
+  // mobile environments will only have a single buffer. Some embedded
+  // environments may have more than one for, e.g., slow/large DRAM and
+  // fast/small SRAM, or for memory associated with particular cores.
   std::vector<std::unique_ptr<uint8_t[]>> non_const_buffers;
-  std::vector<MemoryAllocator> non_const_allocators;
+  std::vector<Span<uint8_t>> non_const_spans;
   size_t num_non_const_buffers = method_meta->num_non_const_buffers();
   for (size_t id = 0; id < num_non_const_buffers; ++id) {
-    size_t buffer_size = method_meta->non_const_buffer_size(id).get();
+    // .get() will always succeed because id < num_non_const_buffers.
+    size_t buffer_size =
+        static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
     ET_LOG(Info, "Setting up non-const buffer %zu, size %zu.", id, buffer_size);
     non_const_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
-    non_const_allocators.push_back(
-        MemoryAllocator(buffer_size, non_const_buffers.back().get()));
-    non_const_allocators.back().enable_profiling("non_const_allocators");
+    non_const_spans.push_back({non_const_buffers.back().get(), buffer_size});
   }
   HierarchicalAllocator non_const_allocator(
-      non_const_allocators.size(), non_const_allocators.data());
+      {non_const_spans.data(), non_const_spans.size()});
 
   // The constant allocator is not currently used. Please initialize with a
   // zero-sized allocator.

test/relocatable_runner.cpp

@@ -89,16 +89,16 @@ MemoryManager* create_memory_manager(
 
   // Create the non-const allocator and the buffers it points to.
   size_t num_non_const_buffers = method_meta->num_non_const_buffers();
-  MemoryAllocator* non_const_allocators =
-      worker_allocator.allocateList<MemoryAllocator>(num_non_const_buffers);
+  Span<uint8_t>* non_const_buffers =
+      worker_allocator.allocateList<Span<uint8_t>>(num_non_const_buffers);
+  ET_CHECK(non_const_buffers != nullptr);
   for (size_t id = 0; id < num_non_const_buffers; ++id) {
     const size_t buffer_size = method_meta->non_const_buffer_size(id).get();
     ET_LOG(
         Info, "Setting up non-const buffer id %zu, size %zu.", id, buffer_size);
     void* buffer = worker_allocator.allocate(buffer_size);
     ET_CHECK(buffer != nullptr);
-    new (&non_const_allocators[id])
-        MemoryAllocator(buffer_size, (uint8_t*)buffer);
+    non_const_buffers[id] = {(uint8_t*)buffer, buffer_size};
     ET_LOG(
         Info,
         "Created non_const_allocators with size %zu and addr %p",
@@ -109,7 +109,7 @@ MemoryManager* create_memory_manager(
       worker_allocator.allocateInstance<HierarchicalAllocator>();
   ET_CHECK(non_const_allocator != nullptr);
   new (non_const_allocator)
-      HierarchicalAllocator(num_non_const_buffers, non_const_allocators);
+      HierarchicalAllocator({non_const_buffers, num_non_const_buffers});
 
   // The constant allocator is not currently used, but must be provided.
   auto* const_allocator = worker_allocator.allocateInstance<MemoryAllocator>();

test/size_test.cpp

@@ -39,12 +39,8 @@ int main(int argc, char** argv) {
   MemoryAllocator temp_allocator{MemoryAllocator(0, nullptr)};
   temp_allocator.enable_profiling("temp allocator");
 
-  MemoryAllocator non_const_allocators[1]{
-      MemoryAllocator(kMemoryAmount, activation_pool)};
-  non_const_allocators[0].enable_profiling("non_const_allocators");
-
-  HierarchicalAllocator non_const_allocator{
-      HierarchicalAllocator(1, non_const_allocators)};
+  Span<uint8_t> non_const_buffers[1]{{activation_pool, kMemoryAmount}};
+  HierarchicalAllocator non_const_allocator({non_const_buffers, 1});
 
   MemoryManager memory_manager{MemoryManager(
       &const_allocator,