diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index f9b4b947506..9fe373f1be5 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -92,6 +92,19 @@ def get_extra_aoti_compile_context_manager(cls): """Return extra context manager to apply during aoti_compile stage. By default returns an empty context manager.""" return contextlib.nullcontext() + @classmethod + def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None: + """Sign the compiled .so before packing into .pte. + + Called after AOTInductor compilation, before the .so bytes are read + and packed into the named data store. Override in platform-specific + backends to apply code signing (e.g., macOS codesign for Hardened + Runtime compatibility). + + Default: no-op. + """ + return + @classmethod @contextlib.contextmanager def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]): @@ -226,6 +239,9 @@ def preprocess( f"Could not find required files in compiled paths, got {paths}" ) + # Sign the .so for platform-specific requirements (e.g., macOS Hardened Runtime) + cls.codesign_so(so_path, compile_specs) + # Read SO file with open(so_path, "rb") as f: so_data = f.read() diff --git a/backends/apple/metal/metal_backend.py b/backends/apple/metal/metal_backend.py index 5ddd5e13d88..8a2ac34d137 100644 --- a/backends/apple/metal/metal_backend.py +++ b/backends/apple/metal/metal_backend.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import subprocess import typing from typing import Any, Dict, final, List @@ -78,3 +79,16 @@ def get_aoti_compile_options( inductor_configs["aot_inductor.custom_ops_to_c_shims"] = torchao_op_c_shim return inductor_configs + + @classmethod + def codesign_so(cls, so_path: str, compile_specs: List[CompileSpec]) -> None: + """Sign the compiled .so for macOS Hardened Runtime compatibility. + + Only signs if a ``codesign_identity`` compile spec is provided. + Pass ``"-"`` for ad-hoc signing or a Developer ID for distribution. + """ + for spec in compile_specs: + if spec.key == "codesign_identity": + identity = spec.value.decode("utf-8") + subprocess.run(["codesign", "-f", "-s", identity, so_path], check=True) + return diff --git a/backends/apple/metal/tests/test_modules.py b/backends/apple/metal/tests/test_modules.py index 2bf14a0d10b..10e5a4b97a9 100644 --- a/backends/apple/metal/tests/test_modules.py +++ b/backends/apple/metal/tests/test_modules.py @@ -27,6 +27,7 @@ from executorch.backends.apple.metal.metal_backend import MetalBackend from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.compile_spec_schema import CompileSpec from torch import nn from torch.export import export from torch.nn.attention import SDPBackend @@ -787,21 +788,25 @@ def linear_filter(m, fqn): def export_model_to_metal( - model: nn.Module, example_inputs: Tuple[torch.Tensor, ...] + model: nn.Module, + example_inputs: Tuple[torch.Tensor, ...], + codesign_identity: Optional[str] = None, ) -> Any: """Export model through the Metal backend pipeline.""" method_name = "forward" + compile_specs = [MetalBackend.generate_method_name_compile_spec(method_name)] + if codesign_identity: + compile_specs.append( + CompileSpec("codesign_identity", codesign_identity.encode("utf-8")) + ) + with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): aten_dialect = export(model, example_inputs, strict=False) edge_program = to_edge_transform_and_lower( aten_dialect, - partitioner=[ - MetalPartitioner( - [MetalBackend.generate_method_name_compile_spec(method_name)] - ) - ], + partitioner=[MetalPartitioner(compile_specs)], ) executorch_program = edge_program.to_executorch() @@ -1154,6 +1159,23 @@ def run_test_in_directory(test_dir: Path) -> None: with tempfile.TemporaryDirectory() as tmpdir: run_test_in_directory(Path(tmpdir)) + @unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON) + @unittest.skipIf(not IS_MACOS, "codesign requires macOS") + def test_codesign_export(self): + """Test that export with codesign_identity='-' signs the .so and succeeds.""" + model, example_inputs = get_model_and_inputs("add", dtype=torch.float32) + # codesign -f -s - runs during export; check=True means CalledProcessError + # is raised (and the test fails) if signing fails. + program = export_model_to_metal(model, example_inputs, codesign_identity="-") + self.assertGreater(len(program.buffer), 0) + + @unittest.skipIf(SKIP_EXPORT_TESTS, SKIP_REASON) + def test_export_without_codesign(self): + """Test that export without codesign_identity skips signing.""" + model, example_inputs = get_model_and_inputs("add", dtype=torch.float32) + program = export_model_to_metal(model, example_inputs) + self.assertGreater(len(program.buffer), 0) + # ============================================================================= # Dynamically generate test methods for each module and dtype in MODULE_REGISTRY diff --git a/examples/models/parakeet/export_parakeet_tdt.py b/examples/models/parakeet/export_parakeet_tdt.py index c35e17eed59..6a18cd58218 100644 --- a/examples/models/parakeet/export_parakeet_tdt.py +++ b/examples/models/parakeet/export_parakeet_tdt.py @@ -498,10 +498,11 @@ def _linear_bias_decomposition(input, weight, bias=None): return out -def _create_metal_partitioners(programs): +def _create_metal_partitioners(programs, codesign_identity=None): """Create Metal partitioners for all programs except preprocessor.""" from executorch.backends.apple.metal.metal_backend import MetalBackend from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner + from executorch.exir.backend.compile_spec_schema import CompileSpec print("\nLowering to ExecuTorch with Metal...") @@ -521,6 +522,10 @@ def _create_metal_partitioners(programs): partitioner[key] = [] else: compile_specs = [MetalBackend.generate_method_name_compile_spec(key)] + if codesign_identity: + compile_specs.append( + CompileSpec("codesign_identity", codesign_identity.encode("utf-8")) + ) partitioner[key] = [MetalPartitioner(compile_specs)] return partitioner, updated_programs @@ -586,12 +591,18 @@ def _create_vulkan_partitioners(programs, vulkan_force_fp16=False): def lower_to_executorch( - programs, metadata=None, backend="portable", vulkan_force_fp16=False + programs, + metadata=None, + backend="portable", + vulkan_force_fp16=False, + codesign_identity=None, ): if backend == "xnnpack": partitioner, programs = _create_xnnpack_partitioners(programs) elif backend == "metal": - partitioner, programs = _create_metal_partitioners(programs) + partitioner, programs = _create_metal_partitioners( + programs, codesign_identity=codesign_identity + ) elif backend == "mlx": partitioner, programs = _create_mlx_partitioners(programs) elif backend in ("cuda", "cuda-windows"): @@ -714,6 +725,13 @@ def main(): ) parser.add_argument("--vulkan_force_fp16", action="store_true") + parser.add_argument( + "--codesign-identity", + default=None, + help="macOS code signing identity for the Metal backend .so. " + "Use '-' for ad-hoc or a Developer ID for notarized apps. " + "If omitted, the .so is not signed.", + ) args = parser.parse_args() @@ -767,6 +785,7 @@ def main(): metadata=metadata, backend=args.backend, vulkan_force_fp16=args.vulkan_force_fp16, + codesign_identity=args.codesign_identity, ) pte_path = os.path.join(args.output_dir, "model.pte") diff --git a/examples/models/voxtral_realtime/export_voxtral_rt.py b/examples/models/voxtral_realtime/export_voxtral_rt.py index 824c4485662..3dfa53af16a 100644 --- a/examples/models/voxtral_realtime/export_voxtral_rt.py +++ b/examples/models/voxtral_realtime/export_voxtral_rt.py @@ -379,7 +379,7 @@ def _linear_bias_decomposition(input, weight, bias=None): return out -def lower_to_executorch(programs, metadata, backend="xnnpack"): +def lower_to_executorch(programs, metadata, backend="xnnpack", codesign_identity=None): """Lower exported programs to ExecuTorch.""" transform_passes = None @@ -397,6 +397,7 @@ def lower_to_executorch(programs, metadata, backend="xnnpack"): elif backend == "metal": from executorch.backends.apple.metal.metal_backend import MetalBackend from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner + from executorch.exir.backend.compile_spec_schema import CompileSpec print("\nLowering to ExecuTorch with Metal...") @@ -411,6 +412,10 @@ def lower_to_executorch(programs, metadata, backend="xnnpack"): partitioner = {} for key in programs: compile_specs = [MetalBackend.generate_method_name_compile_spec(key)] + if codesign_identity: + compile_specs.append( + CompileSpec("codesign_identity", codesign_identity.encode("utf-8")) + ) partitioner[key] = [MetalPartitioner(compile_specs)] elif backend in ("cuda", "cuda-windows"): from executorch.backends.cuda.cuda_backend import CudaBackend @@ -577,6 +582,13 @@ def main(): choices=["fp32", "bf16"], help="Model dtype (default: fp32).", ) + parser.add_argument( + "--codesign-identity", + default=None, + help="macOS code signing identity for the Metal backend .so. " + "Use '-' for ad-hoc or a Developer ID for notarized apps. " + "If omitted, the .so is not signed.", + ) args = parser.parse_args() backend_for_export = "cuda" if args.backend == "cuda-windows" else args.backend @@ -638,7 +650,12 @@ def main(): metadata["delay_tokens"] = args.delay_tokens # Lower - et = lower_to_executorch(programs, metadata, backend=args.backend) + et = lower_to_executorch( + programs, + metadata, + backend=args.backend, + codesign_identity=args.codesign_identity, + ) # Save pte_path = os.path.join(args.output_dir, "model.pte")