|
27 | 27 | from executorch.backends.apple.metal.metal_backend import MetalBackend |
28 | 28 | from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner |
29 | 29 | from executorch.exir import to_edge_transform_and_lower |
| 30 | +from executorch.exir.backend.compile_spec_schema import CompileSpec |
30 | 31 | from torch import nn |
31 | 32 | from torch.export import export |
32 | 33 | from torch.nn.attention import SDPBackend |
@@ -959,21 +960,25 @@ def linear_filter(m, fqn): |
959 | 960 |
|
960 | 961 |
|
961 | 962 | def export_model_to_metal( |
962 | | - model: nn.Module, example_inputs: Tuple[torch.Tensor, ...] |
| 963 | + model: nn.Module, |
| 964 | + example_inputs: Tuple[torch.Tensor, ...], |
| 965 | + codesign_identity: Optional[str] = None, |
963 | 966 | ) -> Any: |
964 | 967 | """Export model through the Metal backend pipeline.""" |
965 | 968 | method_name = "forward" |
966 | 969 |
|
| 970 | + compile_specs = [MetalBackend.generate_method_name_compile_spec(method_name)] |
| 971 | + if codesign_identity: |
| 972 | + compile_specs.append( |
| 973 | + CompileSpec("codesign_identity", codesign_identity.encode("utf-8")) |
| 974 | + ) |
| 975 | + |
967 | 976 | with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): |
968 | 977 | aten_dialect = export(model, example_inputs, strict=False) |
969 | 978 |
|
970 | 979 | edge_program = to_edge_transform_and_lower( |
971 | 980 | aten_dialect, |
972 | | - partitioner=[ |
973 | | - MetalPartitioner( |
974 | | - [MetalBackend.generate_method_name_compile_spec(method_name)] |
975 | | - ) |
976 | | - ], |
| 981 | + partitioner=[MetalPartitioner(compile_specs)], |
977 | 982 | ) |
978 | 983 |
|
979 | 984 | executorch_program = edge_program.to_executorch() |
@@ -1326,6 +1331,23 @@ def run_test_in_directory(test_dir: Path) -> None: |
1326 | 1331 | with tempfile.TemporaryDirectory() as tmpdir: |
1327 | 1332 | run_test_in_directory(Path(tmpdir)) |
1328 | 1333 |
|
| 1334 | + @unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON) |
| 1335 | + @unittest.skipIf(not IS_MACOS, "codesign requires macOS") |
| 1336 | + def test_codesign_export(self): |
| 1337 | + """Test that export with codesign_identity='-' signs the .so and succeeds.""" |
| 1338 | + model, example_inputs = get_model_and_inputs("add", dtype=torch.float32) |
| 1339 | + # codesign -f -s - runs during export; check=True means CalledProcessError |
| 1340 | + # is raised (and the test fails) if signing fails. |
| 1341 | + program = export_model_to_metal(model, example_inputs, codesign_identity="-") |
| 1342 | + self.assertGreater(len(program.buffer), 0) |
| 1343 | + |
| 1344 | + @unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON) |
| 1345 | + def test_export_without_codesign(self): |
| 1346 | + """Test that export without codesign_identity skips signing.""" |
| 1347 | + model, example_inputs = get_model_and_inputs("add", dtype=torch.float32) |
| 1348 | + program = export_model_to_metal(model, example_inputs) |
| 1349 | + self.assertGreater(len(program.buffer), 0) |
| 1350 | + |
1329 | 1351 |
|
1330 | 1352 | # ============================================================================= |
1331 | 1353 | # Dynamically generate test methods for each module and dtype in MODULE_REGISTRY |
|
0 commit comments