support multimethod by method name runtime

Differential Revision: https://internalfb.com/D49335926

fbshipit-source-id: 2eb55cd9ee40d82a4ac33aaa10950c9044192c80

Author: Songhao Jia

bundled_program/tests/common.py

@@ -12,7 +12,6 @@
 import executorch.exir as exir
 import torch
 from executorch.bundled_program.config import BundledConfig
-from executorch.exir import CaptureConfig
 from executorch.exir.schema import Program
 
 # @manual=fbsource//third-party/pypi/typing-extensions:typing-extensions

sdk/runners/executor_runner.cpp

@@ -75,6 +75,11 @@ DEFINE_string(
 
 DEFINE_int32(num_threads, 1, "Number of threads to use.");
 
+DEFINE_string(
+    method_name,
+    "forward",
+    "Name of method to run. Only used by bundled program mode.");
+
 DEFINE_int32(
     testset_idx,
     0,
@@ -347,7 +352,7 @@ int main(int argc, char** argv) {
           *method,
           program_data.bundled_program_data(),
           &bundled_input_allocator,
-          method_index,
+          method_name,
           FLAGS_testset_idx);
       ET_CHECK_MSG(
           status == Error::Ok,
@@ -376,7 +381,7 @@ int main(int argc, char** argv) {
           *method,
           program_data.bundled_program_data(),
           &bundled_input_allocator,
-          method_index,
+          method_name,
           FLAGS_testset_idx,
           FLAGS_rtol,
           FLAGS_atol);

util/bundled_program_verification.cpp

@@ -166,28 +166,45 @@ bool tensors_are_close(
   }
 }
 
+Result<executorch_flatbuffer::BundledExecutionPlanTest*> get_method_test(
+    const executorch_flatbuffer::BundledProgram* bundled_program,
+    const char* method_name) {
+  auto method_tests = bundled_program->execution_plan_tests();
+  for (size_t i = 0; i < method_tests->size(); i++) {
+    auto m_test = method_tests->GetMutableObject(i);
+    if (std::strcmp(m_test->method_name()->c_str(), method_name) == 0) {
+      return m_test;
+    }
+  }
+  ET_LOG(Error, "No method named '%s' in given bundled program", method_name);
+  return Error::InvalidArgument;
+}
+
 } // namespace
 
 // Load testset_idx-th bundled data into the Method
 __ET_NODISCARD Error LoadBundledInput(
     Method& method,
     serialized_bundled_program* bundled_program_ptr,
     MemoryAllocator* memory_allocator,
-    size_t method_idx,
+    const char* method_name,
     size_t testset_idx) {
   ET_CHECK_OR_RETURN_ERROR(
       executorch_flatbuffer::BundledProgramBufferHasIdentifier(
           bundled_program_ptr),
       NotSupported,
       "The input buffer should be a bundled program.");
 
+  auto method_test = get_method_test(
+      executorch_flatbuffer::GetBundledProgram(bundled_program_ptr),
+      method_name);
+
+  if (!method_test.ok()) {
+    return method_test.error();
+  }
+
   auto bundled_inputs =
-      executorch_flatbuffer::GetBundledProgram(bundled_program_ptr)
-          ->execution_plan_tests()
-          ->Get(method_idx)
-          ->test_sets()
-          ->Get(testset_idx)
-          ->inputs();
+      method_test.get()->test_sets()->Get(testset_idx)->inputs();
 
   for (size_t input_idx = 0; input_idx < method.inputs_size(); input_idx++) {
     auto bundled_input = bundled_inputs->GetMutableObject(input_idx);
@@ -263,7 +280,7 @@ __ET_NODISCARD Error VerifyResultWithBundledExpectedOutput(
     Method& method,
     serialized_bundled_program* bundled_program_ptr,
     MemoryAllocator* memory_allocator,
-    size_t method_idx,
+    const char* method_name,
     size_t testset_idx,
     double rtol,
     double atol) {
@@ -273,13 +290,16 @@ __ET_NODISCARD Error VerifyResultWithBundledExpectedOutput(
       NotSupported,
       "The input buffer should be a bundled program.");
 
+  auto method_test = get_method_test(
+      executorch_flatbuffer::GetBundledProgram(bundled_program_ptr),
+      method_name);
+
+  if (!method_test.ok()) {
+    return method_test.error();
+  }
+
   auto bundled_expected_outputs =
-      executorch_flatbuffer::GetBundledProgram(bundled_program_ptr)
-          ->execution_plan_tests()
-          ->Get(method_idx)
-          ->test_sets()
-          ->Get(testset_idx)
-          ->expected_outputs();
+      method_test.get()->test_sets()->Get(testset_idx)->expected_outputs();
 
   for (size_t output_idx = 0; output_idx < method.outputs_size();
        output_idx++) {

util/bundled_program_verification.h

@@ -26,7 +26,7 @@ using serialized_bundled_program = const void;
  *
  * @param[in] method The Method to verify.
  * @param[in] bundled_program_ptr The bundled program contains expected output.
- * @param[in] method_idx  The index of the Method being verified.
+ * @param[in] method_name  The name of the Method being verified.
  * @param[in] testset_idx  The index of input needs to be set into given Method.
  *
  * @returns Return Error::Ok if load successfully, or the error happens during
@@ -36,7 +36,7 @@ __ET_NODISCARD Error LoadBundledInput(
     Method& method,
     serialized_bundled_program* bundled_program_ptr,
     MemoryAllocator* memory_allocator,
-    size_t method_idx,
+    const char* method_name,
     size_t testset_idx);
 
 /**
@@ -45,7 +45,7 @@ __ET_NODISCARD Error LoadBundledInput(
  *
  * @param[in] method The Method to extract outputs from.
  * @param[in] bundled_program_ptr The bundled program contains expected output.
- * @param[in] method_idx  The index of the Method being verified.
+ * @param[in] method_name  The name of the Method being verified.
  * @param[in] testset_idx  The index of expected output needs to be compared.
  * @param[in] rtol Relative tolerance used for data comparsion.
  * @param[in] atol Absolute tolerance used for data comparsion.
@@ -57,7 +57,7 @@ __ET_NODISCARD Error VerifyResultWithBundledExpectedOutput(
     Method& method,
     serialized_bundled_program* bundled_program_ptr,
     MemoryAllocator* memory_allocator,
-    size_t method_idx,
+    const char* method_name,
     size_t testset_idx,
     double rtol = 1e-5,
     double atol = 1e-8);