1#!/bin/bash 2############################################################################### 3# This script tests the flow to build libtorch locally with optimized binary 4# size for mobile devices and the flow to integrate it with a simple predictor 5# in c++. 6# 7# Supported custom build types: 8# 9# 1. `TEST_DEFAULT_BUILD=1 ./build.sh` - it is similar to the prebuilt libtorch 10# libraries released for Android and iOS (same CMake build options + host 11# toolchain), which doesn't contain autograd function nor backward ops thus is 12# smaller than full LibTorch. 13# 14# 2. `TEST_CUSTOM_BUILD_STATIC=1 ./build.sh` - optimizes libtorch size by only 15# including ops used by a specific model. It relies on the static dispatch + 16# linker to prune code. 17# 18############################################################################### 19 20set -ex -o pipefail 21 22SRC_ROOT="$( cd "$(dirname "$0")"/../../.. ; pwd -P)" 23TEST_SRC_ROOT="${SRC_ROOT}/test/mobile/custom_build" 24BUILD_ROOT="${BUILD_ROOT:-${SRC_ROOT}/build_test_custom_build}" 25mkdir -p "${BUILD_ROOT}" 26cd "${BUILD_ROOT}" 27 28prepare_model_and_dump_root_ops() { 29 cd "${BUILD_ROOT}" 30 MODEL="${BUILD_ROOT}/MobileNetV2.pt" 31 ROOT_OPS="${BUILD_ROOT}/MobileNetV2.yaml" 32 33 python "${TEST_SRC_ROOT}/prepare_model.py" 34} 35 36run_default_build() { 37 LIBTORCH_BUILD_ROOT="${BUILD_ROOT}/build_default_libtorch" 38 LIBTORCH_INSTALL_PREFIX="${LIBTORCH_BUILD_ROOT}/install" 39 40 BUILD_ROOT="${LIBTORCH_BUILD_ROOT}" \ 41 "${SRC_ROOT}/scripts/build_mobile.sh" 42} 43 44run_custom_build_with_static_dispatch() { 45 LIBTORCH_BUILD_ROOT="${BUILD_ROOT}/build_custom_libtorch_static" 46 LIBTORCH_INSTALL_PREFIX="${LIBTORCH_BUILD_ROOT}/install" 47 48 # Here it generates registration code for used ROOT ops only, whose unboxing 49 # kernels are still needed by the JIT runtime. The intermediate ops will be 50 # automatically kepted by the linker as they are statically referenced by the 51 # static dispatch code, for which we can bypass the registration. 52 # We don't set '-DSTATIC_DISPATCH_BACKEND=CPU' explicitly to test automatic 53 # fallback to static dispatch. 54 BUILD_ROOT="${LIBTORCH_BUILD_ROOT}" \ 55 "${SRC_ROOT}/scripts/build_mobile.sh" \ 56 -DCMAKE_CXX_FLAGS="-DSTRIP_ERROR_MESSAGES" \ 57 -DSELECTED_OP_LIST="${ROOT_OPS}" 58} 59 60build_predictor() { 61 PREDICTOR_BUILD_ROOT="${BUILD_ROOT}/predictor" 62 63 rm -rf "${PREDICTOR_BUILD_ROOT}" && mkdir -p "${PREDICTOR_BUILD_ROOT}" 64 cd "${PREDICTOR_BUILD_ROOT}" 65 66 cmake "${TEST_SRC_ROOT}" \ 67 -DCMAKE_PREFIX_PATH="${LIBTORCH_INSTALL_PREFIX}" \ 68 -DCMAKE_BUILD_TYPE=Release 69 70 make 71} 72 73run_predictor() { 74 cd "${PREDICTOR_BUILD_ROOT}" 75 ./Predictor "${MODEL}" > output.txt 76 77 if cmp -s output.txt "${TEST_SRC_ROOT}/expected_output.txt"; then 78 echo "Test result is the same as expected." 79 else 80 echo "Test result is DIFFERENT from expected!" 81 diff output.txt "${TEST_SRC_ROOT}/expected_output.txt" 82 exit 1 83 fi 84} 85 86test_default_build() { 87 prepare_model_and_dump_root_ops 88 run_default_build 89 build_predictor 90 run_predictor 91} 92 93test_custom_build_with_static_dispatch() { 94 prepare_model_and_dump_root_ops 95 run_custom_build_with_static_dispatch 96 build_predictor 97 run_predictor 98} 99 100if [ -n "${TEST_DEFAULT_BUILD}" ]; then 101 test_default_build 102fi 103 104if [ -n "${TEST_CUSTOM_BUILD_STATIC}" ]; then 105 test_custom_build_with_static_dispatch 106fi 107