(
module_class: Type[nn.Module],
*,
backend_id: str,
extract_delegate_segments: bool,
constant_tensor_alignment: Optional[int] = None,
delegate_alignment: Optional[int] = None,
method_name: str = "forward",
external_constants: bool = False,
)
| 130 | |
| 131 | |
| 132 | def export_module_to_program( |
| 133 | module_class: Type[nn.Module], |
| 134 | *, |
| 135 | backend_id: str, |
| 136 | extract_delegate_segments: bool, |
| 137 | constant_tensor_alignment: Optional[int] = None, |
| 138 | delegate_alignment: Optional[int] = None, |
| 139 | method_name: str = "forward", |
| 140 | external_constants: bool = False, |
| 141 | ) -> ExecutorchProgramManager: |
| 142 | eager_module = module_class().eval() |
| 143 | inputs = () |
| 144 | if hasattr(eager_module, "get_random_inputs"): |
| 145 | inputs = eager_module.get_random_inputs() # type: ignore[operator] |
| 146 | |
| 147 | class WrapperModule(torch.nn.Module): |
| 148 | def __init__(self, fn, method_name=method_name): |
| 149 | super().__init__() |
| 150 | self.fn = fn |
| 151 | self.method_name = method_name |
| 152 | |
| 153 | def forward(self, *args, **kwargs): |
| 154 | return getattr(self.fn, self.method_name)(*args, **kwargs) |
| 155 | |
| 156 | if method_name != "forward": |
| 157 | # Only require wrapper module if we're exporting a specific method other than forward. |
| 158 | exported_program = export(WrapperModule(eager_module), args=inputs) |
| 159 | else: |
| 160 | exported_program = export(eager_module, args=inputs) |
| 161 | |
| 162 | edge_config = EdgeCompileConfig(_check_ir_validity=False) |
| 163 | et_config = exir.ExecutorchBackendConfig( |
| 164 | extract_delegate_segments=extract_delegate_segments, |
| 165 | constant_tensor_alignment=constant_tensor_alignment, |
| 166 | delegate_alignment=delegate_alignment, |
| 167 | external_constants=external_constants, |
| 168 | ) |
| 169 | |
| 170 | if backend_id == "XnnpackBackend": |
| 171 | from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( |
| 172 | XnnpackPartitioner, |
| 173 | ) |
| 174 | |
| 175 | if external_constants: |
| 176 | tagged_module = exported_program.module() |
| 177 | delegate_external_constants_pass_unlifted( |
| 178 | module=tagged_module, |
| 179 | gen_tag_fn=lambda x: module_class.__name__, |
| 180 | ) |
| 181 | exported_program = export(tagged_module, args=inputs) |
| 182 | executorch_program = to_edge_transform_and_lower( |
| 183 | exported_program, |
| 184 | compile_config=edge_config, |
| 185 | partitioner=[XnnpackPartitioner()], |
| 186 | ).to_executorch(config=et_config) |
| 187 | else: |
| 188 | edge: exir.EdgeProgramManager = to_edge(exported_program) |
| 189 | lowered_module = to_backend( # type: ignore[call-arg] |
no test coverage detected