Enable android runner with custom sdpa op (#363)

kimishpatel · malfet · commit 1f7f01ec0b3c · 2024-07-17T09:55:42.000-07:00
Summary:
This diff
- refactors install_et.sh into a bunch of utils
- Uses those utils in build_android.sh to minimize duplication.
- Makes sure taht we are building with custom sdpa op

Test Plan:

Model export
python export.py --quant '{"linear:a8w4dq" : {"groupsize": 256}}'
--checkpoint-path /home/kimishpatel/models/llama2/stories/stories110M.pt
--params-path /home/kimishpatel/models/llama2/stories/params.json
--output-pte-path /tmp/stories110m_a8w4dq.pte
python utils/tokenizer.py --tokenizer-model=/tmp/tokenizer.model

linux:
./scripts/install_et.sh
rm -rf build/cmake-out/
cmake -S ./runner-et -B build/cmake-out -G Ninja
cmake --build ./build/cmake-out

./build/cmake-out/runner_et /tmp/stories110m_a8w4dq.pte -z /tmp/tokenizer.bin -t 0 -n 120

android:
./runner-et/build_android.sh
adb push ./build/cmake-out-android/runner_et /data/local/tmp/
adb push /tmp/stories110m_a8w4dq.pte /data/local/tmp/
adb push /tmp/tokenizer.bin /data/local/tmp/
adb shell "cd /data/local/tmp &amp;&amp; ./runner_et ./stories110m_a8w4dq.pte -z
./tokenizer.bin -t 0 -n 120"

Will add build commands to ci in the next PR

Reviewers:

Subscribers:

Tasks:

Tags:
diff --git a/export_et_util.py b/export_et_util.py
@@ -80,6 +80,7 @@ def forward(self, x, freqs_cis, mask, input_pos=None):
 
 
 def replace_attention_with_custom_sdpa_attention(module: nn.Module):
+    from executorch.examples.models.llama2.custom_ops import sdpa_with_kv_cache  # noqa
     for name, child in module.named_children():
         if isinstance(child, Attention):
             setattr(module, name, CustomSDPAAttention(child))
diff --git a/runner-et/CMakeLists.txt b/runner-et/CMakeLists.txt
@@ -41,7 +41,7 @@ target_link_libraries(
   run PRIVATE
         executorch
         extension_module
-        ${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/lib/libextension_data_loader.a # This one gets installed in build directory by ExecuTorch
+        ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/${CMAKE_OUT_DIR}/extension/data_loader/libextension_data_loader.a # This one does not get installed by ExecuTorch
         optimized_kernels
         portable_kernels
         cpublas
@@ -58,5 +58,32 @@ target_link_options_shared_lib(xnnpack_backend)
 target_link_options_shared_lib(XNNPACK)
 target_link_options_shared_lib(pthreadpool)
 target_link_options_shared_lib(cpuinfo)
-target_link_options_shared_lib(executorch)
-target_link_libraries(run PRIVATE "$<LINK_LIBRARY:WHOLE_ARCHIVE,${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops.a>")
+# Not clear why linking executorch as whole-archive outside android/apple is leading
+# to double registration. Most likely because of linkage issues.
+# Will figure this out later. Until then use this.
+if(ANDROID OR APPLE)
+  target_link_options_shared_lib(executorch)
+endif()
+target_link_libraries(runner_et PRIVATE
+"$<LINK_LIBRARY:WHOLE_ARCHIVE,${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops.a>")
+
+# This one is needed for cpuinfo where it uses android specific log lib
+if(ANDROID)
+  target_link_libraries(runner_et PRIVATE log)
+endif()
+
+# Adding target_link_options_shared_lib as commented out below leads to this:
+#
+# CMake Error at Utils.cmake:22 (target_link_options):
+#   Cannot specify link options for target
+#   "/Users/scroy/etorch/torchchat/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops_lib.a"
+#   which is not built by this project.
+# Call Stack (most recent call first):
+#   Utils.cmake:30 (macos_kernel_link_options)
+#   CMakeLists.txt:41 (target_link_options_shared_lib)
+#
+#target_link_options_shared_lib("${TORCHCHAT_ROOT}/et-build/src/executorch/${CMAKE_OUT_DIR}/examples/models/llama2/custom_ops/libcustom_ops_lib.a") # This one does not get installed by ExecuTorch
+
+# This works on mac, but appears to run into issues on linux
+# It is needed to solve:
+# E 00:00:00.055965 executorch:method.cpp:536] Missing operator: [8] llama::sdpa_with_kv_cache.out
diff --git a/runner-et/build_android.sh b/runner-et/build_android.sh
@@ -7,6 +7,8 @@
 
 set -ex
 
+source "$(dirname "${BASH_SOURCE[0]}")/../scripts/install_utils.sh"
+
 if ["${ANDROID_NDK}" == ""]; then
   echo "Please set ANDROID_NDK enviornment variable."
   echo "For example it can be /Users/guest/Desktop/android-ndk-r26."
@@ -17,45 +19,24 @@ else
 fi
 
 export CMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake
-export DANDROID_ABI=arm64-v8a
-export DANDROID_PLATFORM=android-23 
+export ANDROID_ABI=arm64-v8a
+export ANDROID_PLATFORM=android-23
 export ET_BUILD_DIR="et-build-android"
 export CMAKE_OUT_DIR="cmake-out-android"
+export EXECUTORCH_BUILD_CUSTOM_OPS_AOT="OFF"
+export EXECUTORCH_BUILD_CUSTOM="ON"
+export CMAKE_OUT_DIR="cmake-out-android"
 # export DCMAKE_INSTALL_PREFIX=cmake-out-android
 #
 
-install_executorch() {
-  echo "Cloning executorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
-  ET_BUILD_DIR="${TORCHCHAT_ROOT}/${ET_BUILD_DIR}"
-  rm -rf ${ET_BUILD_DIR}
-  mkdir -p ${ET_BUILD_DIR}/src
-  pushd ${ET_BUILD_DIR}/src
-  git clone https://github.com/pytorch/executorch.git
-  cd executorch
-  git checkout viable/strict
-  echo "Install executorch: submodule update"
-  git submodule sync
-  git submodule update --init
-
-  echo "Applying fixes"
-  cp ${TORCHCHAT_ROOT}/scripts/fixes_et/module.cpp ${ET_BUILD_DIR}/src/executorch/extension/module/module.cpp # ET uses non-standard C++ that does not compile in GCC
-  cp ${TORCHCHAT_ROOT}/scripts/fixes_et/managed_tensor.h ${ET_BUILD_DIR}/src/executorch/extension/runner_util/managed_tensor.h # ET is missing headers for vector/memory.  This causes downstream issues when building runner-et.
-
-  CMAKE_OUT_DIR="cmake-out-android"
-  echo "Building and installing C++ libraries"
-  echo "Inside: ${PWD}"
-  mkdir ${CMAKE_OUT_DIR}
-  cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-23 -DCMAKE_INSTALL_PREFIX=cmake-out-android -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON -S . -B ${CMAKE_OUT_DIR} -G Ninja
-  cmake --build ${CMAKE_OUT_DIR}
-  cmake --install ${CMAKE_OUT_DIR} --prefix ${ET_BUILD_DIR}/install
-  popd
-}
-
 build_runner_et() {
   rm -rf build/cmake-out-android
+  echo "ET BUILD DIR IS ${ET_BUILD_DIR}"
   cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-23 -S ./runner-et -B build/cmake-out-android -G Ninja
   cmake --build build/cmake-out-android/ -j16 --config Release
 }
 
-# install_executorch
+find_cmake_prefix_path
+clone_executorch
+install_executorch
 build_runner_et
diff --git a/scripts/install_et.sh b/scripts/install_et.sh
@@ -5,55 +5,20 @@
 # 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
+set -ex
 
-install_pip_dependencies() {
-  echo "Intalling common pip packages"
-
-  pip3 install wheel
-  pip3 install "cmake>=3.19"
-  pip3 install ninja
-  pip3 install zstd
-  pushd ${TORCHCHAT_ROOT}
-  pip3 install -r ./requirements.txt
-  popd
-}
-
-install_executorch() {
-  echo "Cloning executorch to ${TORCHCHAT_ROOT}/et-build/src"
-  rm -rf ${TORCHCHAT_ROOT}/et-build
-  mkdir -p ${TORCHCHAT_ROOT}/et-build/src
-  pushd ${TORCHCHAT_ROOT}/et-build/src
-  git clone https://github.com/pytorch/executorch.git
-  cd executorch
-  git checkout viable/strict
-  echo "Install executorch: submodule update"
-  git submodule sync
-  git submodule update --init
-
-  echo "Building and installing python libraries"
-  if [ "${ENABLE_ET_PYBIND}" = false ]; then
-      echo "Not installing pybind"
-      bash ./install_requirements.sh
-  else
-      echo "Installing pybind"
-      bash ./install_requirements.sh --pybind xnnpack
-  fi
-  pip3 list
-
-  echo "Building and installing C++ libraries"
-  echo "Inside: ${PWD}"
-  mkdir cmake-out
-  cmake -DCMAKE_BUILD_TYPE=Release -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_CUSTOM=ON -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON -S . -B cmake-out -G Ninja
-  cmake --build cmake-out
-  cmake --install cmake-out --prefix ${TORCHCHAT_ROOT}/et-build/install
-  popd
-}
+source "$(dirname "${BASH_SOURCE[0]}")/install_utils.sh"
 
+if [ "${ET_BUILD_DIR}" == "" ]; then
+  ET_BUILD_DIR="et-build"
+fi
 
 ENABLE_ET_PYBIND="${1:-true}"
 
 pushd ${TORCHCHAT_ROOT}
+find_cmake_prefix_path
 install_pip_dependencies
-install_executorch $ENABLE_ET_PYBIND
+clone_executorch
+install_executorch_python_libs $ENABLE_ET_PYBIND
+install_executorch
 popd
diff --git a/scripts/install_utils.sh b/scripts/install_utils.sh
@@ -0,0 +1,107 @@
+#!/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 -ex
+
+install_pip_dependencies() {
+  echo "Intalling common pip packages"
+
+  pip3 install wheel
+  pip3 install "cmake>=3.19"
+  pip3 install ninja
+  pip3 install zstd
+  pushd ${TORCHCHAT_ROOT}
+  pip3 install -r ./requirements.txt
+  popd
+}
+
+function find_cmake_prefix_path() {
+  path=`python -c "from distutils.sysconfig import get_python_lib;print(get_python_lib())"`
+  MY_CMAKE_PREFIX_PATH=$path
+}
+
+clone_executorch() {
+  echo "Cloning executorch to ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src"
+  rm -rf ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}
+  mkdir -p ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
+  pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
+  git clone https://github.com/pytorch/executorch.git
+  cd executorch
+  git checkout viable/strict
+  echo "Install executorch: submodule update"
+  git submodule sync
+  git submodule update --init
+
+  echo "Applying fixes"
+  cp ${TORCHCHAT_ROOT}/scripts/fixes_et/module.cpp ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/extension/module/module.cpp # ET uses non-standard C++ that does not compile in GCC
+  cp ${TORCHCHAT_ROOT}/scripts/fixes_et/managed_tensor.h ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src/executorch/extension/runner_util/managed_tensor.h # ET is missing headers for vector/memory.  This causes downstream issues when building runner-et.
+  popd
+}
+
+install_executorch_python_libs() {
+  if [ ! -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}" ]; then
+    echo "Directory ${TORCHCHAT_ROOT}/${ET_BUILD_DIR} does not exist."
+    echo "Make sur eyou run clone_executorch"
+    exit 1
+  fi
+  pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
+  cd executorch
+
+  echo "Building and installing python libraries"
+  echo "Building and installing python libraries"
+  if [ "${ENABLE_ET_PYBIND}" = false ]; then
+      echo "Not installing pybind"
+      bash ./install_requirements.sh
+  else
+      echo "Installing pybind"
+      bash ./install_requirements.sh --pybind xnnpack
+  fi
+  pip3 list
+  popd
+}
+
+install_executorch() {
+  # AOT lib has to be build for model export
+  # So by default it is built, and you can explicitly opt-out
+  EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR=OFF
+  if [ "${EXECUTORCH_BUILD_CUSTOM_OPS_AOT}" == "" ]; then
+    EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR=ON
+  fi
+
+  # but for runner not
+  EXECUTORCH_BUILD_CUSTOM_VAR=OFF
+  if [ ! ["${EXECUTORCH_BUILD_CUSTOM}" == ""] ]; then
+    EXECUTORCH_BUILD_CUSTOM_VAR=ON
+  fi
+  echo "${EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR}"
+  echo "${EXECUTORCH_BUILD_CUSTOM_VAR}"
+  if [ ! -d "${TORCHCHAT_ROOT}/${ET_BUILD_DIR}" ]; then
+    echo "Directory ${TORCHCHAT_ROOT}/${ET_BUILD_DIR} does not exist."
+    echo "Make sure you run clone_executorch"
+    exit 1
+  fi
+  pushd ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/src
+  cd executorch
+
+  if [ "${CMAKE_OUT_DIR}" == "" ]; then
+    CMAKE_OUT_DIR="cmake-out"
+  fi
+
+  CROSS_COMPILE_ARGS=""
+  if [ "${CMAKE_OUT_DIR}" == "cmake-out-android" ]; then
+    CROSS_COMPILE_ARGS="-DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE} -DANDROID_ABI=${ANDROID_ABI} -DANDROID_PLATFORM=${ANDROID_PLATFORM}"
+  fi
+
+  echo "Building and installing C++ libraries"
+  echo "Inside: ${PWD}"
+  rm -rf ${CMAKE_OUT_DIR}
+  mkdir ${CMAKE_OUT_DIR}
+  cmake -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} -DCMAKE_BUILD_TYPE=Release -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=Info -DEXECUTORCH_BUILD_CUSTOM_OPS_AOT=${EXECUTORCH_BUILD_CUSTOM_OPS_AOT_VAR} -DEXECUTORCH_BUILD_CUSTOM=${EXECUTORCH_BUILD_CUSTOM_VAR} -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_OPTIMIZED=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON -DEXECUTORCH_BUILD_XNNPACK=ON ${CROSS_COMPILE_ARGS} -S . -B ${CMAKE_OUT_DIR} -G Ninja
+  cmake --build ${CMAKE_OUT_DIR}
+  cmake --install ${CMAKE_OUT_DIR} --prefix ${TORCHCHAT_ROOT}/${ET_BUILD_DIR}/install
+  popd
+}
