Migrate relocatable_runner to use MethodMeta (#316)

Summary:
Pull Request resolved: #316

Make this use the new APIs. Note that MethodMeta::num_non_const_buffers is zero-based, so we can avoid the hacks that skip index zero.

Reviewed By: JacobSzwejbka

Differential Revision: D48762565

fbshipit-source-id: fcb70ca0d85eb6c37438c31278d348fc996decdc

Author: dbort

test/relocatable_runner.cpp

@@ -29,10 +29,10 @@ using namespace torch::executor;
  * power down and then back up) in between two inference requests.
  *
  * For ExecuTorch to work efficiently in these environments, we want to
- * initialize the execution plan once once for the model and avoid
- * re-initializing it for every inference. This can be achieved by restricting
- * the runtime contexts (torch::executor::Program and torch::executor::Method)
- * to live in a pre-allocated, shared, and persistent memory.
+ * initialize the Method once once for the model and avoid re-initializing it
+ * for every inference. This can be achieved by restricting the runtime contexts
+ * (torch::executor::Program and torch::executor::Method) to live in a
+ * pre-allocated, shared, and persistent memory.
  *
  * This tool demonstrates that the memory can be managed this way.
  */
@@ -79,8 +79,7 @@ Program* load_program(
 }
 
 MemoryManager* create_memory_manager(
-    Program* program,
-    const char* method_name,
+    MethodMeta* method_meta,
     MemoryAllocator& worker_allocator) {
   // Create the runtime allocator.
   auto* runtime_allocator =
@@ -89,18 +88,16 @@ MemoryManager* create_memory_manager(
   new (runtime_allocator) MemoryAllocator(sizeof(runtime_pool), runtime_pool);
 
   // Create the non-const allocator and the buffers it points to.
-  size_t num_non_const_buffers =
-      program->num_non_const_buffers(method_name).get();
+  size_t num_non_const_buffers = method_meta->num_non_const_buffers();
   MemoryAllocator* non_const_allocators =
-      worker_allocator.allocateList<MemoryAllocator>(num_non_const_buffers - 1);
-  for (size_t id = 1; id < num_non_const_buffers; ++id) {
-    const size_t buffer_size =
-        program->get_non_const_buffer_size(id, method_name).get();
+      worker_allocator.allocateList<MemoryAllocator>(num_non_const_buffers);
+  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 - 1])
+    new (&non_const_allocators[id])
         MemoryAllocator(buffer_size, (uint8_t*)buffer);
     ET_LOG(
         Info,
@@ -112,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 - 1, non_const_allocators);
+      HierarchicalAllocator(num_non_const_buffers, non_const_allocators);
 
   // The constant allocator is not currently used, but must be provided.
   auto* const_allocator = worker_allocator.allocateInstance<MemoryAllocator>();
@@ -140,8 +137,11 @@ Method* init_method(
     MemoryAllocator& worker_allocator,
     std::vector<size_t>& input_sizes,
     std::vector<size_t>& output_sizes) {
+  Result<MethodMeta> method_meta = program->method_meta(method_name);
+  ET_CHECK(method_meta.ok());
+
   MemoryManager* memory_manager =
-      create_memory_manager(program, method_name, worker_allocator);
+      create_memory_manager(&method_meta.get(), worker_allocator);
 
   //
   // Create and load a method from the program, using the provided
@@ -227,7 +227,7 @@ void inference_loop(
   Error status = method->execute();
   ET_CHECK_MSG(
       status == Error::Ok,
-      "plan->execute() failed with status 0x%" PRIx32,
+      "method->execute() failed with status 0x%" PRIx32,
       status);
   ET_LOG(Info, "Model executed successfully.");
 }
@@ -285,8 +285,7 @@ int main(int argc, char** argv) {
   const char* method_name = nullptr;
   {
     // Use the first method in the program.
-    const size_t plan_index = 0;
-    const auto method_name_result = program->get_method_name(plan_index);
+    const auto method_name_result = program->get_method_name(0);
     ET_CHECK_MSG(method_name_result.ok(), "Program has no methods");
     method_name = *method_name_result;
   }