diff --git a/exir/backend/utils.py b/exir/backend/utils.py index eb9aeb19756..57847f8ebb0 100644 --- a/exir/backend/utils.py +++ b/exir/backend/utils.py @@ -11,7 +11,7 @@ import operator from collections import defaultdict, OrderedDict from functools import lru_cache -from typing import Dict, Iterable, List, Optional, Set, Tuple, Union +from typing import Dict, Iterable, List, Optional, Set, Tuple, Union, Any import torch from executorch.exir.backend.backend_details import ExportedProgram @@ -22,6 +22,9 @@ from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as exir_ops +# get Executorch +from executorch.exir import ExecutorchProgram + from executorch.exir.lowered_backend_module import create_submodule_from_nodes from tabulate import tabulate from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param @@ -431,6 +434,23 @@ def is_shape_dynamic(node: torch.fx.Node) -> bool: return has_free_symbols(node.meta["val"].shape) +def get_delegated_payload(exec_program: ExecutorchProgram) -> Dict[str, Any]: + """ + Get the payload of the executorch program which contains all the delegated modules. + """ + delegated_payload_list = {} + for delegate in exec_program.executorch_program.execution_plan[0].delegates: + backend_id = delegate.id + delegate_index = delegate.processed.index + # delegate_data_size[delegate_index] = (delegate.id, len(data)) + delegated_payload_list[delegate_index] = ( + backend_id, + delegate.compile_specs, + exec_program.executorch_program.backend_delegate_data[delegate_index].data, + ) + return delegated_payload_list + + # TODO - style: use templated types class DelegateMappingBuilder: """