Refactor pull workflow to avoid code duplication and add quantization export test

.ci/scripts/gather_test_models.py

@@ -10,6 +10,12 @@
 from typing import Any
 
 from examples.models import MODEL_NAME_TO_MODEL
+from examples.quantization.example import QUANT_MODEL_NAME_TO_MODEL
+
+BUILD_TOOLS = [
+    "buck2",
+    "cmake",
+]
 
 
 def set_output(name: str, val: Any) -> None:
@@ -31,7 +37,13 @@ def export_models_for_ci() -> None:
     """
     # This is the JSON syntax for configuration matrix used by GitHub
     # https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs
-    models = {"include": [{"model": name} for name in MODEL_NAME_TO_MODEL.keys()]}
+    models = {"include": []}
+    for name in MODEL_NAME_TO_MODEL.keys():
+        quantization = name in QUANT_MODEL_NAME_TO_MODEL
+        for build_tool in BUILD_TOOLS:
+            models["include"].append(
+                {"build-tool": build_tool, "model": name, "quantization": quantization}
+            )
     set_output("models", json.dumps(models))
 
 

.ci/scripts/setup-linux.sh

@@ -0,0 +1,24 @@
+#!/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.
+
+set -exu
+
+# shellcheck source=/dev/null
+source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
+
+BUILD_TOOL=$1
+if [[ -z "${BUILD_TOOL:-}" ]]; then
+  echo "Missing build tool (require buck2 or cmake), exiting..."
+  exit 1
+else
+  echo "Setup Linux for ${BUILD_TOOL} ..."
+fi
+
+# As Linux job is running inside a Docker container, all of its dependencies
+# have already been installed
+install_executorch
+build_executorch_runner "${BUILD_TOOL}"

.ci/scripts/setup-macos.sh

@@ -10,6 +10,14 @@ set -exu
 # shellcheck source=/dev/null
 source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
+BUILD_TOOL=$1
+if [[ -z "${BUILD_TOOL:-}" ]]; then
+  echo "Missing build tool (require buck2 or cmake), exiting..."
+  exit 1
+else
+  echo "Setup MacOS for ${BUILD_TOOL} ..."
+fi
+
 install_buck() {
   if ! command -v zstd &> /dev/null; then
     brew install zstd
@@ -36,6 +44,10 @@ install_buck() {
   fi
 }
 
-install_buck
+if [[ "${BUILD_TOOL}" == "buck2" ]]; then
+  install_buck
+fi
 install_conda
 install_pip_dependencies
+install_executorch
+build_executorch_runner "${BUILD_TOOL}"

.ci/scripts/test.sh

@@ -14,25 +14,46 @@ MODEL_NAME=$1
 if [[ -z "${MODEL_NAME:-}" ]]; then
   echo "Missing model name, exiting..."
   exit 1
-else
-  echo "Testing ${MODEL_NAME} ..."
+fi
+
+BUILD_TOOL=$2
+if [[ -z "${BUILD_TOOL:-}" ]]; then
+  echo "Missing build tool (require buck2 or cmake), exiting..."
+  exit 1
+fi
+
+QUANTIZATION=$3
+if [[ -z "${QUANTIZATION:-}" ]]; then
+  QUANTIZATION=false
 fi
 
 test_model() {
   python -m examples.export.export_example --model_name="${MODEL_NAME}"
 
   # Run test model
-  buck2 run //examples/executor_runner:executor_runner -- --model_path "./${MODEL_NAME}.pte"
+  if [[ "${BUILD_TOOL}" == "buck2" ]]; then
+    buck2 run //examples/executor_runner:executor_runner -- --model_path "./${MODEL_NAME}.pte"
+  elif [[ "${BUILD_TOOL}" == "cmake" ]]; then
+    CMAKE_OUTPUT_DIR=cmake-out
+    ./"${CMAKE_OUTPUT_DIR}"/executor_runner --model_path "./${MODEL_NAME}.pte"
+  else
+    echo "Invalid build tool ${BUILD_TOOL}. Only buck2 and cmake are supported atm"
+    exit 1
+  fi
 }
 
-build_and_test_executorch() {
-  # Build executorch runtime
-  buck2 build //examples/executor_runner:executor_runner
-
-  which python
-  # Test the select model
-  test_model
+test_quantized_model() {
+  python -m examples.quantization.example --model_name="${MODEL_NAME}"
 }
 
-install_executorch
-build_and_test_executorch
+which python
+
+echo "Testing ${MODEL_NAME} with ${BUILD_TOOL}..."
+# Test the select model
+test_model
+
+if [[ "${QUANTIZATION}" == true ]]; then
+  test_quantized_model
+else
+  echo "The model ${MODEL_NAME} doesn't support quantization yet"
+fi

.ci/scripts/utils.sh

@@ -38,3 +38,36 @@ install_pip_dependencies() {
     --index-url https://download.pytorch.org/whl/nightly/cpu
   popd || return
 }
+
+build_executorch_runner_buck2() {
+  # Build executorch runtime
+  buck2 build //examples/executor_runner:executor_runner
+}
+
+build_executorch_runner_cmake() {
+  CMAKE_OUTPUT_DIR=cmake-out
+  # Build executorch runtime using cmake
+  rm -rf "${CMAKE_OUTPUT_DIR}" && mkdir "${CMAKE_OUTPUT_DIR}"
+
+  pushd "${CMAKE_OUTPUT_DIR}" || return
+  cmake -DBUCK2=buck2 -DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}" ..
+  popd || return
+
+  if [ "$(uname)" == "Darwin" ]; then
+    CMAKE_JOBS=$(( $(sysctl -n hw.ncpu) - 1 ))
+  else
+    CMAKE_JOBS=$(( $(nproc) - 1 ))
+  fi
+  cmake --build "${CMAKE_OUTPUT_DIR}" -j "${CMAKE_JOBS}"
+}
+
+build_executorch_runner() {
+  if [[ $1 == "buck2" ]]; then
+    build_executorch_runner_buck2
+  elif [[ $1 == "cmake" ]]; then
+    build_executorch_runner_cmake
+  else
+    echo "Invalid build tool $1. Only buck2 and cmake are supported atm"
+    exit 1
+  fi
+}

.github/workflows/pull.yml

@@ -35,8 +35,8 @@ jobs:
 
           PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py
 
-  buck-build-test-linux:
-    name: buck-build-test-linux
+  test-models-linux:
+    name: test-models-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     needs: gather-models
     strategy:
@@ -53,15 +53,16 @@ jobs:
         # here, as it's there in the container
         export PATH="/opt/conda/envs/py_${PYTHON_VERSION}/bin:${PATH}"
 
-        # The name of model we are going to test
         MODEL_NAME=${{ matrix.model }}
+        BUILD_TOOL=${{ matrix.build-tool }}
+        QUANTIZATION=${{ matrix.quantization }}
+
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Build and test Executorch
-        bash .ci/scripts/test.sh "${MODEL_NAME}"
-        # Test custom ops
-        bash examples/custom_ops/test_custom_ops.sh buck2
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${QUANTIZATION}"
 
-  buck-build-test-macos:
-    name: buck-build-test-macos
+  test-models-macos:
+    name: test-models-macos
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
     needs: gather-models
     strategy:
@@ -73,25 +74,26 @@ jobs:
       ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
       script: |
         WORKSPACE=$(pwd)
-
         pushd "${WORKSPACE}/pytorch/executorch"
-        # Setup MacOS dependencies as there is no Docker support on MacOS atm
-        bash .ci/scripts/setup-macos.sh
 
-        # The name of model we are going to test
         MODEL_NAME=${{ matrix.model }}
+        BUILD_TOOL=${{ matrix.build-tool }}
+        QUANTIZATION=${{ matrix.quantization }}
+
+        # Setup MacOS dependencies as there is no Docker support on MacOS atm
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
         # Build and test Executorch
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}"
-        # Test custom ops
-        PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh buck2
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${QUANTIZATION}"
         popd
 
-  cmake-build-test-linux:
-    name: cmake-build-test-linux
+  test-custom-ops-linux:
+    name: test-custom-ops-linux
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
-    needs: gather-models
     strategy:
-      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      matrix:
+        include:
+          - build-tool: buck2
+          - build-tool: cmake
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -104,19 +106,19 @@ jobs:
         # here, as it's there in the container
         export PATH="/opt/conda/envs/py_${PYTHON_VERSION}/bin:${PATH}"
 
-        # The name of model we are going to test
-        MODEL_NAME=${{ matrix.model }}
-        # Build and test Executorch
-        bash .ci/scripts/test-cmake.sh "${MODEL_NAME}"
-        # Build and test custom ops
-        bash examples/custom_ops/test_custom_ops.sh cmake
+        BUILD_TOOL=${{ matrix.build-tool }}
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
+        # Test custom ops
+        PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh "${BUILD_TOOL}"
 
-  cmake-build-test-macos:
-    name: cmake-build-test-macos
+  test-custom-ops-macos:
+    name: test-custom-ops-macos
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
-    needs: gather-models
     strategy:
-      matrix: ${{ fromJSON(needs.gather-models.outputs.models) }}
+      matrix:
+        include:
+          - build-tool: buck2
+          - build-tool: cmake
       fail-fast: false
     with:
       runner: macos-m1-12
@@ -126,13 +128,10 @@ jobs:
         WORKSPACE=$(pwd)
 
         pushd "${WORKSPACE}/pytorch/executorch"
-        # Setup MacOS dependencies as there is no Docker support on MacOS atm
-        bash .ci/scripts/setup-macos.sh
 
-        # The name of model we are going to test
-        MODEL_NAME=${{ matrix.model }}
-        # Build and test Executorch
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test-cmake.sh "${MODEL_NAME}"
+        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}"
         # Build and test custom ops
-        PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh cmake
+        PYTHON_EXECUTABLE=python bash examples/custom_ops/test_custom_ops.sh "${BUILD_TOOL}"
         popd

examples/quantization/example.py

@@ -30,6 +30,11 @@
 
 from ..models import MODEL_NAME_TO_MODEL
 
+# Note: for mv3, the mul op is not supported in XNNPACKQuantizer, that could be supported soon
+QUANT_MODEL_NAME_TO_MODEL = {
+    name: MODEL_NAME_TO_MODEL[name] for name in ["linear", "add", "add_mul", "mv2"]
+}
+
 
 def quantize(model_name, model, example_inputs):
     """This is the official recommended flow for quantization in pytorch 2.0 export"""
@@ -94,11 +99,6 @@ def verify_xnnpack_quantizer_matching_fx_quant_model(model_name, model, example_
 
 
 if __name__ == "__main__":
-    # Note: for mv3, the mul op is not supported in XNNPACKQuantizer, that could be supported soon
-    QUANT_MODEL_NAME_TO_MODEL = {
-        name: MODEL_NAME_TO_MODEL[name] for name in ["linear", "add", "add_mul", "mv2"]
-    }
-
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "-m",