executorch 0.2.1__cp39-cp39-macosx_11_0_arm64.whl

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

@@ -0,0 +1,9 @@
+# Copyright © 2023 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from .coreml_preprocess import CoreMLBackend
+
+__all__ = [
+    CoreMLBackend,
+]

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

@@ -0,0 +1,376 @@
+#  Copyright © 2023 Apple Inc. All rights reserved.
+
+# CoreML backend for delegating a EdgeProgram to CoreML.
+
+import json
+
+import shutil
+import uuid
+from dataclasses import asdict, dataclass
+from enum import Enum
+
+from pathlib import Path
+
+from typing import Any, Dict, final, List, Optional, Tuple
+
+import coremltools as ct
+import executorchcoreml
+
+from executorch.exir.backend.backend_details import (
+    BackendDetails,
+    ExportedProgram,
+    PreprocessResult,
+)
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+
+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"
+
+
+class MODEL_PATHS(Enum):
+    MODEL = "model.mlpackage"
+    COMPILED_MODEL = "model.mlmodelc"
+    METADATA = "metadata.json"
+    DEBUG_INFO = "debug_info.json"
+
+
+@dataclass
+class ModelMetadata:
+    # The model input names.
+    inputNames: List[str]
+    # The model output names.
+    outputNames: List[str]
+    # The model identifier.
+    identifier: str
+
+
+@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: 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``.
+        """
+        return CompileSpec(
+            COMPILE_SPEC_KEYS.MIN_DEPLOYMENT_TARGET.value,
+            str(min_deployment_target.value).encode("utf-8"),
+        )
+
+    @staticmethod
+    def min_deployment_target_from_compile_specs(
+        compile_specs: List[CompileSpec],
+    ) -> 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:
+                compile_spec_value: int = int(compile_spec.value.decode("utf-8"))
+                return ct.target(compile_spec_value)
+
+        return ct.target.iOS15
+
+    @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_compile_specs(
+        compute_unit: ct.ComputeUnit = ct.ComputeUnit.ALL,
+        minimum_deployment_target: ct.target = ct.target.iOS15,
+        compute_precision: ct.precision = ct.precision.FLOAT16,
+        model_type: MODEL_TYPE = MODEL_TYPE.MODEL,
+    ) -> 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))
+
+        return compile_specs
+
+    @staticmethod
+    def model_metadata_from_spec(
+        model_spec: ct.proto.Model_pb2, identifier: str
+    ) -> Dict[str, str]:
+        input_names: List[str] = [input.name for input in model_spec.description.input]
+        output_names = [output.name for output in model_spec.description.output]
+
+        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: ModelMetadata, 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 preprocess_model(
+        mlmodel: ct.models.MLModel, model_type: MODEL_TYPE
+    ) -> PreprocessResult:
+        identifier = str(uuid.uuid4())
+        dir_path: Path = Path("tmp") / identifier
+        model_dir_path: Path = dir_path / "lowered_module"
+        model_spec: ct.proto.Model_pb2 = mlmodel.get_spec()
+        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(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())
+        )
+
+        debug_handle_map: Optional[Dict[str, Tuple[int]]] = None
+        if model_debug_info is not None:
+            debug_handle_map = model_debug_info.debugSymbolToHandles
+
+        shutil.rmtree(str(dir_path.resolve()))
+        return PreprocessResult(
+            processed_bytes=processed_bytes,
+            debug_handle_map=debug_handle_map,
+        )
+
+    @classmethod
+    def preprocess(
+        cls,
+        edge_program: ExportedProgram,
+        module_compile_specs: List[CompileSpec],
+    ) -> PreprocessResult:
+        model_type: CoreMLBackend.MODEL_TYPE = (
+            CoreMLBackend.model_type_from_compile_specs(
+                module_compile_specs,
+            )
+        )
+
+        model_compute_precision: ct.precision = (
+            CoreMLBackend.model_compute_precision_from_compile_specs(
+                module_compile_specs
+            )
+        )
+
+        minimum_deployment_target: ct.target = (
+            CoreMLBackend.min_deployment_target_from_compile_specs(module_compile_specs)
+        )
+
+        mlmodel = ct.convert(
+            model=edge_program,
+            source="pytorch",
+            convert_to="mlprogram",
+            pass_pipeline=ct.PassPipeline.DEFAULT,
+            skip_model_load=False,
+            compute_precision=model_compute_precision,
+            minimum_deployment_target=minimum_deployment_target,
+        )
+
+        return CoreMLBackend.preprocess_model(mlmodel, model_type=model_type)

executorch/backends/apple/coreml/partition/coreml_partitioner.py

@@ -0,0 +1,95 @@
+# Copyright © 2023 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+import logging
+from typing import List, Optional
+
+import coremltools as ct
+
+import torch
+
+from executorch.backends.apple.coreml.compiler import CoreMLBackend
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.backend.utils import tag_constant_data
+from torch.export.exported_program import ExportedProgram
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+from torch.fx.passes.operator_support import OperatorSupportBase
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class OperatorsSupportedForCoreMLBackend(OperatorSupportBase):
+    def __init__(
+        self, skip_ops_for_coreml_delegation: Optional[List[str]] = None
+    ) -> None:
+        if skip_ops_for_coreml_delegation is None:
+            skip_ops_for_coreml_delegation = []
+        super().__init__()
+        self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
+
+    def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
+        # get_attr node can always be supported on any backend
+        if node.op == "get_attr":
+            return True
+        # check if the PyTorch op get called is supported in Core ML
+        elif node.op == "call_function":
+            # skip ops if specified by user
+            node_target_name = node.target.__name__.lower()
+            if node_target_name in self.skip_ops_for_coreml_delegation:
+                return False
+            # query coremltools to see if node is supported
+            return ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)
+        # cowardly refuse to support all other types of node:
+        # 1. placeholder / output nodes should not be tagged
+        #    reference: https://github.com/pytorch/executorch/pull/1398
+        # 2. call_module / call_method should have been replaced with call_function?
+        else:
+            return False
+
+
+class CoreMLPartitioner(Partitioner):
+
+    def __init__(
+        self,
+        skip_ops_for_coreml_delegation: Optional[List[str]] = None,
+        compile_specs: Optional[List[CompileSpec]] = None,
+    ) -> None:
+        if skip_ops_for_coreml_delegation is None:
+            skip_ops_for_coreml_delegation = []
+        self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation
+        self.delegation_spec = DelegationSpec(
+            backend_id=CoreMLBackend.__name__,
+            compile_specs=compile_specs if compile_specs is not None else [],
+        )
+
+    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
+        # Run the CapabilityBasedPartitioner to return the largest possible
+        # subgraphs containing the nodes with the tags
+        logger.info("CoreMLPartitioner::partition")
+        partition_tags = {}
+
+        capability_partitioner = CapabilityBasedPartitioner(
+            exported_program.graph_module,
+            OperatorsSupportedForCoreMLBackend(self.skip_ops_for_coreml_delegation),
+            allows_single_node_partition=True,
+        )
+        partition_list = capability_partitioner.propose_partitions()
+        for partition in partition_list:
+            for node in partition.nodes:
+                tag = f"tag{partition.id}"
+                node.meta["delegation_tag"] = tag
+                partition_tags[tag] = self.delegation_spec
+
+        tag_constant_data(exported_program)
+
+        return PartitionResult(
+            tagged_exported_program=exported_program, partition_tags=partition_tags
+        )

executorch/backends/apple/coreml/quantizer/coreml_quantizer.py

@@ -0,0 +1,7 @@
+# Copyright © 2024 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+from coremltools.optimize.torch.quantization._coreml_quantizer import (  # noqa: FLAKE8 F401
+    CoreMLQuantizer,
+)

executorch/backends/apple/coreml/runtime/inmemoryfs/setup.py

@@ -0,0 +1,52 @@
+import os
+
+import pathlib
+import sys
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent.parent.parent.parent
+PYBIND11_DIR_PATH = REPO_ROOT / "third-party" / "pybind11"
+sys.path.append(str(PYBIND11_DIR_PATH.absolute()))
+
+from pybind11.setup_helpers import build_ext, Pybind11Extension
+from setuptools import setup
+
+__version__ = "0.0.1"
+
+cxx_std = int(os.environ.get("CMAKE_CXX_STANDARD", "17"))
+
+ext_modules = [
+    Pybind11Extension(
+        "executorchcoreml",
+        [
+            "../util/json_util.cpp",
+            "inmemory_filesystem.cpp",
+            "inmemory_filesystem_py.cpp",
+            "inmemory_filesystem_utils.cpp",
+            "memory_buffer.cpp",
+            "memory_stream.cpp",
+            "reversed_memory_stream.cpp",
+        ],
+        define_macros=[("VERSION_INFO", __version__)],
+        cxx_std=cxx_std,
+        extra_compile_args=["-mmacosx-version-min=10.15", "-g"],
+        include_dirs=[
+            "../../third-party/nlohmann_json/single_include/nlohmann",
+            ".",
+            "../util",
+        ],
+    ),
+]
+
+setup(
+    name="executorchcoreml",
+    version=__version__,
+    description="CoreML extension for executorch",
+    long_description="",
+    author="Apple Inc.",
+    ext_modules=ext_modules,
+    extras_require={"test": "pytest"},
+    cmdclass={"build_ext": build_ext},
+    include_package_data=True,
+    zip_safe=False,
+    python_requires=">=3.9",
+)

executorch/backends/apple/coreml/test/test_coreml_partitioner.py

@@ -0,0 +1,78 @@
+# Copyright © 2023 Apple Inc. All rights reserved.
+#
+# Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+import unittest
+
+import executorch.exir
+
+import torch
+import torchvision
+
+from executorch.backends.apple.coreml.partition.coreml_partitioner import (
+    CoreMLPartitioner,
+)
+
+
+class TestCoreMLPartitioner(unittest.TestCase):
+    def test_add_sub_skip_mm(self):
+        class Model(torch.nn.Module):
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a = z - a
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        model = Model()
+        model.eval()
+
+        example_inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+        exir_program_aten = torch.export.export(model, example_inputs)
+        edge_program_manager = executorch.exir.to_edge(exir_program_aten)
+        delegated_program_manager = edge_program_manager.to_backend(
+            CoreMLPartitioner(skip_ops_for_coreml_delegation=["aten.mm.default"])
+        )
+
+        assert [
+            node.target.__name__
+            for node in delegated_program_manager.exported_program().graph.nodes
+            if node.op == "call_function"
+        ] == [
+            "aten.mm.default",
+            "executorch_call_delegate",
+            "getitem",
+            "aten.mm.default",
+            "executorch_call_delegate",
+            "getitem",
+        ]
+
+    def test_vit_skip_conv(self):
+        model = torchvision.models.vit_b_16(weights="IMAGENET1K_V1")
+        model.eval()
+
+        example_inputs = (torch.randn(1, 3, 224, 224),)
+        exir_program_aten = torch.export.export(model, example_inputs)
+        edge_program_manager = executorch.exir.to_edge(exir_program_aten)
+        delegated_program_manager = edge_program_manager.to_backend(
+            CoreMLPartitioner(
+                skip_ops_for_coreml_delegation=["aten.convolution.default"]
+            )
+        )
+
+        assert [
+            node.target.__name__
+            for node in delegated_program_manager.exported_program().graph.nodes
+            if node.op == "call_function"
+        ] == [
+            "aten.convolution.default",
+            "executorch_call_delegate",
+            "getitem",
+        ]
+
+
+if __name__ == "__main__":
+    test_runner = TestCoreMLPartitioner()
+    test_runner.test_add_sub_skip_mm()
+    test_runner.test_vit_skip_conv()