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="$(pwd)/${MODEL_NAME}_mps_float16_bundled.pte"
+
+  if [[ ! -f "${EXPORTED_MODEL}" ]]; then
+    echo "[error] exported model not found: ${EXPORTED_MODEL}"
+    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

.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
 

backends/xnnpack/third-party/cpuinfo

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

examples/apple/mps/executor_runner/mps_executor_runner.mm renamed to examples/apple/mps/executor_runner/mps_executor_runner.cpp

@@ -17,7 +17,7 @@ def define_common_targets():
         runtime.cxx_binary(
             name = "mps_executor_runner",
             srcs = [
-                "mps_executor_runner.mm",
+                "mps_executor_runner.cpp",
             ],
             deps = [
                 "//executorch/backends/apple/mps:mps_schema",

examples/apple/mps/executor_runner/targets.bzl

@@ -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,26 @@ 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 \
+      -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/apple/mps/scripts/build_mps_executor_runner.sh

@@ -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