-
Notifications
You must be signed in to change notification settings - Fork 571
Commit 4ed3e44
authored
[ExecuTorch][to_backend] Enable to_backend API to leverage preprocess_multimethod (#9824)
### 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_method1 parent 50ee3a5 commit 4ed3e44Copy full SHA for 4ed3e44
File tree
5 files changed
+1257
-67
lines changedFilter options
- exir
- backend
- test
5 files changed
+1257
-67
lines changed
0 commit comments