Refactor android demo job (pytorch#4288)

Summary:
We would want to reuse the same demo app to benchmark as many models ad possible. It may be not easy to create super generic app for all types of models, but we can reuse our existing demo apps to swap in different models of performing same task, e.g. our llama demo should be able to benchmark different casual LLMs w/o problems. To do this, we need to organize the build vertically by the demo apps. Currently we have two demo apps for android (ios demo app would follow the same rule), this PR is to address the llama demo. The android job 'build-llm-demo' is going to build different flavors of the same app by android-abi and tokenizer library. In the downstream, an app built for arm with bpe tokenizer could be used to benchmark all LLMs using bpe tokenizer on a physical android device.

Pull Request resolved: pytorch#4288

Reviewed By: huydhn, kirklandsign

Differential Revision: D59874919

Pulled By: guangy10

fbshipit-source-id: 11bf280765af9ddd4e5459e47c859cc8d37b3848

.github/workflows/android.yml

@@ -22,15 +22,14 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  build-demo-android:
-    name: build-demo-android
+  build-llm-demo:
+    name: build-llm-demo
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     strategy:
       matrix:
-          tiktoken: [OFF, ON]
+          tokenizer: [bpe, tiktoken]
     with:
-      # NB: The example model dl3 requires lots of memory (T161064121)
-      runner: linux.12xlarge
+      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 }}
@@ -42,28 +41,8 @@ jobs:
         # 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}"
-
-        # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh buck2
-        # Build Android library
-        export EXECUTORCH_USE_TIKTOKEN=${{ matrix.tiktoken }}
-        bash build/build_android_library.sh
-        # Build Android demo app
-        bash build/test_android_ci.sh
+        export ARTIFACTS_DIR_NAME=artifacts-to-be-uploaded
 
-        mkdir -p artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN
-        mkdir -p artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/arm64-v8a/
-        mkdir -p artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/x86_64/
-        # Copy the jar to S3
-        cp extension/android/build/libs/executorch.jar artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/
-        # Copy the app and its test suite to S3
-        cp examples/demo-apps/android/LlamaDemo/app/build/outputs/apk/debug/*.apk artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/
-        cp examples/demo-apps/android/LlamaDemo/app/build/outputs/apk/androidTest/debug/*.apk artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/
-        # Also copy the libraries
-        cp cmake-out-android-arm64-v8a/lib/*.a artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/arm64-v8a/
-        cp cmake-out-android-arm64-v8a/extension/android/*.so artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/arm64-v8a/
-        cp cmake-out-android-x86_64/lib/*.a artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/x86_64/
-        cp cmake-out-android-x86_64/extension/android/*.so artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/x86_64/
-        # Copyp AAR to S3
-        cp executorch.aar artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/
-        cp executorch-llama.aar artifacts-to-be-uploaded/tiktoken_$EXECUTORCH_USE_TIKTOKEN/
+        # Build LLM Demo for Android
+        bash build/build_android_llm_demo.sh ${{ matrix.tokenizer }} ${ARTIFACTS_DIR_NAME}

build/build_android_library.sh → build/build_android_llm_demo.sh

@@ -11,14 +11,21 @@ build_jar() {
   pushd extension/android
   ./gradlew build
   popd
-  cp extension/android/build/libs/executorch.jar "${BUILD_AAR_DIR}/libs"
+  mkdir -p "${BUILD_AAR_DIR}/libs"
+  cp extension/android/build/libs/executorch.jar "${BUILD_AAR_DIR}/libs/"
 }
 
 build_android_native_library() {
   ANDROID_ABI="$1"
+  TOKENIZER="$2"
   ANDROID_NDK="${ANDROID_NDK:-/opt/ndk}"
-  CMAKE_OUT="cmake-out-android-$1"
-  EXECUTORCH_USE_TIKTOKEN="${EXECUTORCH_USE_TIKTOKEN:-OFF}"
+  CMAKE_OUT="cmake-out-android-${ANDROID_ABI}"
+  if [[ $TOKENIZER == "tiktoken" ]]; then
+    EXECUTORCH_USE_TIKTOKEN=ON
+  else
+    EXECUTORCH_USE_TIKTOKEN=OFF
+  fi
+
   cmake . -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
     -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK}/build/cmake/android.toolchain.cmake" \
     -DANDROID_ABI="${ANDROID_ABI}" \
@@ -65,32 +72,66 @@ build_android_native_library() {
 
   cmake --build "${CMAKE_OUT}"/extension/android -j "${CMAKE_JOBS}" --config Release
 
-  cp "${CMAKE_OUT}"/extension/android/*.so "${BUILD_AAR_DIR}/jni/$1/"
+  # Copy artifacts to ABI specific directory
+  mkdir -p "${BUILD_AAR_DIR}/jni/${ANDROID_ABI}"
+  cp "${CMAKE_OUT}"/extension/android/*.so "${BUILD_AAR_DIR}/jni/${ANDROID_ABI}/"
 }
 
 build_aar() {
   echo \<manifest xmlns:android=\"http://schemas.android.com/apk/res/android\" \
-   package=\"org.pytorch.executorch\"\> \
-   \<uses-sdk android:minSdkVersion=\"19\" /\> \
-   \</manifest\> > "${BUILD_AAR_DIR}/AndroidManifest.xml"
+  package=\"org.pytorch.executorch\"\> \
+  \<uses-sdk android:minSdkVersion=\"19\" /\> \
+  \</manifest\> > "${BUILD_AAR_DIR}/AndroidManifest.xml"
   pushd "${BUILD_AAR_DIR}"
   # Rename libexecutorch_jni.so to libexecutorch.so for soname consistency
   # between Java and JNI
-  mv jni/arm64-v8a/libexecutorch_jni.so jni/arm64-v8a/libexecutorch.so
-  mv jni/x86_64/libexecutorch_jni.so jni/x86_64/libexecutorch.so
-  zip -r executorch.aar libs jni/arm64-v8a/libexecutorch.so jni/x86_64/libexecutorch.so AndroidManifest.xml
+  find jni -type f -name "libexecutorch_jni.so" -exec bash -c 'mv "$1" "${1/_jni/}"' bash {} \;
+  # Zip all necessary files into the AAR file
+  zip -r executorch.aar libs jni/*/libexecutorch.so AndroidManifest.xml
+  zip -r executorch-llama.aar libs jni/*/libexecutorch_llama_jni.so AndroidManifest.xml
+  popd
+}
 
-  rm jni/arm64-v8a/libexecutorch.so jni/x86_64/libexecutorch.so
-  zip -r executorch-llama.aar libs jni/arm64-v8a/libexecutorch_llama_jni.so jni/x86_64/libexecutorch_llama_jni.so AndroidManifest.xml
+build_android_llm_demo_app() {
+  mkdir -p examples/demo-apps/android/LlamaDemo/app/libs
+  cp ${BUILD_AAR_DIR}/executorch-llama.aar examples/demo-apps/android/LlamaDemo/app/libs
+  pushd examples/demo-apps/android/LlamaDemo
+  ANDROID_HOME="${ANDROID_SDK:-/opt/android/sdk}" ./gradlew build assembleAndroidTest
   popd
-  cp "${BUILD_AAR_DIR}/executorch-llama.aar" .
-  cp "${BUILD_AAR_DIR}/executorch.aar" .
+}
+
+collect_artifacts_to_be_uploaded() {
+  TOKENIZER="$1"
+  ARTIFACTS_DIR_NAME="$2"
+  DEMO_APP_DIR="${ARTIFACTS_DIR_NAME}/llm_demo_${TOKENIZER}"
+  # The app directory is named using its build flavor as a suffix.
+  mkdir -p "${DEMO_APP_DIR}"
+  # Collect the app and its test suite
+  cp examples/demo-apps/android/LlamaDemo/app/build/outputs/apk/debug/*.apk "${DEMO_APP_DIR}"
+  cp examples/demo-apps/android/LlamaDemo/app/build/outputs/apk/androidTest/debug/*.apk "${DEMO_APP_DIR}"
+  # Collect all ABI specific libraries
+  for ANDROID_ABI in "${ANDROID_ABIS[@]}"; do
+    mkdir -p "${DEMO_APP_DIR}/${ANDROID_ABI}"
+    cp cmake-out-android-${ANDROID_ABI}/lib/*.a "${DEMO_APP_DIR}/${ANDROID_ABI}/"
+    cp cmake-out-android-${ANDROID_ABI}/extension/android/*.so "${DEMO_APP_DIR}/${ANDROID_ABI}/"
+  done
+  # Collect JAR and AAR
+  cp extension/android/build/libs/executorch.jar "${DEMO_APP_DIR}"
+  find "${BUILD_AAR_DIR}/" -name 'executorch*.aar' -exec cp {} "${DEMO_APP_DIR}" \;
 }
 
 BUILD_AAR_DIR="$(mktemp -d)"
 export BUILD_AAR_DIR
-mkdir -p "${BUILD_AAR_DIR}/jni/arm64-v8a" "${BUILD_AAR_DIR}/jni/x86_64" "${BUILD_AAR_DIR}/libs"
+ANDROID_ABIS=("arm64-v8a" "x86_64")
+export ANDROID_ABIS
+
+TOKENIZER="${1:-bpe}"
+ARTIFACTS_DIR_NAME="$2"
+
 build_jar
-build_android_native_library arm64-v8a
-build_android_native_library x86_64
+for ANDROID_ABI in "${ANDROID_ABIS[@]}"; do
+  build_android_native_library ${ANDROID_ABI} ${TOKENIZER}
+done
 build_aar
+build_android_llm_demo_app
+collect_artifacts_to_be_uploaded ${TOKENIZER} ${ARTIFACTS_DIR_NAME}

build/test_android_ci.sh

@@ -26,15 +26,5 @@ build_android_demo_app() {
   popd
 }
 
-build_android_llama_demo_app() {
-  mkdir -p examples/demo-apps/android/LlamaDemo/app/libs
-  cp executorch-llama.aar examples/demo-apps/android/LlamaDemo/app/libs
-  pushd examples/demo-apps/android/LlamaDemo
-  ANDROID_HOME=/opt/android/sdk ./gradlew build
-  ANDROID_HOME=/opt/android/sdk ./gradlew assembleAndroidTest
-  popd
-}
-
 export_model
 build_android_demo_app
-build_android_llama_demo_app