Cover MPS models in the CI (#6036)

Summary: Pull Request resolved: #6036

Reviewed By: kirklandsign

Differential Revision: D64126006

Pulled By: guangy10

fbshipit-source-id: c5aa5a69b9371391f8c300d8f0ee3409ffc76a1f
(cherry picked from commit 554eebe)

Author: Guang Yang

.ci/scripts/test_model.sh

@@ -191,6 +191,11 @@ test_model_with_coreml() {
   EXPORTED_MODEL=$(find "." -type f -name "${MODEL_NAME}*.pte" -print -quit)
 }
 
+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)
+}
+
 if [[ "${BACKEND}" == "portable" ]]; then
   echo "Testing ${MODEL_NAME} with portable kernels..."
   test_model
@@ -206,6 +211,12 @@ elif [[ "${BACKEND}" == "coreml" ]]; then
   if [[ $? -eq 0 ]]; then
     prepare_artifacts_upload
   fi
+elif [[ "${BACKEND}" == "mps" ]]; then
+  echo "Testing ${MODEL_NAME} with mps..."
+  test_model_with_mps
+  if [[ $? -eq 0 ]]; then
+    prepare_artifacts_upload
+  fi
 elif [[ "${BACKEND}" == "xnnpack" ]]; then
   echo "Testing ${MODEL_NAME} with xnnpack..."
   WITH_QUANTIZATION=true

.github/workflows/apple-perf.yml

@@ -169,6 +169,8 @@ jobs:
             DELEGATE_CONFIG="xnnpack+custom+qe"
           elif [[ ${{ matrix.delegate }} == "coreml" ]]; then
             DELEGATE_CONFIG="coreml"
+          elif [[ ${{ matrix.delegate }} == "mps" ]]; then
+            DELEGATE_CONFIG="mps"
           fi
           PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
             bash .ci/scripts/test_llama.sh "${{ matrix.model }}" "${BUILD_MODE}" "${DTYPE}" "${DELEGATE_CONFIG}" "${ARTIFACTS_DIR_NAME}"

.github/workflows/trunk.yml

@@ -322,9 +322,9 @@ jobs:
         PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh
         PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn"
 
-  test-coreml-model:
-    name: test-coreml-model
-    uses: pytorch/test-infra/.github/workflows/macos_job.yml@release/2.5
+  test-apple-model:
+    name: test-apple-model
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
     strategy:
       fail-fast: false
     with:
@@ -335,20 +335,25 @@ jobs:
       timeout: 90
       script: |
         BUILD_TOOL=cmake
-        BACKEND=coreml
 
         bash .ci/scripts/setup-conda.sh
 
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/coreml/scripts/install_requirements.sh
         echo "Finishing installing coreml."
+        PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/mps/install_requirements.sh
+        echo "Finishing installing mps."
 
         # Build and test coreml model
         MODELS=(mv3 ic4 resnet50 edsr mobilebert w2l)
         for MODEL_NAME in "${MODELS[@]}"; do
           echo "::group::Exporting coreml model: $MODEL_NAME"
-          PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}"
+          PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "coreml"
+          echo "::endgroup::"
+
+          echo "::group::Exporting mps model: $MODEL_NAME"
+          PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "mps"
           echo "::endgroup::"
         done
 

examples/apple/mps/scripts/mps_example.py

@@ -195,18 +195,15 @@ def get_model_config(args):
             edge_compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
         ).to_executorch(config=ExecutorchBackendConfig(extract_delegate_segments=False))
 
-    model_name = f"{args.model_name}_mps"
+    dtype = "float16" if args.use_fp16 else "float32"
+    model_name = f"{args.model_name}_mps_{dtype}"
 
     if args.bundled:
         expected_output = model(*example_inputs)
         bundled_program_buffer = get_bundled_program(
             executorch_program, example_inputs, expected_output
         )
-        model_name = f"{model_name}_bundled"
-        extension = "fp16"
-        if not args.use_fp16:
-            extension = "fp32"
-        model_name = f"{model_name}_{extension}.pte"
+        model_name = f"{model_name}_bundled.pte"
 
     if args.generate_etrecord:
         etrecord_path = "etrecord.bin"