From f1e673f608de89eb226da9e84fc1c7389493bf28 Mon Sep 17 00:00:00 2001
From: Chen Lai <chenlai@meta.com>
Date: Fri, 5 Jul 2024 12:19:27 -0700
Subject: [PATCH] add mps stories end to end in ci (#4137)

Summary:
Add the end to end test for stories fp32 in ci

Pull Request resolved: https://github.com/pytorch/executorch/pull/4137

Reviewed By: kimishpatel

Differential Revision: D59305037

Pulled By: cccclai

fbshipit-source-id: 9c132721e1120c00f4cb969b02d494879e181f32
---
 .ci/scripts/test_llama.sh   | 12 ++++++++++++
 .github/workflows/trunk.yml | 20 ++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/.ci/scripts/test_llama.sh b/.ci/scripts/test_llama.sh
index 18d0c67f88..a376f5ad2e 100644
--- a/.ci/scripts/test_llama.sh
+++ b/.ci/scripts/test_llama.sh
@@ -55,6 +55,14 @@ else
   QE=OFF
 fi
 
+if [[ "${MODE}" =~ .*mps.* ]]; then
+  MPS=ON
+else
+  MPS=OFF
+fi
+
+echo "MPS option ${MPS}"
+
 if [[ -z "${BUCK:-}" ]]; then
   BUCK=buck2
 fi
@@ -77,6 +85,7 @@ cmake_install_executorch_libraries() {
         -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
         -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
         -DEXECUTORCH_BUILD_XNNPACK="$XNNPACK" \
+        -DEXECUTORCH_BUILD_MPS="$MPS" \
         -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
         -Bcmake-out .
     cmake --build cmake-out -j9 --target install --config Debug
@@ -142,6 +151,9 @@ fi
 if [[ "${QE}" == "ON" ]]; then
   EXPORT_ARGS="${EXPORT_ARGS} --embedding-quantize 8,1024"
 fi
+if [[ "${MPS}" == "ON" ]]; then
+  EXPORT_ARGS="${EXPORT_ARGS} -kv -v --mps --disable_dynamic_shape"
+fi
 # Add dynamically linked library location
 $PYTHON_EXECUTABLE -m examples.models.llama2.export_llama ${EXPORT_ARGS}
 
diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml
index db5dc6cd38..9414003347 100644
--- a/.github/workflows/trunk.yml
+++ b/.github/workflows/trunk.yml
@@ -225,7 +225,7 @@ jobs:
       matrix:
         dtype: [fp32]
         build-tool: [buck2, cmake]
-        mode: [portable, xnnpack+kv+custom]
+        mode: [portable, xnnpack+kv+custom, mps]
       fail-fast: false
     with:
       runner: macos-m1-stable
@@ -234,15 +234,31 @@ jobs:
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       timeout: 900
       script: |
-        bash .ci/scripts/setup-conda.sh
 
         DTYPE=${{ matrix.dtype }}
         BUILD_TOOL=${{ matrix.build-tool }}
         MODE=${{ matrix.mode }}
 
+        if [[ "${BUILD_TOOL}" == "buck2" ]]; then
+          # TODO: Will add more modes that don't support buck2
+          if [[ "${MODE}" == "mps" ]]; then
+            echo "mps doesn't support buck2."
+            exit 0
+          fi
+        fi
+
+        bash .ci/scripts/setup-conda.sh
+
         # Setup executorch
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
 
+        if [[ "${MODE}" == "mps" ]]; then
+          PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/mps/install_requirements.sh
+          echo "Finishing installing mps."
+        else
+          echo "Not mps mode, skip installing mps."
+        fi
+
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash examples/models/llama2/install_requirements.sh
         # Test llama2