Update on "[executorch] Add vectorized scalar path for single-element Tensor passed to optimized mul"

We are currently doing slow broadcasting for this case. After this diff, we should get nice vectorization.

Differential Revision: [D61560825](https://our.internmc.facebook.com/intern/diff/D61560825/)

[ghstack-poisoned]

Author: swolchok

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-b556d31586845fb1e296a975d2b85d9d325205c9
+c42ac54d9e817bf0a0366eb78e6c8beba4d5eff5

.ci/docker/common/install_linter.sh

@@ -13,3 +13,7 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 # NB: Install all linter dependencies, the caching of lintrunner init could be
 # done after Executorch becomes public
 pip_install -r requirements-lintrunner.txt
+
+# Install google-java-format
+curl -L --retry 3 https://github.com/google/google-java-format/releases/download/v1.23.0/google-java-format_linux-x86-64 > /opt/google-java-format
+chmod +x /opt/google-java-format

.ci/scripts/build-qnn-sdk.sh

@@ -5,15 +5,43 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-set -ex
+set -eux
 
 build_qnn_backend() {
   echo "Start building qnn backend."
   export ANDROID_NDK_ROOT=/opt/ndk
   export QNN_SDK_ROOT=/tmp/qnn/2.23.0.240531
-  export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
+  export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)"
 
   bash backends/qualcomm/scripts/build.sh --skip_aarch64 --job_number 2 --release
 }
 
+set_up_aot() {
+  cd $EXECUTORCH_ROOT
+  if [ ! -d "cmake-out" ]; then
+      mkdir cmake-out
+  fi
+  pushd cmake-out
+  cmake .. \
+      -DCMAKE_INSTALL_PREFIX=$PWD \
+      -DEXECUTORCH_BUILD_QNN=ON \
+      -DQNN_SDK_ROOT=${QNN_SDK_ROOT} \
+      -DEXECUTORCH_BUILD_SDK=ON \
+      -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+      -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
+      -DPYTHON_EXECUTABLE=python3 \
+      -DEXECUTORCH_SEPARATE_FLATCC_HOST_PROJECT=OFF
+  cmake --build $PWD --target "PyQnnManagerAdaptor" "PyQnnWrapperAdaptor" -j$(nproc)
+  # install Python APIs to correct import path
+  # The filename might vary depending on your Python and host version.
+  cp -f backends/qualcomm/PyQnnManagerAdaptor.cpython-310-x86_64-linux-gnu.so $EXECUTORCH_ROOT/backends/qualcomm/python
+  cp -f backends/qualcomm/PyQnnWrapperAdaptor.cpython-310-x86_64-linux-gnu.so $EXECUTORCH_ROOT/backends/qualcomm/python
+  popd
+
+  # Workaround for fbs files in exir/_serialize
+  cp schema/program.fbs exir/_serialize/program.fbs
+  cp schema/scalar_type.fbs exir/_serialize/scalar_type.fbs
+}
+
 build_qnn_backend
+set_up_aot

.ci/scripts/setup-linux.sh

@@ -20,6 +20,5 @@ fi
 
 # As Linux job is running inside a Docker container, all of its dependencies
 # have already been installed
-install_flatc_from_source
 install_executorch
 build_executorch_runner "${BUILD_TOOL}"

.ci/scripts/setup-macos.sh

@@ -128,7 +128,5 @@ if [[ -z "${GITHUB_RUNNER:-}" ]]; then
 fi
 
 print_cmake_info
-install_pytorch_and_domains
-install_flatc_from_source
 install_executorch
 build_executorch_runner "${BUILD_TOOL}"

.ci/scripts/test.sh

@@ -28,9 +28,25 @@ if [[ -z "${BACKEND:-}" ]]; then
   exit 1
 fi
 
+UPLOAD_DIR=${4:-}
+
+if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then
+  PYTHON_EXECUTABLE=python3
+fi
 which "${PYTHON_EXECUTABLE}"
+
 # Just set this variable here, it's cheap even if we use buck2
 CMAKE_OUTPUT_DIR=cmake-out
+EXPORTED_MODEL=${MODEL_NAME}
+
+prepare_artifacts_upload() {
+  if [ -n "$UPLOAD_DIR" ]; then
+    echo "Preparing for uploading generated artifacs"
+    zip -j model.zip "${EXPORTED_MODEL}"
+    mkdir -p "${UPLOAD_DIR}"
+    mv model.zip "${UPLOAD_DIR}"
+  fi
+}
 
 build_cmake_executor_runner() {
   echo "Building executor_runner"
@@ -114,6 +130,7 @@ test_model_with_xnnpack() {
   fi
 
   OUTPUT_MODEL_PATH="${MODEL_NAME}_xnnpack_${SUFFIX}.pte"
+  EXPORTED_MODEL=${OUTPUT_MODEL_PATH}
 
   # Run test model
   if [[ "${BUILD_TOOL}" == "buck2" ]]; then
@@ -129,9 +146,36 @@ test_model_with_xnnpack() {
   fi
 }
 
+test_model_with_qnn() {
+  source "$(dirname "${BASH_SOURCE[0]}")/build-qnn-sdk.sh"
+  echo "ANDROID_NDK_ROOT: $ANDROID_NDK_ROOT"
+  echo "QNN_SDK_ROOT: $QNN_SDK_ROOT"
+  echo "EXECUTORCH_ROOT: $EXECUTORCH_ROOT"
+
+  export LD_LIBRARY_PATH=$QNN_SDK_ROOT/lib/x86_64-linux-clang/
+  export PYTHONPATH=$EXECUTORCH_ROOT/..
+
+  if [[ "${MODEL_NAME}" == "dl3" ]]; then
+    "${PYTHON_EXECUTABLE}" -m examples.qualcomm.scripts.deeplab_v3 -b ${CMAKE_OUTPUT_DIR} -m SM8550 --compile_only --download
+    EXPORTED_MODEL=./deeplab_v3/dlv3_qnn.pte
+  fi
+}
+
 if [[ "${BACKEND}" == "portable" ]]; then
   echo "Testing ${MODEL_NAME} with portable kernels..."
   test_model
+elif [[ "${BACKEND}" == "qnn" ]]; then
+  echo "Testing ${MODEL_NAME} with qnn..."
+  test_model_with_qnn
+  if [[ $? -eq 0 ]]; then
+    prepare_artifacts_upload
+  fi
+elif [[ "${BACKEND}" == "xnnpack" ]]; then
+  echo "Testing ${MODEL_NAME} with xnnpack..."
+  test_model_with_xnnpack true true
+  if [[ $? -eq 0 ]]; then
+    prepare_artifacts_upload
+  fi
 else
   set +e
   if [[ "${BACKEND}" == *"quantization"* ]]; then
@@ -153,5 +197,7 @@ else
   if [[ -n "${Q_ERROR:-}" ]] || [[ -n "${D_ERROR:-}" ]] || [[ -n "${Q_D_ERROR:-}" ]]; then
     echo "Portable q8 ${Q_ERROR:-ok}," "Delegation fp32 ${D_ERROR:-ok}," "Delegation q8 ${Q_D_ERROR:-ok}"
     exit 1
+  else
+    prepare_artifacts_upload
   fi
 fi

.ci/scripts/test_llama.sh

@@ -9,7 +9,7 @@ set -exu
 # shellcheck source=/dev/null
 source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
-MODEL_NAME=$1 # stories110M.pt
+MODEL_NAME=$1 # stories110M
 BUILD_TOOL=$2 # buck2 or cmake
 DTYPE=$3 # fp16 or fp32
 MODE=${4:-"xnnpack+custom"} # portable or xnnpack+custom or xnnpack+custom+qe
@@ -140,7 +140,7 @@ cmake_build_llama_runner() {
 
 cleanup_files() {
   echo "Deleting downloaded and generated files"
-  rm "${MODEL_NAME}"
+  rm "${CHECKPOINT_FILE_NAME}"
   rm tokenizer.model
   rm tokenizer.bin
   rm "${EXPORTED_MODEL_NAME}"
@@ -159,8 +159,10 @@ prepare_artifacts_upload() {
 
 # Download and create artifacts.
 PARAMS="params.json"
+CHECKPOINT_FILE_NAME=""
 touch "${PARAMS}"
-if [[ "${MODEL_NAME}" == "stories110M.pt" ]]; then
+if [[ "${MODEL_NAME}" == "stories110M" ]]; then
+  CHECKPOINT_FILE_NAME="stories110M.pt"
   download_stories_model_artifacts
 else
   echo "Unsupported model name ${MODEL_NAME}"
@@ -181,7 +183,7 @@ fi
 # Export model.
 EXPORTED_MODEL_NAME="${EXPORTED_MODEL_NAME}.pte"
 echo "Exporting ${EXPORTED_MODEL_NAME}"
-EXPORT_ARGS="-c stories110M.pt -p ${PARAMS} -d ${DTYPE} -n ${EXPORTED_MODEL_NAME} -kv"
+EXPORT_ARGS="-c ${CHECKPOINT_FILE_NAME} -p ${PARAMS} -d ${DTYPE} -n ${EXPORTED_MODEL_NAME} -kv"
 if [[ "${XNNPACK}" == "ON" ]]; then
   EXPORT_ARGS="${EXPORT_ARGS} -X -qmode 8da4w -G 128"
 fi

.ci/scripts/utils.sh

@@ -19,11 +19,9 @@ retry () {
 install_executorch() {
   which pip
   # Install executorch, this assumes that Executorch is checked out in the
-  # current directory. The --extra-index-url options tell pip to look on the
-  # pytorch servers for nightly and pre-release versions of torch packages.
-  pip install . --no-build-isolation -v \
-      --extra-index-url https://download.pytorch.org/whl/test/cpu \
-      --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+  # current directory.
+  # TODO(T199538337): clean up install scripts to use install_requirements.sh
+  ./install_requirements.sh --pybind xnnpack
   # Just print out the list of packages for debugging
   pip list
 }
@@ -35,42 +33,6 @@ install_pip_dependencies() {
   popd || return
 }
 
-install_domains() {
-  echo "Install torchvision and torchaudio"
-  pip install --no-use-pep517 --user "git+https://github.com/pytorch/audio.git@${TORCHAUDIO_VERSION}"
-  pip install --no-use-pep517 --user "git+https://github.com/pytorch/vision.git@${TORCHVISION_VERSION}"
-}
-
-install_pytorch_and_domains() {
-  pushd .ci/docker || return
-  TORCH_VERSION=$(cat ci_commit_pins/pytorch.txt)
-  popd || return
-
-  git clone https://github.com/pytorch/pytorch.git
-
-  # Fetch the target commit
-  pushd pytorch || return
-  git checkout "${TORCH_VERSION}"
-  git submodule update --init --recursive
-
-  export _GLIBCXX_USE_CXX11_ABI=0
-  # Then build and install PyTorch
-  python setup.py bdist_wheel
-  pip install "$(echo dist/*.whl)"
-
-  # Grab the pinned audio and vision commits from PyTorch
-  TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
-  export TORCHAUDIO_VERSION
-  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
-  export TORCHVISION_VERSION
-
-  install_domains
-
-  popd || return
-  # Print sccache stats for debugging
-  sccache --show-stats || true
-}
-
 install_flatc_from_source() {
   # NB: This function could be used to install flatbuffer from source
   pushd third-party/flatbuffers || return

.github/workflows/android-perf.yml

@@ -135,24 +135,39 @@ jobs:
       fail-fast: false
     with:
       runner: linux.2xlarge
-      docker-image: executorch-ubuntu-22.04-clang12
+      docker-image: executorch-ubuntu-22.04-clang12-android
       submodules: 'true'
       timeout: 60
       upload-artifact: android-models
       script: |
         # The generic Linux job chooses to use base env, not the one setup by the image
+        echo "::group::Setting up dev environment"
         CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
         conda activate "${CONDA_ENV}"
-
+        if [[ ${{ matrix.delegate }} == "qnn" ]]; then
+            PYTHON_EXECUTABLE=python bash .ci/scripts/setup-qnn-deps.sh
+            PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh
+        fi
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "cmake"
-        echo "Exporting model: ${{ matrix.model }}"
-        export ARTIFACTS_DIR_NAME=artifacts-to-be-uploaded/${{ matrix.model }}_${{ matrix.delegate }}
+        ARTIFACTS_DIR_NAME=artifacts-to-be-uploaded/${{ matrix.model }}_${{ matrix.delegate }}
+        echo "::endgroup::"
 
-        # TODO(T197546696): Note that the following scripts/steps only work for llama. It's expected to fail for other models+delegates.
-        # Install requirements for export_llama
-        PYTHON_EXECUTABLE=python bash examples/models/llama2/install_requirements.sh
-        # Test llama2
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh "${{ matrix.model }}.pt" "cmake" "fp32" "xnnpack+custom+qe" "${ARTIFACTS_DIR_NAME}"\
+        echo "::group::Exporting ${{ matrix.delegate }} model: ${{ matrix.model }}"
+        BUILD_MODE="cmake"
+        DTYPE="fp32"
+
+        if [[ ${{ matrix.model }} =~ ^stories* ]]; then
+            # Install requirements for export_llama
+            PYTHON_EXECUTABLE=python bash examples/models/llama2/install_requirements.sh
+            # Test llama2
+            if [[ ${{ matrix.delegate }} == "xnnpack" ]]; then
+                DELEGATE_CONFIG="xnnpack+custom+qe"
+            fi
+            PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh "${{ matrix.model }}" "${BUILD_MODE}" "${DTYPE}" "${DELEGATE_CONFIG}" "${ARTIFACTS_DIR_NAME}"
+        else
+            PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh "${{ matrix.model }}" "${BUILD_MODE}" "${{ matrix.delegate }}" "${ARTIFACTS_DIR_NAME}"
+        fi
+        echo "::endgroup::"
 
   # Upload models to S3. The artifacts are needed not only by the device farm but also TorchChat
   upload-models:

.github/workflows/lint.yml

@@ -54,3 +54,20 @@ jobs:
           lint.json || true
 
         exit $RC
+
+  android-java-format:
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-linter
+      fetch-depth: 0
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        FILES_NEEDS_FORMAT=$(/opt/google-java-format -n extension/android/src/main/java/org/pytorch/executorch/*.java \
+          examples/demo-apps/android/ExecuTorchDemo/app/src/main/java/com/example/executorchdemo/*.java \
+          examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/*.java)
+        if [ -n "$FILES_NEEDS_FORMAT" ]; then
+          echo "Warning: The following files need formatting. Please use google-java-format."
+          echo "$FILES_NEEDS_FORMAT"
+          exit 1
+        fi

.github/workflows/pull.yml

@@ -112,7 +112,7 @@ jobs:
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python bash examples/models/llama2/install_requirements.sh
         # Test llama2
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M.pt "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
 
   test-llama-runner-linux-android:
     name: test-llama-runner-linux-android
@@ -406,4 +406,4 @@ jobs:
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python bash examples/models/llama2/install_requirements.sh
         # Test llama2
-        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M.pt "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh stories110M "${BUILD_TOOL}" "${DTYPE}" "${MODE}"

.github/workflows/trunk.yml

@@ -9,6 +9,7 @@ on:
       - ciflow/trunk/*
   pull_request:
     paths:
+      - .ci/docker/ci_commit_pins/pytorch.txt
       - .ci/scripts/**
   workflow_dispatch:
 
@@ -142,7 +143,6 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         source .ci/scripts/utils.sh
-        install_flatc_from_source
         install_executorch
 
         install_arm
@@ -168,7 +168,6 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         source .ci/scripts/utils.sh
-        install_flatc_from_source
         install_executorch
 
         install_arm
@@ -269,4 +268,27 @@ jobs:
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash examples/models/llama2/install_requirements.sh
         # Test llama2
-        PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_llama.sh stories110M.pt "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
+        PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_llama.sh stories110M "${BUILD_TOOL}" "${DTYPE}" "${MODE}"
+
+  test-qnn-model:
+    name: test-qnn-model
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    strategy:
+      matrix:
+        dtype: [fp32]
+        model: [dl3]
+      fail-fast: false
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-clang12-android
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 900
+      script: |
+        # The generic Linux job chooses to use base env, not the one setup by the image
+        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+        conda activate "${CONDA_ENV}"
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh cmake
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-qnn-deps.sh
+        PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh
+        PYTHON_EXECUTABLE=python bash .ci/scripts/test.sh ${{ matrix.model }} "cmake" "qnn"

.github/workflows/update-viablestrict.yml

@@ -20,6 +20,6 @@ jobs:
         with:
           repository: pytorch/executorch
           stable-branch: viable/strict
-          requires: '[\"pull\", \"lint\", \"trunk\", \"Build documentation\", \"Android\", \"Apple\"]'
+          requires: '[\"pull\", \"lint\", \"trunk\", \"Build documentation\", \"^Android$\", \"^Apple$\"]'
           secret-bot-token: ${{ secrets.UPDATEBOT_TOKEN }}
           rockset-api-key: ${{ secrets.ROCKSET_API_KEY }}