executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/backends/apple/coreml/compiler/coreml_preprocess.py

@@ -0,0 +1,993 @@
+#  Copyright © 2023 Apple Inc. All rights reserved.
+
+# CoreML backend for delegating a EdgeProgram to CoreML.
+
+import json
+import logging
+
+import shutil
+import tempfile
+import uuid
+from dataclasses import asdict, dataclass
+from enum import Enum
+
+from pathlib import Path
+
+from typing import Any, Dict, final, List, Optional, Tuple, Union
+
+import coremltools as ct
+import coremltools.optimize as cto
+from executorch.backends.apple.coreml import executorchcoreml
+from executorch.backends.apple.coreml.compiler.enumerated_shape_utils import (
+    _get_ct_inputs,
+    _SymbolicShapeToEnumeratedShapeMap,
+)
+from executorch.backends.apple.coreml.logging import get_coreml_log_level
+from executorch.exir.backend.backend_details import (
+    BackendDetails,
+    ExportedProgram,
+    PreprocessResult,
+)
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+from executorch.backends.apple.coreml.compiler.torch_ops import *  # noqa: F401, F403
+
+logger = logging.getLogger(__name__)
+logger.setLevel(get_coreml_log_level(default_level=logging.WARNING))
+
+
+def _gathers_sharing_a_weight(mlmodel: ct.models.MLModel) -> List[str]:
+    """
+    Names of gather ops whose table is also consumed by an op of another type.
+
+    Such a constant reaches the compressor under two different configurations at once,
+    which it rejects rather than resolves. Naming these gathers lets them be configured
+    the same way as the op they share the constant with; every other gather keeps the
+    opt-out. Returns nothing for a model whose program is not available, which leaves
+    the behaviour as it was.
+    """
+    program = getattr(mlmodel, "_mil_program", None)
+    if program is None:
+        return []
+
+    consumers: Dict[str, List[Any]] = {}
+    for function in program.functions.values():
+        for op in function.operations:
+            for value in op.inputs.values():
+                for var in value if isinstance(value, (list, tuple)) else [value]:
+                    source = getattr(var, "op", None)
+                    if source is not None and source.op_type == "const":
+                        consumers.setdefault(var.name, []).append(op)
+
+    shared: List[str] = []
+    for ops in consumers.values():
+        gathers = [op for op in ops if op.op_type == "gather"]
+        if gathers and len(gathers) < len(ops):
+            shared += [op.name for op in gathers]
+    return shared
+
+
+class COMPILE_SPEC_KEYS(Enum):
+    COMPUTE_UNITS = "compute_units"
+    MODEL_TYPE = "model_type"
+    MIN_DEPLOYMENT_TARGET = "min_deployment_target"
+    MODEL_COMPUTE_PRECISION = "model_compute_precision"
+    OP_LINEAR_QUANTIZER_CONFIG = "op_linear_quantizer_config"
+    ENUMERATED_SHAPES = "enumerated_shapes"
+    PASS_PIPELINE = "pass_pipeline"
+    MULTIMETHOD_WEIGHT_SHARING_STRATEGY = "multimethod_weight_sharing_strategy"
+
+
+class MULTIMETHOD_WEIGHT_SHARING_STRATEGY(Enum):
+    # Methods are processed independently with no weight sharing.
+    DISABLED = "disabled"
+    # Partitions must align positionally across methods; enables weight sharing
+    # via NamedDataStore. Raises an error if partition counts don't match.
+    POSITIONAL = "positional"
+
+
+class MODEL_PATHS(Enum):
+    MODEL = "model.mlpackage"
+    COMPILED_MODEL = "model.mlmodelc"
+    METADATA = "metadata.json"
+    DEBUG_INFO = "debug_info.json"
+
+
+@dataclass
+class MethodMetadata:
+    # The method input names.
+    inputNames: List[str]
+    # The method output names.
+    outputNames: List[str]
+
+
+@dataclass
+class ModelMetadata:
+    # The model input names (for single-method models).
+    inputNames: List[str]
+    # The model output names (for single-method models).
+    outputNames: List[str]
+    # The model identifier.
+    identifier: str
+
+
+@dataclass
+class MultifunctionModelMetadata:
+    # The model identifier.
+    identifier: str
+    # Per-method metadata (method name -> MethodMetadata).
+    methods: Dict[str, MethodMetadata]
+
+
+@dataclass
+class ModelDebugInfo:
+    # Version info.
+    versionInfo: Dict[str, str]
+    # Mapping from debug symbol to operation path.
+    debugSymbolToOperationPath: Dict[str, List[Dict[str, str]]]
+    # Mapping from debug symbol to handle.
+    debugSymbolToHandles: Dict[str, List[int]]
+
+
+@final
+class CoreMLBackend(BackendDetails):
+    class MODEL_TYPE(Enum):
+        MODEL = "model"
+        COMPILED_MODEL = "compiled_model"
+
+    @staticmethod
+    def generate_model_type_compile_spec(model_type: MODEL_TYPE) -> CompileSpec:
+        """
+        Returns the compile spec representing the given model type.
+
+        If the model type is ``MODEL_TYPE.Model`` then the ``CoreMLBackend`` returns
+        the in-memory representation of the ``mlpackage`` contents.
+
+        If the model type is ``MODEL_TYPE.COMPILED_MODEL`` then the ``CoreMLBackend`` compiles the model
+        and returns the in-memory representation of ``mlmodelc`` (compiled model) contents.
+        """
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.MODEL_TYPE.value, model_type.value.encode("utf-8")
+        )
+
+    @staticmethod
+    def model_type_from_compile_specs(compile_specs: List[CompileSpec]) -> MODEL_TYPE:
+        """
+        Returns the model type by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.MODEL_TYPE.value:
+                return CoreMLBackend.MODEL_TYPE(compile_spec.value.decode("utf-8"))
+
+        return CoreMLBackend.MODEL_TYPE.MODEL
+
+    @staticmethod
+    def generate_compute_precision_compile_spec(
+        compute_precision: ct.precision,
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the model compute precision, for additional details
+        please refer to the documentation for ``coremltools.precision``.
+        """
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.MODEL_COMPUTE_PRECISION.value,
+            compute_precision.value.encode("utf-8"),
+        )
+
+    @staticmethod
+    def model_compute_precision_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> ct.precision:
+        """
+        Returns the model's compute precision by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.MODEL_COMPUTE_PRECISION.value:
+                return ct.precision(compile_spec.value.decode("utf-8"))
+
+        return ct.precision.FLOAT16
+
+    @staticmethod
+    def generate_minimum_deployment_target_compile_spec(
+        min_deployment_target: Optional[ct.target],
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the minimum deployment target on which the model can run,
+        for additional details please refer to the documentation for ``coremltools.target``.
+        """
+        value = str("").encode("utf-8")
+        if min_deployment_target is not None:
+            value = str(min_deployment_target.value).encode("utf-8")
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_TARGET.value,
+            value,
+        )
+
+    @staticmethod
+    def min_deployment_target_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> Optional[ct.target]:
+        """
+        Returns the minimum deployment target by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_TARGET.value:
+                value = compile_spec.value.decode("utf-8")
+                if value == "":
+                    return None
+                compile_spec_value: int = int(value)
+                return ct.target(compile_spec_value)
+
+        return None
+
+    @staticmethod
+    def compute_unit_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> ct.ComputeUnit:
+        """
+        Returns the minimum deployment target by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.COMPUTE_UNITS.value:
+                return ct.ComputeUnit[compile_spec.value.decode("utf-8").upper()]
+
+        return ct.ComputeUnit.ALL
+
+    @staticmethod
+    def generate_compute_unit_compile_spec(
+        compute_unit: ct.ComputeUnit,
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the compute units on which the model can run, for additional details
+        please refer to the documentation for ``coremltools.ComputeUnit`.
+        """
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.COMPUTE_UNITS.value,
+            compute_unit.name.lower().encode("utf-8"),
+        )
+
+    @staticmethod
+    def generate_op_linear_quantizer_config_compile_spec(
+        op_linear_quantizer_config: Dict,
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the model post conversion quantization,
+        which is a dict that will construct cto.coreml.OpLinearQuantizerConfig
+        """
+        str_representation = json.dumps(op_linear_quantizer_config)
+        byte_representation = str_representation.encode("utf-8")
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.OP_LINEAR_QUANTIZER_CONFIG.value,
+            byte_representation,
+        )
+
+    @staticmethod
+    def op_linear_quantizer_config_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> cto.coreml.OpLinearQuantizerConfig:
+        """
+        Returns the model's post conversion quantization by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.OP_LINEAR_QUANTIZER_CONFIG.value:
+                config_dict_str = compile_spec.value.decode("utf-8")
+                config_dict = json.loads(config_dict_str)
+                config = cto.coreml.OpLinearQuantizerConfig._from_dict(config_dict)
+                return config
+
+        return None
+
+    @staticmethod
+    def generate_pass_pipeline_compile_spec(pass_names: List[str]) -> CompileSpec:
+        """
+        Creates a compile spec representing the pass pipeline to be used by the CoreML backend
+        :param pass_names: the list of pass names
+        """
+        str_representation = json.dumps(pass_names)
+        byte_representation = str_representation.encode("utf-8")
+        return CompileSpec(COMPILE_SPEC_KEYS.PASS_PIPELINE.value, byte_representation)
+
+    @staticmethod
+    def pass_pipeline_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> ct.PassPipeline:
+        """
+        Creates a PassPipeline from the list of compile specs, or returns the default if none are provided.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.PASS_PIPELINE.value:
+                pass_names_str = compile_spec.value.decode("utf-8")
+                pass_names = json.loads(pass_names_str)
+                return ct.PassPipeline(
+                    pass_names, pipeline_name="executorch_user_pipeline"
+                )
+
+        return ct.PassPipeline.DEFAULT
+
+    @staticmethod
+    def generate_multimethod_weight_sharing_strategy_compile_spec(
+        strategy: "MULTIMETHOD_WEIGHT_SHARING_STRATEGY",
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the multimethod weight sharing strategy.
+
+        Args:
+            strategy: The weight sharing strategy to use when combining methods.
+                POSITIONAL: Partitions must align positionally across methods; enables
+                    weight sharing via NamedDataStore. Raises error if partitions don't align.
+                DISABLED: Methods are processed independently with no weight sharing.
+        """
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.MULTIMETHOD_WEIGHT_SHARING_STRATEGY.value,
+            strategy.value.encode("utf-8"),
+        )
+
+    @staticmethod
+    def multimethod_weight_sharing_strategy_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> "MULTIMETHOD_WEIGHT_SHARING_STRATEGY":
+        """
+        Returns the multimethod weight sharing strategy by parsing the list of compile specs.
+        Defaults to POSITIONAL if not specified.
+        """
+        for compile_spec in compile_specs:
+            if (
+                compile_spec.key
+                == COMPILE_SPEC_KEYS.MULTIMETHOD_WEIGHT_SHARING_STRATEGY.value
+            ):
+                return MULTIMETHOD_WEIGHT_SHARING_STRATEGY(
+                    compile_spec.value.decode("utf-8")
+                )
+
+        return MULTIMETHOD_WEIGHT_SHARING_STRATEGY.POSITIONAL
+
+    @staticmethod
+    def generate_enumerated_shapes_compile_spec(
+        ep: ExportedProgram,
+        enumerated_shapes: Dict[str, List[List[int]]],
+    ) -> CompileSpec:
+        """
+        Returns the compile spec representing the model enumerated shapes
+        enumerated_shapes is a dictionary for each input to its enumerated shapes, e.g.,
+
+        enumerated_shapes = {
+         {"x": [[1, 1, 24], [8, 9, 24]]
+         {"y": [[1, 6], [30, 6]],
+        ]
+
+        means the model can handle x can be shape [1, 1, 24] or [8, 9, 24] and y can be shape [1, 6] or [30, 6].
+
+        Only multiple inputs can have enumerated shapes if using iOS18 or later.
+        In this case, each input must have the same number of enumerated shapes, and these shapes are tied together
+        by their order in the list. For example, the model above can handle x with shape [1, 1, 24] and y with shape [1, 6],
+        or x with shape [8, 9, 24] and y with shape [30, 6], but not x with shape [1, 1, 24] and y with shape [30, 6].
+
+        Passing incorrect shapes at runtime will result in an error.
+        """
+        emap = _SymbolicShapeToEnumeratedShapeMap.from_exported_program(
+            ep,
+            enumerated_shapes,
+        )
+        str_representation = emap.to_json()
+        byte_representation = str_representation.encode("utf-8")
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.ENUMERATED_SHAPES.value,
+            byte_representation,
+        )
+
+    @staticmethod
+    def enumerated_shapes_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> cto.coreml.OpLinearQuantizerConfig:
+        """
+        Returns the model's post conversion quantization by parsing the list of compile specs.
+        """
+        for compile_spec in compile_specs:
+            if compile_spec.key == COMPILE_SPEC_KEYS.ENUMERATED_SHAPES.value:
+                emap_json = compile_spec.value.decode("utf-8")
+                emap = _SymbolicShapeToEnumeratedShapeMap.from_json(emap_json)
+                return emap
+        return None
+
+    @staticmethod
+    def generate_compile_specs(
+        compute_unit: ct.ComputeUnit = ct.ComputeUnit.ALL,
+        minimum_deployment_target: Optional[ct.target] = None,
+        compute_precision: ct.precision = ct.precision.FLOAT16,
+        model_type: MODEL_TYPE = MODEL_TYPE.MODEL,
+        op_linear_quantizer_config: Optional[Dict] = None,
+        pass_names: Optional[List[str]] = None,
+    ) -> List[CompileSpec]:
+        """
+        Returns the list of compile specs that's used by CoreMLBackend to lower the module.
+        """
+        compile_specs: List[CompileSpec] = []
+        compile_specs.append(
+            CoreMLBackend.generate_compute_unit_compile_spec(compute_unit)
+        )
+        compile_specs.append(
+            CoreMLBackend.generate_minimum_deployment_target_compile_spec(
+                minimum_deployment_target
+            )
+        )
+        compile_specs.append(
+            CoreMLBackend.generate_compute_precision_compile_spec(compute_precision)
+        )
+        compile_specs.append(CoreMLBackend.generate_model_type_compile_spec(model_type))
+        if op_linear_quantizer_config is not None:
+            compile_specs.append(
+                CoreMLBackend.generate_op_linear_quantizer_config_compile_spec(
+                    op_linear_quantizer_config
+                )
+            )
+        if pass_names is not None:
+            compile_specs.append(
+                CoreMLBackend.generate_pass_pipeline_compile_spec(pass_names)
+            )
+
+        return compile_specs
+
+    @staticmethod
+    def model_metadata_from_spec(
+        model_spec: ct.proto.Model_pb2, identifier: str  # pyre-ignore
+    ) -> ModelMetadata:
+        input_names: List[str] = [input.name for input in model_spec.description.input]
+        output_names = [output.name for output in model_spec.description.output]
+
+        if len(output_names) == 0:
+            raise ValueError("Cannot lower a model with no outputs in CoreML.")
+        if len(input_names) == 0:
+            assert (
+                model_spec.specificationVersion >= 9
+            ), "Deploying a model with no inputs in CoreML requires you set minimum_deployment_target to iOS18 or later in the CoreMLPartitioner."
+
+        return ModelMetadata(
+            inputNames=input_names, outputNames=output_names, identifier=identifier
+        )
+
+    @staticmethod
+    def get_debug_symbol(operation_path: List[Dict[str, str]]) -> Optional[str]:
+        if len(operation_path) == 0:
+            return None
+
+        operator_name: Optional[str] = operation_path[-1].get("Operator", None)
+        output_name: Optional[str] = operation_path[-1].get("Output", None)
+        if output_name is None or operator_name is None:
+            return None
+
+        return output_name + ":" + operator_name
+
+    @staticmethod
+    def get_model_debug_info(model_package_dir: Path) -> Optional[ModelDebugInfo]:
+        delegate_info_file = model_package_dir / "executorch_debug_handle_mapping.json"
+
+        if not delegate_info_file.is_file():
+            return None
+
+        delegate_info: Optional[Dict[str, Any]] = None
+
+        try:
+            with open(delegate_info_file) as f:
+                delegate_info = json.load(f)
+        except ValueError:
+            return None
+
+        if delegate_info is None:
+            return None
+
+        debug_handle_to_operation_path_mapping: Optional[Dict[str, Any]] = (
+            delegate_info.get("mapping", None)
+        )
+
+        if debug_handle_to_operation_path_mapping is None:
+            return None
+
+        debug_symbol_to_operation_path: Dict[str, List[Dict[str, str]]] = {}
+        debug_symbol_to_handles: Dict[str, List[int]] = {}
+        for (
+            debug_handle,
+            operation_paths,
+        ) in debug_handle_to_operation_path_mapping.items():
+            debug_handle_value: Optional[int] = None
+            try:
+                debug_handle_value = int(debug_handle)
+            except ValueError:
+                debug_handle_value = None
+
+            if debug_handle_value is None:
+                continue
+
+            for operation_path in operation_paths:
+                debug_symbol: Optional[str] = CoreMLBackend.get_debug_symbol(
+                    operation_path=operation_path
+                )
+
+                if debug_symbol is None:
+                    continue
+
+                debug_handle_values: List[int] = debug_symbol_to_handles.get(
+                    debug_symbol, []
+                )
+                debug_handle_values.append(debug_handle_value)
+                debug_symbol_to_handles[debug_symbol] = debug_handle_values
+
+                debug_symbol_to_operation_path[debug_symbol] = operation_path
+
+        version_info: Dict[str, str] = delegate_info.get("version", {})
+
+        return ModelDebugInfo(
+            versionInfo=version_info,
+            debugSymbolToOperationPath=debug_symbol_to_operation_path,
+            debugSymbolToHandles=debug_symbol_to_handles,
+        )
+
+    @staticmethod
+    def save_model_metadata(
+        model_metadata: Union[ModelMetadata, MultifunctionModelMetadata],
+        model_dir_path: Path,
+    ):
+        # Store model metadata.
+        model_metadata_path = Path(model_dir_path) / MODEL_PATHS.METADATA.value
+        model_metadata_json = json.dumps(asdict(model_metadata))
+        with open(model_metadata_path, "w") as outfile:
+            outfile.write(model_metadata_json)
+
+    @staticmethod
+    def save_model_debug_info(model_debug_info: ModelDebugInfo, model_dir_path: Path):
+        # Store model debug info.
+        model_debug_info_path = Path(model_dir_path) / MODEL_PATHS.DEBUG_INFO.value
+        model_debug_info_json = json.dumps(asdict(model_debug_info))
+        with open(model_debug_info_path, "w") as outfile:
+            outfile.write(model_debug_info_json)
+
+    @staticmethod
+    def _convert_to_mlmodel(
+        edge_program: ExportedProgram,
+        compile_specs: List[CompileSpec],
+        skip_model_load: bool = True,
+    ) -> ct.models.MLModel:
+        """
+        Convert an ExportedProgram to a CoreML MLModel.
+
+        Args:
+            edge_program: The edge program to convert
+            compile_specs: Compile specs for this conversion
+            skip_model_load: Whether to skip loading the model (for efficiency)
+
+        Returns:
+            The converted MLModel
+        """
+        model_compute_precision = (
+            CoreMLBackend.model_compute_precision_from_compile_specs(compile_specs)
+        )
+        minimum_deployment_target = (
+            CoreMLBackend.min_deployment_target_from_compile_specs(compile_specs)
+        )
+        compute_units = CoreMLBackend.compute_unit_from_compile_specs(compile_specs)
+        pass_pipeline = CoreMLBackend.pass_pipeline_from_compile_specs(compile_specs)
+        enumerated_shapes = CoreMLBackend.enumerated_shapes_from_compile_specs(
+            compile_specs
+        )
+
+        # If using enumerated shapes, pass inputs explicitly to CoreML's convert()
+        ct_inputs = None
+        if enumerated_shapes is not None:
+            ct_inputs = _get_ct_inputs(edge_program, enumerated_shapes)
+
+            # Check there are not multiple enumerated inputs if iOS is below 18
+            if (minimum_deployment_target is None) or (
+                minimum_deployment_target < ct.target.iOS18
+            ):
+                n_enumerated_inputs = sum(
+                    1
+                    for ct_in in ct_inputs
+                    if isinstance(ct_in.shape, ct.EnumeratedShapes)
+                )
+                if n_enumerated_inputs > 1:
+                    raise ValueError(
+                        f"Your program has {n_enumerated_inputs} enumerated inputs, "
+                        f"but minimum_deployment_target is {minimum_deployment_target}. "
+                        "Multiple enumerated inputs requires iOS18 or later."
+                    )
+
+        mlmodel = ct.convert(
+            model=edge_program,
+            source="pytorch",
+            convert_to="mlprogram",
+            pass_pipeline=pass_pipeline,
+            skip_model_load=skip_model_load,
+            compute_precision=model_compute_precision,
+            minimum_deployment_target=minimum_deployment_target,
+            compute_units=compute_units,
+            inputs=ct_inputs,
+        )
+
+        # Apply quantization if specified
+        op_linear_quantizer_config = (
+            CoreMLBackend.op_linear_quantizer_config_from_compile_specs(compile_specs)
+        )
+        if op_linear_quantizer_config is not None:
+            logger.warning(
+                "Core ML Backend op_linear_quantizer_config API is experimental"
+            )
+            # Gathers are left alone so that embedding tables are not quantized by this
+            # config, except where the table is also some other op's weight. A tied
+            # embedding is one constant feeding both a gather and a linear, and coremltools
+            # refuses to compress a constant its consumers disagree about, so opting the
+            # gather out there does not skip the table, it fails the whole lowering.
+            tied = _gathers_sharing_a_weight(mlmodel)
+            config = cto.coreml.OptimizationConfig(
+                global_config=op_linear_quantizer_config,
+                op_type_configs={"gather": None},
+                op_name_configs={name: op_linear_quantizer_config for name in tied},
+            )
+            mlmodel = cto.coreml.linear_quantize_weights(mlmodel, config=config)
+
+        return mlmodel
+
+    @staticmethod
+    def preprocess_model(
+        mlmodel: ct.models.MLModel, model_type: MODEL_TYPE
+    ) -> PreprocessResult:
+        identifier = "executorch_" + str(uuid.uuid4())
+        dir_path: Path = Path(tempfile.gettempdir()) / identifier
+        model_dir_path: Path = dir_path / "lowered_module"
+        model_spec: ct.proto.Model_pb2 = mlmodel.get_spec()
+        logger.warning(
+            f"The model with identifier {identifier} was exported with CoreML specification version {model_spec.specificationVersion}, and it will not run on all version of iOS/macOS."
+            " See https://apple.github.io/coremltools/mlmodel/Format/Model.html#model for information on what OS versions are compatible with this specifcation version."
+            " If you want to control the deployment target, please set the minimum_deployment_target compile spec in the CoreMLPartitioner."
+        )
+
+        model_metadata: ModelMetadata = CoreMLBackend.model_metadata_from_spec(
+            model_spec=model_spec,
+            identifier=identifier,
+        )
+
+        # Save model.
+        model_path = model_dir_path / MODEL_PATHS.MODEL.value
+        mlmodel.save(str(model_path))
+        # Extract delegate mapping file.
+        model_debug_info: Optional[ModelDebugInfo] = CoreMLBackend.get_model_debug_info(
+            model_path
+        )
+
+        match model_type:
+            case CoreMLBackend.MODEL_TYPE.COMPILED_MODEL:
+                shutil.rmtree(str(model_path.resolve()))
+                model_path = model_dir_path / MODEL_PATHS.COMPILED_MODEL.value
+                compiled_model_path = mlmodel.get_compiled_model_path()
+                shutil.move(
+                    compiled_model_path,
+                    str(model_path.resolve()),
+                )
+
+            case _:
+                pass
+
+        CoreMLBackend.save_model_metadata(
+            model_metadata=model_metadata, model_dir_path=model_dir_path
+        )
+        if model_debug_info is not None:
+            CoreMLBackend.save_model_debug_info(
+                model_debug_info=model_debug_info, model_dir_path=model_dir_path
+            )
+
+        processed_bytes: bytes = (
+            executorchcoreml.flatten_directory_contents(str(model_dir_path.resolve()))
+            or b""
+        )
+
+        debug_handle_map: Optional[Dict[str, Tuple[int]]] = None
+        if model_debug_info is not None:
+            debug_handle_map = {
+                key: tuple(value)
+                for key, value in model_debug_info.debugSymbolToHandles.items()
+            }
+
+        shutil.rmtree(str(dir_path.resolve()))
+        return PreprocessResult(
+            processed_bytes=processed_bytes,
+            debug_handle_map=debug_handle_map,
+        )
+
+    @staticmethod
+    def preprocess(
+        edge_program: ExportedProgram,
+        compile_specs: List[CompileSpec],
+    ) -> PreprocessResult:
+        logger.info(f"Edge program: {edge_program}")
+        model_type: CoreMLBackend.MODEL_TYPE = (
+            CoreMLBackend.model_type_from_compile_specs(compile_specs)
+        )
+
+        # Load the model if MODEL_TYPE is 'COMPILED_MODEL'. This step is necessary because
+        # get_compiled_model_path() requires a loaded model.
+        skip_model_load = model_type != CoreMLBackend.MODEL_TYPE.COMPILED_MODEL
+
+        mlmodel = CoreMLBackend._convert_to_mlmodel(
+            edge_program, compile_specs, skip_model_load=skip_model_load
+        )
+
+        return CoreMLBackend.preprocess_model(mlmodel, model_type=model_type)
+
+    @classmethod
+    def preprocess_multimethod(  # noqa: C901
+        cls,
+        edge_programs: Dict[str, List[ExportedProgram]],
+        compile_specs: Dict[str, List[List[CompileSpec]]],
+    ) -> Dict[str, List[PreprocessResult]]:
+        """
+        Preprocess multiple methods, optionally combining them into CoreML multifunction models.
+
+        The behavior is controlled by the MULTIMETHOD_WEIGHT_SHARING_STRATEGY compile spec:
+
+        POSITIONAL (default):
+            Converts each method's ExportedPrograms to mlpackages, then combines
+            corresponding partitions across methods using CoreML's multifunction API
+            (ct.utils.save_multifunction). This enables weight sharing on disk between
+            methods (e.g., decode and prefill for LLMs).
+
+            For each partition index, we create one multifunction model that combines
+            that partition from all methods. This requires all methods to have the same
+            number of partitions. Raises ValueError if partition counts don't match.
+
+            To avoid duplication, we store the combined model ONCE in NamedDataStore
+            with a unique key. Each method's processed_bytes contains a JSON reference
+            to the model in NamedDataStore.
+
+        DISABLED:
+            Each method is processed independently with no weight sharing. Falls back
+            to the default BackendDetails.preprocess_multimethod() implementation.
+
+        Args:
+            edge_programs: Dictionary mapping method name to list of partitioned ExportedPrograms
+            compile_specs: Dictionary mapping method name to list of CompileSpecs for each partition.
+                The MULTIMETHOD_WEIGHT_SHARING_STRATEGY is read from the first method's first
+                partition compile specs.
+
+        Returns:
+            Dictionary mapping method name to list of PreprocessResults. When using POSITIONAL
+            strategy, each method's processed_bytes contains a JSON reference to the shared
+            model in NamedDataStore.
+        """
+        from executorch.exir._serialize._named_data_store import NamedDataStore
+
+        method_names = list(edge_programs.keys())
+
+        if len(method_names) <= 1:
+            # Fall back to default implementation for single method
+            return super().preprocess_multimethod(edge_programs, compile_specs)
+
+        # Get compile specs from the first method's first partition
+        first_method = method_names[0]
+        first_compile_specs = compile_specs[first_method][0]
+
+        # Check the weight sharing strategy
+        weight_sharing_strategy = (
+            cls.multimethod_weight_sharing_strategy_from_compile_specs(
+                first_compile_specs
+            )
+        )
+
+        if weight_sharing_strategy == MULTIMETHOD_WEIGHT_SHARING_STRATEGY.DISABLED:
+            # Process each method independently with no weight sharing
+            logger.info(
+                "Multimethod weight sharing is DISABLED. Processing methods independently."
+            )
+            return super().preprocess_multimethod(edge_programs, compile_specs)
+
+        assert weight_sharing_strategy == MULTIMETHOD_WEIGHT_SHARING_STRATEGY.POSITIONAL
+
+        # POSITIONAL strategy: verify all methods have the same number of partitions
+        num_partitions = len(edge_programs[first_method])
+        for method_name, programs in edge_programs.items():
+            if len(programs) != num_partitions:
+                raise ValueError(
+                    f"Method '{method_name}' has {len(programs)} partitions, but "
+                    f"'{first_method}' has {num_partitions}. POSITIONAL weight sharing "
+                    "strategy requires all methods to have the same number of partitions. "
+                    "Use MULTIMETHOD_WEIGHT_SHARING_STRATEGY.DISABLED if methods should "
+                    "be processed independently."
+                )
+
+        model_type: CoreMLBackend.MODEL_TYPE = cls.model_type_from_compile_specs(
+            first_compile_specs
+        )
+
+        # Create a temporary directory for all the mlpackages
+        temp_dir = Path(tempfile.mkdtemp())
+
+        # Structure: method_mlpackage_paths[method_name][partition_idx] = path
+        method_mlpackage_paths: Dict[str, List[Path]] = {
+            method_name: [] for method_name in method_names
+        }
+
+        # Create a NamedDataStore to hold the shared multifunction models
+        named_data_store = NamedDataStore()
+
+        try:
+            # Convert each method's partitions to mlpackages
+            for method_name in method_names:
+                for partition_idx, edge_program in enumerate(
+                    edge_programs[method_name]
+                ):
+                    method_compile_specs = compile_specs[method_name][partition_idx]
+
+                    logger.info(
+                        f"Converting method '{method_name}' partition {partition_idx} to mlpackage..."
+                    )
+
+                    # Convert to CoreML using shared helper
+                    mlmodel = cls._convert_to_mlmodel(
+                        edge_program, method_compile_specs, skip_model_load=True
+                    )
+
+                    # Save the mlpackage
+                    mlpackage_path = (
+                        temp_dir / f"{method_name}_partition_{partition_idx}.mlpackage"
+                    )
+                    mlmodel.save(str(mlpackage_path))
+                    method_mlpackage_paths[method_name].append(mlpackage_path)
+
+                    logger.info(
+                        f"Saved method '{method_name}' partition {partition_idx} to {mlpackage_path}"
+                    )
+
+            # For each partition index, combine that partition from all methods
+            # into a single multifunction model.
+            # Store combined_processed_bytes[partition_idx] = bytes (for first method)
+            combined_processed_bytes: List[bytes] = []
+            debug_handle_maps: List[Optional[Dict[str, Tuple[int]]]] = []
+            model_keys: List[str] = []  # Keys for NamedDataStore lookup
+
+            for partition_idx in range(num_partitions):
+                logger.info(
+                    f"Combining partition {partition_idx} from all methods into multifunction model..."
+                )
+
+                desc = ct.utils.MultiFunctionDescriptor()
+                for method_name in method_names:
+                    mlpackage_path = method_mlpackage_paths[method_name][partition_idx]
+                    desc.add_function(
+                        str(mlpackage_path),
+                        src_function_name="main",
+                        target_function_name=method_name,
+                    )
+
+                # Set the first method as default
+                desc.default_function_name = first_method
+
+                # Save the combined multifunction model for this partition
+                combined_path = (
+                    temp_dir / f"combined_partition_{partition_idx}.mlpackage"
+                )
+                ct.utils.save_multifunction(desc, str(combined_path))
+
+                logger.info(
+                    f"Saved combined multifunction model for partition {partition_idx} to {combined_path}"
+                )
+
+                # Create output directory for this partition's combined model
+                model_dir_path = temp_dir / f"lowered_module_partition_{partition_idx}"
+                model_dir_path.mkdir(exist_ok=True)
+
+                # Handle model type (compiled vs mlpackage)
+                if model_type == CoreMLBackend.MODEL_TYPE.COMPILED_MODEL:
+                    output_model_path = (
+                        model_dir_path / MODEL_PATHS.COMPILED_MODEL.value
+                    )
+                    combined_model_loaded = ct.models.MLModel(str(combined_path))
+                    compiled_path = combined_model_loaded.get_compiled_model_path()
+                    shutil.move(compiled_path, str(output_model_path))
+                else:
+                    output_model_path = model_dir_path / MODEL_PATHS.MODEL.value
+                    shutil.copytree(str(combined_path), str(output_model_path))
+
+                # For multifunction models, we store all method metadata in a single file
+                # with method names as keys. Each method can have different input/output
+                # names (e.g., masks_1023 vs masks_992).
+                identifier = "executorch_" + str(uuid.uuid4())
+
+                # Extract metadata for each method
+                methods_metadata: Dict[str, MethodMetadata] = {}
+                for method_name in method_names:
+                    method_mlpackage_path = method_mlpackage_paths[method_name][
+                        partition_idx
+                    ]
+                    method_model = ct.models.MLModel(
+                        str(method_mlpackage_path), skip_model_load=True
+                    )
+                    method_spec = method_model.get_spec()
+                    input_names = [inp.name for inp in method_spec.description.input]
+                    output_names = [out.name for out in method_spec.description.output]
+                    methods_metadata[method_name] = MethodMetadata(
+                        inputNames=input_names,
+                        outputNames=output_names,
+                    )
+                    logger.info(
+                        f"Extracted metadata for method '{method_name}' partition {partition_idx}: "
+                        f"{len(input_names)} inputs, {len(output_names)} outputs"
+                    )
+
+                # Create consolidated multifunction metadata
+                multifunction_metadata = MultifunctionModelMetadata(
+                    identifier=identifier,
+                    methods={k: asdict(v) for k, v in methods_metadata.items()},
+                )
+
+                # Save consolidated metadata
+                cls.save_model_metadata(multifunction_metadata, model_dir_path)
+
+                # Note: Debug info is not supported for multifunction models.
+                # The combined model's debug mapping doesn't accurately map back to
+                # individual methods, so we skip it rather than provide incorrect info.
+
+                # Flatten the model directory (with model + metadata) to bytes
+                processed_bytes = (
+                    executorchcoreml.flatten_directory_contents(
+                        str(model_dir_path.resolve())
+                    )
+                    or b""
+                )
+                combined_processed_bytes.append(processed_bytes)
+
+                # Store in NamedDataStore and save the key for later reference
+                model_key = f"coreml_{identifier}"
+                model_keys.append(model_key)
+                named_data_store.add_named_data(model_key, processed_bytes)
+
+                logger.info(
+                    f"Created combined processed bytes for partition {partition_idx} ({len(processed_bytes)} bytes)"
+                )
+                logger.info(f"Stored in NamedDataStore with key '{model_key}'")
+
+                # Debug handle map is not supported for multifunction models
+                debug_handle_maps.append(None)
+
+            # Get the NamedDataStoreOutput to share across PreprocessResults
+            named_data_store_output = named_data_store.get_named_data_store_output()
+
+            # Build PreprocessResults for each method and partition.
+            # All methods get a JSON reference to the model in NamedDataStore.
+            # The model is stored ONLY in NamedDataStore to avoid duplication.
+            # Runtime will detect the JSON reference and load from NamedDataMap.
+            preprocess_results: Dict[str, List[PreprocessResult]] = {
+                method_name: [] for method_name in method_names
+            }
+
+            for partition_idx in range(num_partitions):
+                debug_handle_map = debug_handle_maps[partition_idx]
+
+                for method_name in method_names:
+                    # Create JSON reference for runtime to look up model in NamedDataStore.
+                    # functionName specifies which CoreML function to invoke within the
+                    # multifunction model. The runtime gets the ExecuTorch method name
+                    # separately via BackendInitContext::get_method_name().
+                    #
+                    # We prefix the JSON with a magic number so runtime can identify this
+                    # as a JSON reference rather than raw model bytes.
+                    reference = {
+                        "version": 1,
+                        "key": model_keys[partition_idx],
+                        "functionName": method_name,
+                    }
+                    # Magic number "CMJR" (CoreML JSON Reference) followed by JSON bytes
+                    MAGIC_NUMBER = b"CMJR"
+                    reference_bytes = MAGIC_NUMBER + json.dumps(reference).encode(
+                        "utf-8"
+                    )
+
+                    preprocess_results[method_name].append(
+                        PreprocessResult(
+                            processed_bytes=reference_bytes,
+                            debug_handle_map=debug_handle_map,
+                            data_store_output=named_data_store_output,
+                        )
+                    )
+
+            return preprocess_results
+
+        finally:
+            # Clean up temporary directory
+            shutil.rmtree(str(temp_dir))