Update CI script to run test_mps

.github/workflows/lint_mps.yml

@@ -0,0 +1,61 @@
+name: Lint MPS
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      - release/*
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
+  cancel-in-progress: true
+
+jobs:
+  lintrunner-mps:
+    runs-on: macos-executorch
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          submodules: 'true'
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+          cache: pip
+      - name: Lintrunner
+        id: lint-mps
+        run: |
+          set -eux
+          source .ci/scripts/utils.sh
+          # This is a simple Python script but as it tries to import executorch.examples.models,
+          # it requires a whole bunch of ExecuTorch dependencies on the Docker image
+          install_pip_dependencies
+          install_executorch
+          brew install jq
+
+          source .ci/scripts/utils.sh
+
+          # Install lint depdendencies
+          pip install lintrunner==0.11.0
+          pip install lintrunner-adapters==0.11.0
+
+          CACHE_DIRECTORY="/tmp/.lintbin"
+          # Try to recover the cached binaries
+          if [[ -d "${CACHE_DIRECTORY}" ]]; then
+            # It's ok to fail this as lintrunner init would download these binaries
+            # again if they do not exist
+            cp -r "${CACHE_DIRECTORY}" . || true
+          fi
+
+          # This has already been cached in the docker image
+          lintrunner init 2> /dev/null
+
+          RC=0
+          # Run lintrunner on all files
+          if ! lintrunner --force-color --all-files --tee-json lint.json 2> /dev/null; then
+            echo ""
+            echo -e "\e[1m\e[36mYou can reproduce these results locally by using \`lintrunner\`. (If you don't get the same results, run \'lintrunner init\' to update your local linter)\e[0m"
+            echo -e "\e[1m\e[36mSee https://github.com/pytorch/pytorch/wiki/lintrunner for setup instructions.\e[0m"
+            RC=1
+          fi

.github/workflows/pull_mps.yml

@@ -16,6 +16,7 @@ jobs:
     runs-on: macos-executorch
     outputs:
       models: ${{ steps.gather-models-mps.outputs.models }}
+      mps_models: ${{ steps.gather-models-mps.outputs.mps_models }}
     steps:
       - uses: actions/checkout@v3
         with:
@@ -34,8 +35,8 @@ jobs:
           install_pip_dependencies
           install_executorch
           PYTHONPATH="${PWD}" python -m backends.apple.mps.ci.scripts.gather_test_models_mps
-  test-mps-delegate-macos:
-    name: test-mps-delegate-macos
+  test-mps-models:
+    name: test-mps-models
     runs-on: macos-executorch
     needs: gather-models-mps
     strategy:
@@ -56,5 +57,29 @@ jobs:
           PYTHON_EXECUTABLE=python bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
           PYTHON_EXECUTABLE=python bash backends/apple/mps/install_requirements.sh
           # Build and test ExecuTorch
-          PYTHON_EXECUTABLE=python bash backends/apple/mps/ci/scripts/test-mps.sh "${MODEL_NAME}" "${BUILD_TOOL}"
+          PYTHON_EXECUTABLE=python bash backends/apple/mps/ci/scripts/test-mps.sh "${MODEL_NAME}" "${BUILD_TOOL}" false
+          popd
+  test-mps:
+    name: test-mps
+    runs-on: macos-executorch
+    needs: gather-models-mps
+    strategy:
+      matrix: ${{ fromJSON(needs.gather-models-mps.outputs.mps_models) }}
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          submodules: 'true'
+      - name: Run test ${{ matrix.test }}
+        if: always()
+        run: |
+          WORKSPACE=$(pwd)
+          pushd "${WORKSPACE}"
+          MODEL_NAME=${{ matrix.model }}
+          BUILD_TOOL=${{ matrix.build-tool }}
+          # Setup MacOS dependencies as there is no Docker support on MacOS atm
+          PYTHON_EXECUTABLE=python bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
+          PYTHON_EXECUTABLE=python bash backends/apple/mps/install_requirements.sh
+          # Build and test ExecuTorch
+          PYTHON_EXECUTABLE=python bash backends/apple/mps/ci/scripts/test-mps.sh "${MODEL_NAME}" "${BUILD_TOOL}" true
           popd

backends/apple/mps/ci/scripts/gather_test_models_mps.py

@@ -5,10 +5,13 @@
 
 import json
 import os
+import unittest
 from typing import Any
 
 from examples.models import MODEL_NAME_TO_MODEL
-from examples.xnnpack import MODEL_NAME_TO_OPTIONS
+
+MPS_TEST_SUITE_PATH = "backends/apple/mps/test"
+MPS_SUITE = unittest.defaultTestLoader.discover(MPS_TEST_SUITE_PATH)
 
 BUILD_TOOLS = {
     "cmake": {"macos-14"},
@@ -17,6 +20,7 @@
     "macos-14": "macos-executorch",
 }
 
+
 def parse_args() -> Any:
     from argparse import ArgumentParser
 
@@ -44,6 +48,20 @@ def set_output(name: str, val: Any) -> None:
         print(f"::set-output name={name}::{val}")
 
 
+def gather_mps_test_list(suite, mps_test_list):
+    if hasattr(suite, "__iter__"):
+        for x in suite:
+            gather_mps_test_list(x, mps_test_list)
+    else:
+        mps_test_list.append(suite)
+
+
+def gather_mps_tests(suite):
+    mps_test_list = []
+    gather_mps_test_list(suite, mps_test_list)
+    return mps_test_list
+
+
 def export_models_for_ci() -> None:
     """
     This gathers all the example models that we want to test on GitHub OSS CI
@@ -55,9 +73,6 @@ def export_models_for_ci() -> None:
     # https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs
     models = {"include": []}
     for name in MODEL_NAME_TO_MODEL.keys():
-        delegation_configs = {
-            name in MODEL_NAME_TO_OPTIONS and MODEL_NAME_TO_OPTIONS[name].delegation,
-        }
         for build_tool in BUILD_TOOLS.keys():
             if target_os not in BUILD_TOOLS[build_tool]:
                 continue
@@ -72,6 +87,30 @@ def export_models_for_ci() -> None:
 
     set_output("models", json.dumps(models))
 
+    # Gather the list of all mps tests
+    mps_models = {"include": []}
+    mps_test_list = gather_mps_tests(MPS_SUITE)
+    start_path = ".".join([MPS_TEST_SUITE_PATH.replace("/", "."), "test_mps"])
+    for testcase in mps_test_list:
+        if "test_mps" not in str(testcase.__class__):
+            continue
+        for build_tool in BUILD_TOOLS.keys():
+            if target_os not in BUILD_TOOLS[build_tool]:
+                continue
+
+            cmd = ".".join(
+                [start_path, testcase.__class__.__name__, testcase._testMethodName]
+            )
+            record = {
+                "build-tool": build_tool,
+                "model": cmd,
+                "runner": DEFAULT_RUNNERS.get(target_os),
+            }
+
+            mps_models["include"].append(record)
+
+    set_output("mps_models", json.dumps(mps_models))
+
 
 if __name__ == "__main__":
     export_models_for_ci()

backends/apple/mps/ci/scripts/test-mps.sh

@@ -21,6 +21,12 @@ if [[ -z "${BUILD_TOOL:-}" ]]; then
   exit 1
 fi
 
+MPS_TESTS=$3
+if [[ -z "${MPS_TESTS:-}" ]]; then
+  echo "Missing flag specifying which tests to run, exiting..."
+  exit 1
+fi
+
 which "${PYTHON_EXECUTABLE}"
 CMAKE_OUTPUT_DIR=cmake-out
 
@@ -62,5 +68,28 @@ test_model_with_mps() {
   fi
 }
 
+test_mps_model() {
+  "${PYTHON_EXECUTABLE}" -m unittest "${MODEL_NAME}"
+
+  TEST_NAME="${MODEL_NAME##*.}"
+  OUTPUT_MODEL_PATH="${TEST_NAME}.pte"
+  STR_LEN=${#OUTPUT_MODEL_PATH}
+  FINAL_OUTPUT_MODEL_PATH=${OUTPUT_MODEL_PATH:5:$STR_LEN-5}
+
+  if [[ "${BUILD_TOOL}" == "cmake" ]]; then
+    if [[ ! -f ${CMAKE_OUTPUT_DIR}/examples/apple/mps/mps_executor_runner ]]; then
+      build_cmake_mps_executor_runner
+    fi
+    ./${CMAKE_OUTPUT_DIR}/examples/apple/mps/mps_executor_runner --model_path "${FINAL_OUTPUT_MODEL_PATH}" --bundled_program
+  else
+    echo "Invalid build tool ${BUILD_TOOL}. Only cmake is supported atm"
+    exit 1
+  fi
+}
+
 echo "Testing ${MODEL_NAME} with MPS..."
-test_model_with_mps
+if [[ "${MPS_TESTS}" == false ]]; then
+  test_model_with_mps
+else
+  test_mps_model
+fi

backends/apple/mps/test/test_mps.py

@@ -133,9 +133,7 @@ def forward(self, *args):
     method_test_suites = [
         MethodTestSuite(
             method_name="forward",
-            test_cases=[
-                MethodTestCase(inputs=m_inputs, expected_outputs=model(*m_inputs))
-            ],
+            test_cases=[MethodTestCase(inputs=m_inputs, expected_outputs=m(*m_inputs))],
         )
     ]
 

backends/apple/mps/test/test_mps_utils.py

@@ -195,7 +195,7 @@ def forward(self, *args):
                 method_name="forward",
                 test_cases=[
                     MethodTestCase(
-                        input=sample_inputs, expected_outputs=module(*sample_inputs)
+                        inputs=sample_inputs, expected_outputs=module(*sample_inputs)
                     )
                 ],
             )
@@ -208,7 +208,7 @@ def forward(self, *args):
             bundled_program
         )
 
-        filename = f"{func_name}.bpte"
+        filename = f"{func_name}.pte"
         logging.info(f"Step 5: Saving bundled program to {filename}...")
         with open(filename, "wb") as file:
             file.write(bundled_program_buffer)

examples/apple/mps/CMakeLists.txt

@@ -52,7 +52,8 @@ add_custom_command(
   OUTPUT ${_bundled_program_schema__outputs}
   COMMAND
     ${FLATC_EXECUTABLE} --cpp --cpp-std c++11 --gen-mutable --scoped-enums -o
-    "${_program_schema__include_dir}/executorch/sdk/bundled_program/schema" ${_bundled_program_schema__srcs}
+    "${_program_schema__include_dir}/executorch/sdk/bundled_program/schema"
+    ${_bundled_program_schema__srcs}
   WORKING_DIRECTORY ${EXECUTORCH_ROOT}
   DEPENDS ${FLATC_EXECUTABLE} ${_bundled_program_schema__srcs}
   COMMENT "Generating bundled_program headers"
@@ -68,8 +69,12 @@ set(mps_executor_runner_libs "-framework Foundation"
                               "-weak_framework MetalPerformanceShaders"
                               "-weak_framework MetalPerformanceShadersGraph"
                               "-weak_framework Metal")
-list(TRANSFORM _mps_executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/")
-add_executable(mps_executor_runner ${_mps_executor_runner__srcs} ${_bundled_program_schema__outputs})
+list(
+  TRANSFORM _mps_executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/")
+add_executable(
+  mps_executor_runner
+  ${_mps_executor_runner__srcs}
+  ${_bundled_program_schema__outputs})
 target_include_directories(
   mps_executor_runner INTERFACE ${CMAKE_BINARY_DIR}/schema/include/
                           ${EXECUTORCH_ROOT}/third-party/flatbuffers/include)