Update CI script to run test_mps

.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,12 +35,36 @@ 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-macos:
+  #   name: test-mps-models-macos
+  #   runs-on: macos-executorch
+  #   needs: gather-models-mps
+  #   strategy:
+  #     matrix: ${{ fromJSON(needs.gather-models-mps.outputs.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}"
+  #         popd
+  test-mps-macos:
+    name: test-mps-macos
     runs-on: macos-executorch
     needs: gather-models-mps
     strategy:
-      matrix: ${{ fromJSON(needs.gather-models-mps.outputs.models) }}
+      matrix: ${{ fromJSON(needs.gather-models-mps.outputs.mps_models) }}
       fail-fast: false
     steps:
       - uses: actions/checkout@v3
@@ -56,5 +81,5 @@ 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}" 1
           popd

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

@@ -9,6 +9,10 @@
 
 from examples.models import MODEL_NAME_TO_MODEL
 from examples.xnnpack import MODEL_NAME_TO_OPTIONS
+import unittest
+
+MPS_TEST_SUITE_PATH = 'backends/apple/mps/test'
+MPS_SUITE = unittest.defaultTestLoader.discover(MPS_TEST_SUITE_PATH)
 
 BUILD_TOOLS = {
     "cmake": {"macos-14"},
@@ -43,6 +47,17 @@ def set_output(name: str, val: Any) -> None:
     else:
         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:
     """
@@ -55,9 +70,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 +84,26 @@ 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
+
+            record = {
+                "build-tool": build_tool,
+                "model": testcase._testMethodName,
+                "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

@@ -134,7 +134,7 @@ def forward(self, *args):
         MethodTestSuite(
             method_name="forward",
             test_cases=[
-                MethodTestCase(inputs=m_inputs, expected_outputs=model(*m_inputs))
+                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)