Function lower_all_submodules_to_backend

exir/backend/backend_api.py:561–640

Lower all submodules nodes given in the method_to_submodule_nodes map to backend_id.

(
    backend_id: str,
    method_to_submodules_nodes: Dict[str, List[torch.fx.Node]],
    method_to_tagged_edge_program: Dict[str, ExportedProgram],
)

def lower_all_submodules_to_backend(
    backend_id: str,
    method_to_submodules_nodes: Dict[str, List[torch.fx.Node]],
    method_to_tagged_edge_program: Dict[str, ExportedProgram],
) -> None:
    """
    Lower all submodules nodes given in the method_to_submodule_nodes map to backend_id.
    """
    # The created exported program for the submodules are in the call_module node's meta data
    # We just map the method_to_submodule_nodes directly to the method_to_partitioned_exported_programs
    method_to_partitioned_program = {
        method_name: [
            # perform deep copy here in case backends change graph inside preprocess method
            copy.deepcopy(node.meta["submodule_program"])
            for node in call_submodule_nodes
        ]
        for method_name, call_submodule_nodes in method_to_submodules_nodes.items()
    }
    method_to_compile_specs = {
        method_name: [node.meta["compile_spec"] for node in call_submodule_nodes]
        for method_name, call_submodule_nodes in method_to_submodules_nodes.items()
    }

    backend_name_to_subclass = {
        subclass.__name__: subclass for subclass in BackendDetails.__subclasses__()
    }
    if backend_id not in backend_name_to_subclass:
        raise NotImplementedError(f"Backend {backend_id} was not found.")

    method_to_preprocess_result: dict[str, List[PreprocessResult]] = (
        backend_name_to_subclass[backend_id].preprocess_multimethod(
            method_to_partitioned_program, method_to_compile_specs
        )
    )

    for method_name in method_to_preprocess_result.keys():
        owning_program = method_to_tagged_edge_program[method_name]
        list_of_preprocess_results = method_to_preprocess_result[method_name]
        list_of_call_submodule_nodes = method_to_submodules_nodes[method_name]
        list_of_compile_specs = method_to_compile_specs[method_name]
        for preprocess_result, call_submodule_node, compile_spec in zip(
            list_of_preprocess_results,
            list_of_call_submodule_nodes,
            list_of_compile_specs,
        ):
            submodule_program = call_submodule_node.meta["submodule_program"]
            lowered_module = LoweredBackendModule(
                edge_program=submodule_program,
                backend_id=backend_id,
                processed_bytes=preprocess_result.processed_bytes,
                compile_specs=compile_spec,
                named_data_store_output=preprocess_result.data_store_output,
            )
            lowered_module.meta = {
                "debug_handle_map": preprocess_result.debug_handle_map,
            }
            if preprocess_result._delegate_info_meta is not None:
                assert lowered_module.meta is not None