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,17 @@ 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..."
+  ./examples/apple/mps/scripts/build_mps_executor_runner.sh
+  ./cmake-out/examples/apple/mps/mps_executor_runner -bundled_program -model_path "${EXPORTED_MODEL}"
 }
 
 if [[ "${BACKEND}" == "portable" ]]; then

.ci/scripts/wheel/test_macos.py

@@ -19,5 +19,9 @@
                 model=Model.Mv3,
                 backend=Backend.CoreMlExportAndTest,
             ),
+            test_base.ModelTest(
+                model=Model.Mv3,
+                backend=Backend.AppleMPS,
+            ),
         ]
     )

.gitignore

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

examples/apple/mps/scripts/build_mps_executor_runner.sh

@@ -1,6 +1,13 @@
 #!/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
 
@@ -44,19 +51,18 @@ cmake -DBUCK2="$BUCK" \
           -DEXECUTORCH_BUILD_DEVTOOLS=ON \
           -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
           -DEXECUTORCH_BUILD_MPS=ON \
-          -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
+          -DEXECUTORCH_ENABLE_LOGGING=ON \
           -Bcmake-out .
-cmake --build cmake-out -j9 --target install --config "$MODE"
+cmake --build cmake-out -j$(sysctl -n hw.ncpu) --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 --build cmake-out/examples/apple/mps -j$(sysctl -n hw.ncpu) --config "$MODE"
 
 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")