structural api and schema update (#410)

Summary:
Pull Request resolved: #410

Update the AOT API for bundled program to make it more user friendly.

We created two new classes for holding testcases in a more structural way:

MethodTestCase: contains inputs and expected outputs for a single test for method.

MethodTestSuite: Collection of all test cases for a program method.

Detailed design can be found here: https://docs.google.com/document/d/170WJ81dPuNQsvsjqH_wQjafbu2kxtEG1sgX-z6OqCEI/edit#heading=h.69fgezruwl6n

Differential Revision: D49406613

fbshipit-source-id: 532153eb06a9aca7b018b24fb7598cc6ed601fec

Author: Gasoonjia

backends/xnnpack/test/test_xnnpack_utils.py

@@ -6,7 +6,7 @@
 
 import unittest
 from random import randint
-from typing import Any, Tuple
+from typing import Any, List, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -26,7 +26,7 @@
 # import the xnnpack backend implementation
 from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend
 
-from executorch.bundled_program.config import BundledConfig
+from executorch.bundled_program.config import MethodTestCase, MethodTestSuite
 from executorch.bundled_program.core import create_bundled_program
 from executorch.bundled_program.serialize import (
     serialize_from_bundled_program_to_flatbuffer,
@@ -101,14 +101,22 @@ def save_bundled_program(representative_inputs, program, ref_output, output_path
     niter = 1
 
     print("generating bundled program inputs / outputs")
-    inputs_list = [list(representative_inputs) for _ in range(niter)]
-    expected_outputs_list = [
-        [[ref_output] for x in inputs_list],
+
+    method_test_cases: List[MethodTestCase] = []
+    for _ in range(niter):
+        method_test_cases.append(
+            MethodTestCase(
+                inputs=list(representative_inputs),
+                expected_outputs=[ref_output],
+            )
+        )
+
+    method_test_suites = [
+        MethodTestSuite(method_name="forward", method_test_cases=method_test_cases)
     ]
-    bundled_config = BundledConfig([inputs_list], expected_outputs_list)
 
     print("creating bundled program...")
-    bundled_program = create_bundled_program(program, bundled_config)
+    bundled_program = create_bundled_program(program, method_test_suites)
 
     print("serializing bundled program...")
     bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(

bundled_program/config.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import Any, get_args, List, Union
+from typing import get_args, List, Optional, Sequence, Union
 
 import torch
 from executorch.extension.pytree import tree_flatten
@@ -16,7 +16,7 @@
 
 """
 The data types currently supported for element to be bundled. It should be
-consistent with the types in bundled_program.schema.BundledValue.
+consistent with the types in bundled_program.schema.Value.
 """
 ConfigValue: TypeAlias = Union[
     torch.Tensor,
@@ -26,7 +26,7 @@
 ]
 
 """
-All supported types for input/expected output of test set.
+All supported types for input/expected output of MethodTestCase.
 
 Namedtuple is also supported and listed implicity since it is a subclass of tuple.
 """
@@ -35,74 +35,21 @@
 DataContainer: TypeAlias = Union[list, tuple, dict]
 
 
-@dataclass
-class ConfigIOSet:
-    """Type of data BundledConfig stored for each validation set."""
-
-    inputs: List[ConfigValue]
-    expected_outputs: List[ConfigValue]
-
-
-@dataclass
-class ConfigExecutionPlanTest:
-    """All info related to verify execution plan"""
-
-    method_name: str
-    test_sets: List[ConfigIOSet]
-
-
-class BundledConfig:
-    """All information needed to verify a model.
-
-    Public Attributes:
-        execution_plan_tests: inputs, expected outputs, and other info for each execution plan verification.
-        attachments: Other info need to be attached.
-    """
+class MethodTestCase:
+    """Test case with inputs and expected outputs
+    The expected_outputs could be None if user only want to user the test case for profiling."""
 
     def __init__(
-        self,
-        method_names: List[str],
-        # pyre-ignore
-        inputs: List[List[Any]],
-        # pyre-ignore
-        expected_outputs: List[List[Any]],
+        self, inputs: DataContainer, expected_outputs: Optional[DataContainer] = None
     ) -> None:
-        """Contruct the config given inputs and expected outputs
-
-        Args:
-            method_names: All method names need to be verified in program.
-            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.
-
-            expected_outputs: Expected outputs for inputs sharing same index. The size of
-                    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)
-
-        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(
-            method_names, inputs, expected_outputs
-        )
-
-    @staticmethod
-    # TODO(T138930448): Give pyre-ignore commands appropriate warning type and comments.
-    # pyre-ignore
-    def _tree_flatten(unflatten_data: Any) -> List[ConfigValue]:
+        self.inputs: List[ConfigValue] = self._flatten_and_sanity_check(inputs)
+        self.expected_outputs: List[ConfigValue] = []
+        if expected_outputs:
+            self.expected_outputs = self._flatten_and_sanity_check(expected_outputs)
+
+    def _flatten_and_sanity_check(
+        self, unflatten_data: DataContainer
+    ) -> List[ConfigValue]:
         """Flat the given data and check its legality
 
         Args:
@@ -111,6 +58,11 @@ def _tree_flatten(unflatten_data: Any) -> List[ConfigValue]:
         Returns:
             flatten_data: Flatten data with legal type.
         """
+
+        assert isinstance(
+            unflatten_data, get_args(DataContainer)
+        ), f"The input or expected output of MethodTestCase should be in list, tuple or dict, but got {type(unflatten_data)}."
+
         # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`.
         flatten_data, _ = tree_flatten(unflatten_data)
 
@@ -128,70 +80,10 @@ def _tree_flatten(unflatten_data: Any) -> List[ConfigValue]:
 
         return flatten_data
 
-    @staticmethod
-    # pyre-ignore
-    def _check_io_type(test_data_program: List[List[Any]]) -> None:
-        """Check the type of each set of inputs or exepcted_outputs
-
-        Each test set of inputs or expected_outputs will be put into the config
-        should be one of the sub-type in DataContainer.
-
-        Args:
-            test_data_program: inputs or expected_outputs to be put into the config
-                               to verify the whole program.
 
-        """
-        for test_data_execution_plan in test_data_program:
-            for test_set in test_data_execution_plan:
-                assert isinstance(test_set, get_args(DataContainer))
-
-    @staticmethod
-    def _gen_execution_plan_tests(
-        method_names: List[str],
-        # pyre-ignore
-        inputs: List[List[Any]],
-        # pyre-ignore
-        expected_outputs: List[List[Any]],
-    ) -> List[ConfigExecutionPlanTest]:
-        """Generate execution plan test given inputs, expected outputs for verifying each execution plan"""
-
-        execution_plan_tests: List[ConfigExecutionPlanTest] = []
-
-        for (
-            m_name,
-            inputs_per_plan_test,
-            expect_outputs_per_plan_test,
-        ) in zip(method_names, inputs, expected_outputs):
-            test_sets: List[ConfigIOSet] = []
-
-            # transfer I/O sets into ConfigIOSet for each execution plan
-            assert len(inputs_per_plan_test) == len(expect_outputs_per_plan_test), (
-                "The number of input and expected output for identical execution plan should be the same,"
-                + " but got {} and {}".format(
-                    len(inputs_per_plan_test), len(expect_outputs_per_plan_test)
-                )
-            )
-            for unflatten_input, unflatten_expected_output in zip(
-                inputs_per_plan_test, expect_outputs_per_plan_test
-            ):
-                flatten_inputs = BundledConfig._tree_flatten(unflatten_input)
-                flatten_expected_output = BundledConfig._tree_flatten(
-                    unflatten_expected_output
-                )
-                test_sets.append(
-                    ConfigIOSet(
-                        inputs=flatten_inputs, expected_outputs=flatten_expected_output
-                    )
-                )
-
-            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)
+@dataclass
+class MethodTestSuite:
+    """All info related to verify method"""
 
-        return execution_plan_tests
+    method_name: str
+    test_cases: Sequence[MethodTestCase]