Metal backend: add export-time code signing for Hardened Runtime

The Metal AOTI backend extracts a compiled .so from the .pte at runtime
and dlopen's it. macOS Hardened Runtime rejects unsigned dlopen'd code,
making Metal-backend .pte files unusable in notarized apps.

Add a `codesign_so` hook to `AotiBackend` (no-op by default) that runs
after AOTInductor compilation and before the .so is packed into the .pte.
`MetalBackend` overrides it to run `codesign` when a `codesign_identity`
compile spec is provided. Wire `--codesign-identity` through the Voxtral
Realtime and Parakeet export scripts. Add an e2e test in test_modules
that exports a toy model with ad-hoc signing.

Author: mergennachin

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()

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

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

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")

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")