Add XNNPACK with and without quantization tests to CI (#123)

Summary:
TSIA. After #122, we can now add XNNPACK tests to CI.  Note that `xnn_executor_runner` is only available via buck2 here, its cmake support will come in the next PR.

Pull Request resolved: #123

Test Plan:
Spot check for `add` example which supports both XNNPACK and quantization.  Both are tested:

Linux: https://github.com/pytorch/executorch/actions/runs/5970562477/job/16198874573?pr=123
MacOS: https://github.com/pytorch/executorch/actions/runs/5970562477/job/16198877840?pr=123

Reviewed By: guangy10, kirklandsign

Differential Revision: D48670404

Pulled By: huydhn

fbshipit-source-id: 51974887cce141d9abccf3581f5003dfb23fe53d

Author: huydhn

.ci/scripts/gather_test_models.py

@@ -41,9 +41,18 @@ def export_models_for_ci() -> None:
         quantization = (
             name in MODEL_NAME_TO_OPTIONS and MODEL_NAME_TO_OPTIONS[name].quantization
         )
+        xnnpack_delegation = (
+            name in MODEL_NAME_TO_OPTIONS
+            and MODEL_NAME_TO_OPTIONS[name].xnnpack_delegation
+        )
         for build_tool in BUILD_TOOLS:
             models["include"].append(
-                {"build-tool": build_tool, "model": name, "quantization": quantization}
+                {
+                    "build-tool": build_tool,
+                    "model": name,
+                    "quantization": quantization,
+                    "xnnpack_delegation": xnnpack_delegation,
+                }
             )
     set_output("models", json.dumps(models))
 

.ci/scripts/test.sh

@@ -27,37 +27,71 @@ if [[ -z "${QUANTIZATION:-}" ]]; then
   QUANTIZATION=false
 fi
 
+XNNPACK_DELEGATION=$4
+if [[ -z "${XNNPACK_DELEGATION:-}" ]]; then
+  XNNPACK_DELEGATION=false
+fi
+
+which "${PYTHON_EXECUTABLE}"
+# Just set this variable here, it's cheap even if we use buck2
+CMAKE_OUTPUT_DIR=cmake-out
+
 test_model() {
-  python -m examples.export.export_example --model_name="${MODEL_NAME}"
+  "${PYTHON_EXECUTABLE}" -m examples.export.export_example --model_name="${MODEL_NAME}"
 
   # Run test model
   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
 }
 
-which python
-
-echo "Testing ${MODEL_NAME} with ${BUILD_TOOL}..."
-# Test the select model
-test_model
+test_model_with_xnnpack() {
+  WITH_QUANTIZATION=$1
+  if [[ ${WITH_QUANTIZATION} == true ]]; then
+    SUFFIX="q8"
+    "${PYTHON_EXECUTABLE}" -m examples.backend.xnnpack_examples --model_name="${MODEL_NAME}" --delegate --quantize
+  else
+    SUFFIX="fp32"
+    "${PYTHON_EXECUTABLE}" -m examples.backend.xnnpack_examples --model_name="${MODEL_NAME}" --delegate
+  fi
 
-if [[ "${QUANTIZATION}" == true ]]; then
+  OUTPUT_MODEL_PATH="${MODEL_NAME}_xnnpack_${SUFFIX}.pte"
+  # Run test model
   if [[ "${BUILD_TOOL}" == "buck2" ]]; then
-    bash examples/quantization/test_quantize.sh buck2 "${MODEL_NAME}"
+    buck2 run //examples/backend:xnn_executor_runner -- --model_path "${OUTPUT_MODEL_PATH}"
   elif [[ "${BUILD_TOOL}" == "cmake" ]]; then
-    CMAKE_OUTPUT_DIR=cmake-out
-    bash examples/quantization/test_quantize.sh cmake "${MODEL_NAME}"
+    # TODO: Add cmake support for xnn_executor_runner
+    echo "XNNPACK doesn't support cmake yet, skipping..."
   else
     echo "Invalid build tool ${BUILD_TOOL}. Only buck2 and cmake are supported atm"
     exit 1
   fi
+}
+
+echo "Testing ${MODEL_NAME} (fp32, quantized, xnnpack) with ${BUILD_TOOL}..."
+# Test the select model without XNNPACK or quantization
+test_model
+
+# Test quantization
+if [[ "${QUANTIZATION}" == true ]]; then
+  bash examples/quantization/test_quantize.sh "${BUILD_TOOL}" "${MODEL_NAME}"
 else
   echo "The model ${MODEL_NAME} doesn't support quantization yet"
 fi
+
+# Test XNNPACK without quantization
+if [[ "${XNNPACK_DELEGATION}" == true ]]; then
+  test_model_with_xnnpack false
+else
+  echo "The model ${MODEL_NAME} doesn't support XNNPACK yet"
+fi
+
+# Test XNNPACK with quantization
+if [[ "${XNNPACK_DELEGATION}" == true ]] && [[ "${QUANTIZATION}" == true ]]; then
+  test_model_with_xnnpack true
+fi

.github/workflows/pull.yml

@@ -56,10 +56,11 @@ jobs:
         MODEL_NAME=${{ matrix.model }}
         BUILD_TOOL=${{ matrix.build-tool }}
         QUANTIZATION=${{ matrix.quantization }}
+        XNNPACK_DELEGATION=${{ matrix.xnnpack_delegation }}
 
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
         # Build and test Executorch
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${QUANTIZATION}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${QUANTIZATION}" "${XNNPACK_DELEGATION}"
 
   test-models-macos:
     name: test-models-macos
@@ -79,11 +80,12 @@ jobs:
         MODEL_NAME=${{ matrix.model }}
         BUILD_TOOL=${{ matrix.build-tool }}
         QUANTIZATION=${{ matrix.quantization }}
+        XNNPACK_DELEGATION=${{ matrix.xnnpack_delegation }}
 
         # 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}" "${BUILD_TOOL}" "${QUANTIZATION}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${QUANTIZATION}" "${XNNPACK_DELEGATION}"
         popd
 
   test-custom-ops-linux:

exir/backend/utils.py

@@ -30,15 +30,15 @@ def is_same_node(
             and all(
                 is_same_node(arg_left, arg_right)
                 for arg_left, arg_right in zip(
-                    node_left.all_input_nodes, node_right.all_input_nodes, strict=True
+                    node_left.all_input_nodes, node_right.all_input_nodes
                 )
             )
         ):
             return False
     else:
         if len(list(node_left)) != len(list(node_right)):
             return False
-        for n_left, n_right in zip(node_left, node_right, strict=True):
+        for n_left, n_right in zip(node_left, node_right):
             # pyre-fixme[6]: For 1st argument expected `Iterable[Node]` but got `Node`.
             # pyre-fixme[6]: For 2nd argument expected `Iterable[Node]` but got `Node`.
             if not is_same_node(n_left, n_right):
@@ -55,9 +55,7 @@ def is_identical_graph(
     # is not the same.
     if len(list(graph_left.graph.nodes)) != len(list(graph_right.graph.nodes)):
         return False
-    for node_left, node_right in zip(
-        graph_left.graph.nodes, graph_right.graph.nodes, strict=True
-    ):
+    for node_left, node_right in zip(graph_left.graph.nodes, graph_right.graph.nodes):
         if not (is_same_node(node_left, node_right)):
             return False
     return True
@@ -90,7 +88,6 @@ def replace_quantized_partition_with_op(
     partition: SourcePartition,
     replacement_op: torch._ops.OpOverloadPacket,
 ) -> Tuple[torch.fx.Node, List[torch.fx.Node], List[torch.fx.Node]]:
-
     """
     Replaces partition with the op specified by replacement_op. It's also expected that
     the nodes contained in partition are sourced from a quantized module as this function