[ExecuTorch][to_backend] Enable to_backend API to leverage preprocess_multimethod

[mcr229]

Summary

We add a new to_backend api which for multi method models. Specifically, we pass in a dictionary mapping MethodName to ExportedProgram, as well as a dictionary mapping MethodName to Partitioner. We then return a dictionary mapping MethodName to the partitioned and lowered exported program.

In addition, we also provide a new preprocess API for backends to implement. This API is preprocess_multimethod. The signature of the new method is as follows:

    def preprocess_multimethod(
        cls,
        edge_programs: Dict[str, List[ExportedProgram]],
        compile_specs: Dict[str, List[List[CompileSpec]]],
    ) -> Dict[str, list[PreprocessResult]]:
        """
        Runs preprocess on all partitioned Edge Programs across multiple methods. This allows
        backends to share information across partitioned graphs. Backend can serialize shared
        data by putting the shared data into the data_store_output of the preprocess results.
        This will record the shared data used by that specific partition.

        Default implementation is running the existing preprocess implementation on all

        Args:
            edge_programs: Dictionary mapping the method name to a list of all the partitioned
                           edge_programs from that method to be lowered.
            compile_specs: Dictionary mapping the method name to a list of compile_specs. The
                           list of compile specs maps directly to the list of edge_programs for the
                           same given method name i.e. edge_program[method_name][i] --> compile_specs[method_name][i]

        Returns:
            Dictionary mapping the method name to a list of PreprocessResults. The list of
            PreprocessResults maps directly to the list of edge_programs for the same given
            method name. i.e. edge_program[method_name][i] --> result[method_name][i]
        """

This new API enableds backends to preprocess all partitions/methods at once. This way, while processing blobs, they can identify shared components between preprocessed blobs. Shared components can be serialized within the NamedDataStore.

The key change in backend infra, is that when partitioning, we now have to identify all the partitioned graphs to be lowered at once, and pass them to preprocess_multimethod at once. Previously, as we found lowerable partitions, we preprocessed and embedded them into the graph.

Testing

python -m unittest exir.backend.test.test_to_backend_multi_method

[mcr229]

Stack from ghstack (oldest at bottom):

• -> [ExecuTorch][to_backend] Enable to_backend API to leverage preprocess_multimethod #9824
• [Executorch][to_backend] Introduce preprocess_multimethod #9823

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/9824

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 3a29677 with merge base 56c8dc2 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.