support multimethod by method name on AOT (#341)

Summary:

Update AOT api and schema to support multimethod by method name instead of method idx.

Since the following diffs will reformat the AOT apis, this diff doesn't pay much effort on documentation stuff.

Differential Revision: D49246787

Author: Songhao Jia

bundled_program/config.py

@@ -47,6 +47,7 @@ class ConfigIOSet:
 class ConfigExecutionPlanTest:
     """All info related to verify execution plan"""
 
+    method_name: str
     test_sets: List[ConfigIOSet]
 
 
@@ -60,6 +61,7 @@ class BundledConfig:
 
     def __init__(
         self,
+        method_names: List[str],
         # pyre-ignore
         inputs: List[List[Any]],
         # pyre-ignore
@@ -68,34 +70,34 @@ def __init__(
         """Contruct the config given inputs and expected outputs
 
         Args:
-            inputs: All sets of input need to be test on for all execution plans. Each list
-                    of `inputs` is all sets which will be run on the execution plan in the
-                    program sharing same index. Each set of any `inputs` element should
+            method_names: All method names need to be verified in program.  Each method name
+            inputs: All sets of input need to be test on for all methods. Each list
+                    of `inputs` is all sets which will be run on the method in the
+                    program with corresponding method name. Each set of any `inputs` element should
                     contain all inputs required by eager_model with the same inference function
                     as corresponding execution plan for one-time execution.
 
-                    Please note that currently we do not have any consensus about the mapping rule
-                    between inference name in eager_model and execution plan id in executorch
-                    program. Hence, user should take care of the data order in `inputs`: each list
-                    of `inputs` is all sets which will be run on the execution plan with same index,
-                    not the inference function with same index in the result of get_inference_name.
-                    Same as the `expected_outputs` and `metadatas` below.
-
-                    It shouldn't be a problem if there's only one inferenece function per model.
-
             expected_outputs: Expected outputs for inputs sharing same index. The size of
-                    expected_outputs should be the same as the size of inputs.
+                    expected_outputs should be the same as the size of inputs and provided method_names.
         """
         BundledConfig._check_io_type(inputs)
         BundledConfig._check_io_type(expected_outputs)
-        assert len(inputs) == len(expected_outputs), (
-            "length of inputs and expected_outputs should match,"
-            + " but got {} and {}".format(len(inputs), len(expected_outputs))
+
+        for m_name in method_names:
+            assert isinstance(m_name, str)
+
+        assert len(method_names) == len(inputs) == len(expected_outputs), (
+            "length of method_names, inputs and expected_outputs should match,"
+            + " but got {}, {} and {}".format(
+                len(method_names), len(inputs), len(expected_outputs)
+            )
         )
 
         self.execution_plan_tests: List[
             ConfigExecutionPlanTest
-        ] = BundledConfig._gen_execution_plan_tests(inputs, expected_outputs)
+        ] = BundledConfig._gen_execution_plan_tests(
+            method_names, inputs, expected_outputs
+        )
 
     @staticmethod
     # TODO(T138930448): Give pyre-ignore commands appropriate warning type and comments.
@@ -145,6 +147,7 @@ def _check_io_type(test_data_program: List[List[Any]]) -> None:
 
     @staticmethod
     def _gen_execution_plan_tests(
+        method_names: List[str],
         # pyre-ignore
         inputs: List[List[Any]],
         # pyre-ignore
@@ -155,9 +158,10 @@ def _gen_execution_plan_tests(
         execution_plan_tests: List[ConfigExecutionPlanTest] = []
 
         for (
+            m_name,
             inputs_per_plan_test,
             expect_outputs_per_plan_test,
-        ) in zip(inputs, expected_outputs):
+        ) in zip(method_names, inputs, expected_outputs):
             test_sets: List[ConfigIOSet] = []
 
             # transfer I/O sets into ConfigIOSet for each execution plan
@@ -182,7 +186,12 @@ def _gen_execution_plan_tests(
 
             execution_plan_tests.append(
                 ConfigExecutionPlanTest(
+                    method_name=m_name,
                     test_sets=test_sets,
                 )
             )
+
+        # sort the execution plan tests by method name to in line with core program emitter.
+        execution_plan_tests.sort(key=lambda x: x.method_name)
+
         return execution_plan_tests

bundled_program/core.py

@@ -136,32 +136,38 @@ def assert_valid_bundle(
 
     """
 
-    # Check the number of execution plan tests
-    assert len(bundled_config.execution_plan_tests) == len(
-        program.execution_plan
-    ), "The length of execution_plan_tests in config should match the length of execution_plan in program, but get {} and {}.".format(
-        len(bundled_config.execution_plan_tests), len(program.execution_plan)
-    )
+    program_plan_id = 0
+    bp_plan_id = 0
 
     # Check if the inputs' type meet Program's requirement
-    for plan_id in range(len(program.execution_plan)):
+    while bp_plan_id < len(bundled_config.execution_plan_tests):
+
         plan_test: ConfigExecutionPlanTest = bundled_config.execution_plan_tests[
-            plan_id
+            bp_plan_id
         ]
+        plan: ExecutionPlan = program.execution_plan[program_plan_id]
 
-        plan: ExecutionPlan = program.execution_plan[plan_id]
+        # User does not provide testcases for current plan, skip it
+        if plan_test.method_name < plan.name:
+            program_plan_id += 1
+            continue
+
+        # Check if the method name in user provided test matches the one in the original program
+        assert (
+            plan_test.method_name == plan.name
+        ), f"BundledConfig has testcases for method {plan_test.method_name}, but can not find it in the given program. All method names in the program are {', '.join([p.name for p in program.execution_plan])}."
 
         # Check if the type of Program's input is supported
         for index in range(len(plan.inputs)):
             assert (
-                type(get_program_input(program, plan_id, index))
+                type(get_program_input(program, program_plan_id, index))
                 in supported_program_type_table
             ), "The type of program's input isn't supported."
 
         # Check if the type of Program's output is supported
         for index in range(len(plan.outputs)):
             assert (
-                type(get_program_output(program, plan_id, index)) == Tensor
+                type(get_program_output(program, program_plan_id, index)) == Tensor
             ), "Only supports program with output in Tensor type."
 
         # Check if the I/O sets of each execution plan test match program's requirement.
@@ -181,14 +187,14 @@ def assert_valid_bundle(
                 assert (
                     type(cur_plan_test_inputs[j])
                     == supported_program_type_table[
-                        type(get_program_input(program, plan_id, j))
+                        type(get_program_input(program, program_plan_id, j))
                     ]
                 ), "The type {}-th input in {}-th test set of {}-th execution plan does not meet Program's requirement: expected {} but get {}".format(
                     j,
                     i,
-                    plan_id,
+                    program_plan_id,
                     supported_program_type_table[
-                        type(get_program_input(program, plan_id, j))
+                        type(get_program_input(program, program_plan_id, j))
                     ],
                     type(cur_plan_test_inputs[j]),
                 )
@@ -198,10 +204,10 @@ def assert_valid_bundle(
                     # pyre-fixme[16]: Undefined attribute [16]: Item `bool` of `typing.Union[bool, float, int, torch._tensor.Tensor]`
                     # has no attribute `dtype`.
                     assert cur_plan_test_inputs[j].dtype == get_input_dtype(
-                        program, plan_id, j
+                        program, program_plan_id, j
                     ), "The input tensor {} dtype shall be {}, but now is {}".format(
                         cur_plan_test_inputs[j],
-                        get_input_dtype(program, plan_id, j),
+                        get_input_dtype(program, program_plan_id, j),
                         cur_plan_test_inputs[j].dtype,
                     )
                 elif type(cur_plan_test_inputs[j]) in (
@@ -210,9 +216,9 @@ def assert_valid_bundle(
                     float,
                 ):
                     assert type(cur_plan_test_inputs[j]) == get_input_type(
-                        program, plan_id, j
+                        program, program_plan_id, j
                     ), "The input primitive dtype shall be {}, but now is {}".format(
-                        get_input_type(program, plan_id, j),
+                        get_input_type(program, program_plan_id, j),
                         type(cur_plan_test_inputs[j]),
                     )
 
@@ -221,13 +227,16 @@ def assert_valid_bundle(
                 # pyre-fixme[16]: Undefined attribute [16]: Item `bool` of `typing.Union[bool, float, int, torch._tensor.Tensor]`
                 # has no attribute `dtype`.
                 assert cur_plan_test_expected_outputs[j].dtype == get_output_dtype(
-                    program, plan_id, j
+                    program, program_plan_id, j
                 ), "The label tensor {} dtype shall be {}, but now is {}".format(
                     cur_plan_test_expected_outputs[j],
-                    get_output_dtype(program, plan_id, j),
+                    get_output_dtype(program, program_plan_id, j),
                     cur_plan_test_expected_outputs[j].dtype,
                 )
 
+        program_plan_id += 1
+        bp_plan_id += 1
+
 
 def create_bundled_program(
     program: Program,
@@ -245,10 +254,7 @@ def create_bundled_program(
     execution_plan_tests: List[BundledExecutionPlanTest] = []
 
     # Emit data and metadata of bundled tensor
-    for plan_id in range(len(program.execution_plan)):
-        plan_test: ConfigExecutionPlanTest = bundled_config.execution_plan_tests[
-            plan_id
-        ]
+    for plan_test in bundled_config.execution_plan_tests:
         test_sets: List[BundledIOSet] = []
 
         # emit I/O sets for each execution plan test
@@ -283,7 +289,11 @@ def create_bundled_program(
             )
 
         # emit the whole execution plan test
-        execution_plan_tests.append(BundledExecutionPlanTest(test_sets=test_sets))
+        execution_plan_tests.append(
+            BundledExecutionPlanTest(
+                method_name=plan_test.method_name, test_sets=test_sets
+            )
+        )
 
     program_bytes: bytes = _serialize_pte_binary(program)
 

bundled_program/schema.py

@@ -73,6 +73,11 @@ class BundledIOSet:
 class BundledExecutionPlanTest:
     """Context for testing and verifying an exceution plan."""
 
+    # The name of the method to test; e.g., "forward" for the forward() method
+    # of an nn.Module. This name match a method defined by the ExecuTorch
+    # program.
+    method_name: str
+
     # Sets of input/outputs to test with.
     test_sets: List[BundledIOSet]