Buckify Core ML AOT components. (#3701)

Summary:
Pull Request resolved: #3701

.

Reviewed By: kirklandsign

Differential Revision: D57650038

fbshipit-source-id: 581f1afe14262a5fc32117fffadf3b10a5a75f41

Author: shoumikhin

backends/apple/coreml/TARGETS

@@ -1 +1,100 @@
-# This file needs to exist to avoid build system breakage, see https://fburl.com/workplace/jtdlgdmd
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+runtime.python_library(
+    name = "backend",
+    srcs = glob([
+        "compiler/*.py",
+    ]),
+    visibility = [
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        ":executorchcoreml",
+        "//executorch/exir/backend:backend_details",
+        "//executorch/exir/backend:compile_spec_schema",
+        "fbsource//third-party/pypi/coremltools:coremltools",
+    ],
+)
+
+runtime.python_library(
+    name = "partitioner",
+    srcs = glob([
+        "partition/*.py",
+    ]),
+    visibility = [
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        ":backend",
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:compile_spec_schema",
+        "//executorch/exir/backend:partitioner",
+        "//executorch/exir/backend:utils",
+        "fbsource//third-party/pypi/coremltools:coremltools",
+    ],
+)
+
+runtime.python_library(
+    name = "quantizer",
+    srcs = glob([
+        "quantizer/*.py",
+    ]),
+    visibility = [
+        "@EXECUTORCH_CLIENTS",
+    ],
+)
+
+runtime.cxx_python_extension(
+    name = "executorchcoreml",
+    srcs = [
+        "runtime/inmemoryfs/inmemory_filesystem.cpp",
+        "runtime/inmemoryfs/inmemory_filesystem_py.cpp",
+        "runtime/inmemoryfs/inmemory_filesystem_utils.cpp",
+        "runtime/inmemoryfs/memory_buffer.cpp",
+        "runtime/inmemoryfs/memory_stream.cpp",
+        "runtime/inmemoryfs/reversed_memory_stream.cpp",
+        "runtime/util/json_util.cpp",
+    ],
+    headers = glob([
+        "runtime/inmemoryfs/**/*.hpp",
+    ]),
+    preprocessor_flags = [
+        "-Iexecutorch/backends/apple/coreml/runtime/util",
+    ],
+    types = [
+        "executorchcoreml.pyi",
+    ],
+    compiler_flags = [
+        "-std=c++17",
+    ],
+    base_module = "",
+    visibility = [
+        "//executorch/examples/apple/coreml/...",
+    ],
+    external_deps = [
+        "pybind11",
+    ],
+    deps = [
+        "fbsource//third-party/nlohmann-json:nlohmann-json",
+    ],
+)
+
+runtime.python_test(
+    name = "test",
+    srcs = glob([
+        "test/*.py",
+    ]),
+    deps = [
+        ":partitioner",
+        ":quantizer",
+        "//caffe2:torch",
+        "//pytorch/vision:torchvision",
+        "fbsource//third-party/pypi/pytest:pytest",
+    ],
+)

backends/apple/coreml/compiler/coreml_preprocess.py

@@ -184,8 +184,8 @@ def generate_compile_specs(
 
     @staticmethod
     def model_metadata_from_spec(
-        model_spec: ct.proto.Model_pb2, identifier: str
-    ) -> Dict[str, str]:
+        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]
 
@@ -300,7 +300,7 @@ def preprocess_model(
 
         # Save model.
         model_path = model_dir_path / MODEL_PATHS.MODEL.value
-        mlmodel.save(model_path)
+        mlmodel.save(str(model_path))
         # Extract delegate mapping file.
         model_debug_info: Optional[ModelDebugInfo] = CoreMLBackend.get_model_debug_info(
             model_path
@@ -327,40 +327,41 @@ def preprocess_model(
                 model_debug_info=model_debug_info, model_dir_path=model_dir_path
             )
 
-        processed_bytes: bytes = executorchcoreml.flatten_directory_contents(
-            str(model_dir_path.resolve())
+        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 = model_debug_info.debugSymbolToHandles
+            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,
         )
 
-    @classmethod
+    @staticmethod
     def preprocess(
-        cls,
         edge_program: ExportedProgram,
-        module_compile_specs: List[CompileSpec],
+        compile_specs: List[CompileSpec],
     ) -> PreprocessResult:
         model_type: CoreMLBackend.MODEL_TYPE = (
             CoreMLBackend.model_type_from_compile_specs(
-                module_compile_specs,
+                compile_specs,
             )
         )
 
         model_compute_precision: ct.precision = (
-            CoreMLBackend.model_compute_precision_from_compile_specs(
-                module_compile_specs
-            )
+            CoreMLBackend.model_compute_precision_from_compile_specs(compile_specs)
         )
 
         minimum_deployment_target: ct.target = (
-            CoreMLBackend.min_deployment_target_from_compile_specs(module_compile_specs)
+            CoreMLBackend.min_deployment_target_from_compile_specs(compile_specs)
         )
 
         mlmodel = ct.convert(

backends/apple/coreml/executorchcoreml.pyi

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+def flatten_directory_contents(path: str) -> Optional[bytes]: ...
+def unflatten_directory_contents(bytes: bytes, path: str) -> bool: ...

backends/apple/coreml/partition/coreml_partitioner.py

@@ -42,8 +42,8 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         # 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:
+            node_target_name = getattr(node.target, "__name__", "").lower()
+            if node_target_name in (self.skip_ops_for_coreml_delegation or []):
                 return False
             # query coremltools to see if node is supported
             return ct.converters.mil.frontend.torch.is_torch_fx_node_supported(node)

backends/apple/coreml/test/test_coreml_quantizer.py

@@ -51,6 +51,8 @@ def quantize_and_compare(
             prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer)
         elif quantization_type == "QAT":
             prepared_graph = prepare_qat_pt2e(pre_autograd_aten_dialect, quantizer)
+        else:
+            raise ValueError("Invalid quantization type")
 
         prepared_graph(*example_inputs)
         converted_graph = convert_pt2e(prepared_graph)

examples/models/llama2/TARGETS

@@ -82,23 +82,26 @@ runtime.python_library(
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/backends/apple/coreml:backend",
+        "//executorch/backends/apple/coreml:partitioner",
         "//executorch/backends/transforms:duplicate_dynamic_quant_chain",
+        "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
         "//executorch/backends/xnnpack:xnnpack_backend",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
-        "//executorch/backends/vulkan/partitioner:vulkan_partitioner",
         "//executorch/examples/models:model_base",
         "//executorch/examples/models:models",
         "//executorch/examples/models/llama2/custom_ops:custom_ops_aot_py",
         "//executorch/examples/portable:utils",
         "//executorch/exir:lib",
-        "//executorch/sdk/etrecord:etrecord",
-        "//executorch/util:memory_profiler",
-        "//executorch/util:python_profiler",
         # one definition has to be included in the user of the libarary
         # depending on what library the client wants to use
         # "//executorch/extension/pybindings:aten_lib",
         # "//executorch/extension/pybindings:portable_lib",
         # "//executorch/extension/pybindings:portable_lib_plus_custom",
+        "//executorch/sdk/etrecord:etrecord",
+        "//executorch/util:memory_profiler",
+        "//executorch/util:python_profiler",
+        "fbsource//third-party/pypi/coremltools:coremltools",
         "fbsource//third-party/pypi/sentencepiece:sentencepiece",
     ],
 )

examples/models/llama2/lib/partitioner_lib.py

@@ -56,13 +56,8 @@ def get_coreml_partitioner(args):
         args.use_kv_cache is True
     ), "CoreML backend currently only supports static shape and use_kv_cache=True is the only way to support it at the moment"
     try:
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `coremltools`.
         import coremltools as ct
-
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.compiler`
         from executorch.backends.apple.coreml.compiler import CoreMLBackend
-
-        # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition`
         from executorch.backends.apple.coreml.partition import CoreMLPartitioner
     except ImportError:
         raise ImportError(

exir/backend/TARGETS

@@ -109,5 +109,6 @@ runtime.python_library(
         "fbsource//third-party/pypi/pandas:pandas",
         "//caffe2:torch",
         "//executorch/exir:lowered_backend_module",
+        "//executorch/exir/backend/canonical_partitioners:duplicate_constant_node_pass",
     ],
 )