Metal backend: test modules

.github/workflows/metal.yml

@@ -28,9 +28,33 @@ jobs:
         PYTHON_EXECUTABLE=python CMAKE_ARGS="-DEXECUTORCH_BUILD_METAL=ON" ${CONDA_RUN} --no-capture-output ./install_executorch.sh
         echo "::endgroup::"
 
+  test-metal-modules:
+    name: test-metal-backend-modules
+    uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
+    with:
+      runner: macos-m2-stable
+      python-version: '3.11'
+      submodules: 'recursive'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 120
+      script: |
+        set -eux
+
+        echo "::group::Setup ExecuTorch"
+        PYTHON_EXECUTABLE=python ${CONDA_RUN} ./install_executorch.sh
+        echo "::endgroup::"
+
+        echo "::group::Build Metal Runtime"
+        ${CONDA_RUN} backends/apple/metal/tests/run_metal_test.sh --build
+        echo "::endgroup::"
+
+        echo "::group::Run Metal Backend Module Tests"
+        ${CONDA_RUN} python -m unittest backends.apple.metal.tests.test_modules.TestMetalBackendModules
+        echo "::endgroup::"
+
   export-model-metal-artifact:
     name: export-model-metal-artifact
-      # Skip this job if the pull request is from a fork (HuggingFace secrets are not available)
+    # Skip this job if the pull request is from a fork (HuggingFace secrets are not available)
     if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request'
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
     secrets: inherit

backends/apple/metal/tests/run_metal_test.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+# 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.
+
+# Script to build and run Metal backend tests
+# Usage:
+#   ./run_metal_test.sh --build              # Build the Metal runtime
+#   ./run_metal_test.sh --run <pte>          # Run inference with given model file
+#   ./run_metal_test.sh --check-build        # Check if runtime is already built
+
+set -e  # Exit on any error
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+EXECUTORCH_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
+BUILD_DIR="$EXECUTORCH_ROOT/cmake-out"
+EXECUTOR_RUNNER="$BUILD_DIR/executor_runner"
+
+# Function to check if Metal runtime is built
+check_build() {
+    if [[ -f "$EXECUTOR_RUNNER" ]]; then
+        echo "true"
+        return 0
+    else
+        echo "false"
+        return 1
+    fi
+}
+
+# Function to build the Metal runtime
+build_runtime() {
+    echo "Building Metal runtime..."
+
+    # Check if we're on macOS
+    if [[ "$(uname)" != "Darwin" ]]; then
+        echo "Error: Metal backend is only supported on macOS"
+        exit 1
+    fi
+
+    # Create build directory
+    mkdir -p "$BUILD_DIR"
+    cd "$BUILD_DIR"
+
+    # CMake configuration for Metal backend
+    CMAKE_ARGS="-DEXECUTORCH_BUILD_METAL=ON \
+                -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+                -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \
+                -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
+                -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+                -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+                -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON \
+                -DAOTI_METAL=ON \
+                -DEXECUTORCH_LOG_LEVEL=Info \
+                -DCMAKE_BUILD_TYPE=Release"
+
+    echo "Running cmake..."
+    cmake $CMAKE_ARGS "$EXECUTORCH_ROOT"
+
+    echo "Building..."
+    cmake --build . -j$(sysctl -n hw.ncpu)
+
+    cd "$EXECUTORCH_ROOT"
+
+    if [[ -f "$EXECUTOR_RUNNER" ]]; then
+        echo "Build successful: $EXECUTOR_RUNNER"
+    else
+        echo "Error: Build failed - executor_runner not found"
+        exit 1
+    fi
+}
+
+# Function to run inference
+run_inference() {
+    local pte_path="$1"
+
+    if [[ ! -f "$EXECUTOR_RUNNER" ]]; then
+        echo "Error: executor_runner not found at $EXECUTOR_RUNNER"
+        echo "Run '$0 --build' first to build the Metal runtime"
+        exit 1
+    fi
+
+    if [[ ! -f "$pte_path" ]]; then
+        echo "Error: PTE file not found: $pte_path"
+        exit 1
+    fi
+
+    echo "Running inference..."
+    echo "  PTE: $pte_path"
+
+    "$EXECUTOR_RUNNER" --model_path "$pte_path"
+}
+
+# Parse command line arguments
+case "$1" in
+    --build)
+        build_runtime
+        ;;
+    --run)
+        if [[ -z "$2" ]]; then
+            echo "Usage: $0 --run <pte_path>"
+            exit 1
+        fi
+        run_inference "$2"
+        ;;
+    --check-build)
+        check_build
+        ;;
+    *)
+        echo "Metal Backend Test Runner"
+        echo ""
+        echo "Usage:"
+        echo "  $0 --build              Build the Metal runtime"
+        echo "  $0 --run <pte>          Run inference with given model file"
+        echo "  $0 --check-build        Check if runtime is already built"
+        exit 1
+        ;;
+esac