Enable MPS AOT for macOS by default

Author: jathu

.ci/scripts/test_model.sh

@@ -219,7 +219,7 @@ test_model_with_coreml() {
   DTYPE=float16
 
   "${PYTHON_EXECUTABLE}" -m examples.apple.coreml.scripts.export --model_name="${MODEL_NAME}" --compute_precision "${DTYPE}" --use_partitioner
-  EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*.pte" -print -quit)
+  EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*coreml*.pte" -print -quit)
 
   if [ -n "$EXPORTED_MODEL" ]; then
     EXPORTED_MODEL_WITH_DTYPE="${EXPORTED_MODEL%.pte}_${DTYPE}.pte"
@@ -244,8 +244,18 @@ test_model_with_coreml() {
 }
 
 test_model_with_mps() {
-  "${PYTHON_EXECUTABLE}" -m examples.apple.mps.scripts.mps_example --model_name="${MODEL_NAME}" --use_fp16
-  EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*.pte" -print -quit)
+  "${PYTHON_EXECUTABLE}" -m examples.apple.mps.scripts.mps_example --model_name="${MODEL_NAME}" --bundled --use_fp16
+  EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*mps*.pte" -print -quit)
+
+  if [ -z "${EXPORTED_MODEL}" ]; then
+    echo "[error] failed to export model: no .pte file found"
+    exit 1
+  fi
+
+  echo "Testing exported model with mps_executor_runner..."
+  local out_dir=$(mktemp -d)
+  examples/apple/mps/scripts/build_mps_executor_runner.sh --output="${out_dir}"
+  "${out_dir}/examples/apple/mps/mps_executor_runner" --bundled_program --model_path "${EXPORTED_MODEL}"
 }
 
 if [[ "${BACKEND}" == "portable" ]]; then

.ci/scripts/wheel/test_macos.py

@@ -11,13 +11,17 @@
 if __name__ == "__main__":
     test_base.run_tests(
         model_tests=[
+            # test_base.ModelTest(
+            #     model=Model.Mv3,
+            #     backend=Backend.XnnpackQuantizationDelegation,
+            # ),
+            # test_base.ModelTest(
+            #     model=Model.Mv3,
+            #     backend=Backend.CoreMlExportAndTest,
+            # ),
             test_base.ModelTest(
                 model=Model.Mv3,
-                backend=Backend.XnnpackQuantizationDelegation,
-            ),
-            test_base.ModelTest(
-                model=Model.Mv3,
-                backend=Backend.CoreMlExportAndTest,
+                backend=Backend.AppleMPS,
             ),
         ]
     )

.github/workflows/trunk.yml

@@ -486,7 +486,7 @@ jobs:
     strategy:
       fail-fast: false
     with:
-      runner: macos-m1-stable
+      runner: macos-latest-xlarge
       python-version: '3.11'
       submodules: 'recursive'
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}

.gitignore

@@ -25,6 +25,7 @@ pip-out/
 *.bin
 *.model
 *.pte
+*.etdp
 !test_bpe_tokenizer.bin
 !test_tiktoken_tokenizer.model
 

CMakePresets.json

@@ -15,7 +15,7 @@
         "CMAKE_TOOLCHAIN_FILE": "${sourceDir}/third-party/ios-cmake/ios.toolchain.cmake",
         "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/macos.cmake",
         "PLATFORM": "MAC_ARM64",
-        "DEPLOYMENT_TARGET": "10.15"
+        "DEPLOYMENT_TARGET": "12.0"
       },
       "condition": {
         "lhs": "${hostSystemName}",
@@ -77,7 +77,7 @@
       "inherits": ["common"],
       "cacheVariables": {
         "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/pybind.cmake",
-        "CMAKE_OSX_DEPLOYMENT_TARGET": "10.15"
+        "CMAKE_OSX_DEPLOYMENT_TARGET": "12.0"
       },
       "condition": {
         "type": "inList",
@@ -93,7 +93,7 @@
         ],
         "cacheVariables": {
             "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/llm.cmake",
-            "CMAKE_OSX_DEPLOYMENT_TARGET": "10.15"
+            "CMAKE_OSX_DEPLOYMENT_TARGET": "12.0"
         },
         "condition": {
             "type": "inList",

Package.swift

@@ -120,7 +120,7 @@ let package = Package(
   name: "executorch",
   platforms: [
     .iOS(.v17),
-    .macOS(.v10_15),
+    .macOS(.v12),
   ],
   products: packageProducts,
   targets: packageTargets + [

examples/apple/mps/executor_runner/mps_executor_runner.mm

@@ -36,7 +36,7 @@
 #include <chrono>
 using namespace std::chrono;
 
-static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4 MB
+static uint8_t method_allocator_pool[16 * 1024U * 1024U]; // 16 MB
 
 DEFINE_string(model_path, "model.ff", "Model serialized in flatbuffer format.");
 DEFINE_string(
@@ -306,12 +306,16 @@ HierarchicalAllocator planned_memory(
   ET_LOG(Info, "Inputs prepared.");
 
   int num_iterations = FLAGS_num_runs + (FLAGS_skip_warmup ? 0 : 1);
+
+  ET_LOG(Info, "[jathu] num_iterations: %i", num_iterations);
   std::vector<float> exec_times;
-  exec_times.reserve(FLAGS_num_runs);
+  exec_times.reserve(num_iterations);
   for (int i = 0; i < num_iterations; i++) {
     auto start_exec_time = high_resolution_clock::now();
     // Run the model.
+    ET_LOG(Info, "[jathu] i: %i executing method...", i);
     Error status = method->execute();
+    ET_LOG(Info, "[jathu] i: %i done executing method", i);
     auto end_exec_time = high_resolution_clock::now();
     auto duration = duration_cast<microseconds>(end_exec_time - start_exec_time);
     exec_times.push_back(duration.count());

examples/apple/mps/scripts/build_mps_executor_runner.sh

@@ -1,8 +1,15 @@
 #!/usr/bin/env bash
 #  Copyright (c) 2024 Apple Inc. All rights reserved.
 #  Provided subject to the LICENSE file in the top level directory.
+#
+# 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.
 
-set -e
+
+set -eux
 
 MODE="Release"
 OUTPUT="cmake-out"
@@ -38,25 +45,29 @@ done
 
 rm -rf "$OUTPUT"
 
-cmake -DBUCK2="$BUCK" \
-          -DCMAKE_INSTALL_PREFIX=cmake-out \
-          -DCMAKE_BUILD_TYPE="$MODE" \
-          -DEXECUTORCH_BUILD_DEVTOOLS=ON \
-          -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
-          -DEXECUTORCH_BUILD_MPS=ON \
-          -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
-          -Bcmake-out .
-cmake --build cmake-out -j9 --target install --config "$MODE"
-CMAKE_PREFIX_PATH="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags"
-# build mps_executor_runner
-rm -rf cmake-out/examples/apple/mps
-cmake \
-    -DCMAKE_PREFIX_PATH="$CMAKE_PREFIX_PATH" \
-    -DCMAKE_BUILD_TYPE="$MODE" \
-    -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
-    -Bcmake-out/examples/apple/mps \
-    examples/apple/mps
-
-cmake --build cmake-out/examples/apple/mps -j9 --config "$MODE"
+cmake -DCMAKE_INSTALL_PREFIX=${OUTPUT} \
+      -DCMAKE_BUILD_TYPE=${MODE} \
+      -DEXECUTORCH_ENABLE_LOGGING=ON \
+      -DEXECUTORCH_BUILD_DEVTOOLS=ON \
+      -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
+      -Dprotobuf_BUILD_TESTS=OFF \
+      -B ${OUTPUT} \
+      --preset macos
+
+cmake --build ${OUTPUT} \
+      -j $(sysctl -n hw.ncpu) \
+      --config ${MODE} \
+      --target install
+
+cmake -DCMAKE_PREFIX_PATH="${OUTPUT}/lib/cmake/ExecuTorch;${OUTPUT}/third-party/gflags" \
+      -DCMAKE_BUILD_TYPE="$MODE" \
+      -DCMAKE_OSX_DEPLOYMENT_TARGET="12.0" \
+      -B "${OUTPUT}/examples/apple/mps" \
+      -S examples/apple/mps
+
+cmake --build "${OUTPUT}/examples/apple/mps" \
+      -j $(sysctl -n hw.ncpu) \
+      --config ${MODE} \
+      --target mps_executor_runner
 
 echo "Build succeeded!"

examples/models/__init__.py

@@ -46,6 +46,7 @@ class Backend(str, Enum):
     XnnpackQuantizationDelegation = "xnnpack-quantization-delegation"
     CoreMlExportOnly = "coreml"
     CoreMlExportAndTest = "coreml-test"  # AOT export + test with runner
+    AppleMPS = "mps"
 
     def __str__(self) -> str:
         return self.value

tools/cmake/preset/pybind.cmake

@@ -22,6 +22,7 @@ set_overridable_option(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON)
 
 if(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
   set_overridable_option(EXECUTORCH_BUILD_COREML ON)
+  set_overridable_option(EXECUTORCH_BUILD_MPS ON)
 elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
   set_overridable_option(EXECUTORCH_BUILD_COREML ON)
 elseif(CMAKE_SYSTEM_NAME STREQUAL "Windows" OR CMAKE_SYSTEM_NAME STREQUAL "WIN32")