• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/bin/bash
2# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16
17set -e
18set -x
19
20SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
21ROOT_DIR="$(cd "${SCRIPT_DIR}/../../../" && pwd)"
22
23function print_usage {
24  echo "Usage:"
25  echo "  $(basename ${BASH_SOURCE}) \\"
26  echo "    --input_models=model1.tflite,model2.tflite \\"
27  echo "    --target_archs=x86,x86_64,arm64-v8a,armeabi-v7a \\"
28  echo "    --tflite_custom_ops_srcs=file1.cc,file2.h \\"
29  echo "    --tflite_custom_ops_deps=dep1,dep2"
30  echo ""
31  echo "Where: "
32  echo "  --input_models: Supported TFLite models. "
33  echo "  --target_archs: Supported arches included in the aar file."
34  echo "  --tflite_custom_ops_srcs: The src files for building additional TFLite custom ops if any."
35  echo "  --tflite_custom_ops_deps: Dependencies for building additional TFLite custom ops if any."
36  echo ""
37  exit 1
38}
39
40function generate_list_field {
41  local name="$1"
42  local list_string="$2"
43  local list=(${list_string//,/ })
44
45  local message=("$name=[")
46  for item in "${list[@]}"
47  do
48    message+=("\"$item\",")
49  done
50  message+=('],')
51  printf '%s' "${message[@]}"
52}
53
54function print_output {
55  if [ -z OMIT_PRINTING_OUTPUT_PATHS ]; then
56    echo "Output can be found here:"
57    for i in "$@"
58    do
59      # Check if the file exist.
60      ls -1a ${ROOT_DIR}/$i
61    done
62  fi
63}
64
65function generate_tflite_aar {
66  pushd ${TMP_DIR} > /dev/null
67  # Generate the BUILD file.
68  message=(
69    'load("//tensorflow/lite:build_def.bzl", "tflite_custom_android_library")'
70    'load("//tensorflow/lite/java:aar_with_jni.bzl", "aar_with_jni")'
71    ''
72    'tflite_custom_android_library('
73    '    name = "custom_tensorflowlite",'
74  )
75  message+=('    '$(generate_list_field "models" $MODEL_NAMES))
76  message+=('    '$(generate_list_field "srcs" $TFLITE_OPS_SRCS))
77  message+=('    '$(generate_list_field "deps" $FLAG_TFLITE_OPS_DEPS))
78  message+=(
79    ')'
80    ''
81    'aar_with_jni('
82    '    name = "tensorflow-lite",'
83    '    android_library = ":custom_tensorflowlite",'
84    ')'
85    ''
86  )
87  printf '%s\n' "${message[@]}" >> BUILD
88
89  # Build the aar package.
90  popd > /dev/null
91  bazel ${CACHE_DIR_FLAG} build -c opt --cxxopt='--std=c++17' \
92        --fat_apk_cpu=${TARGET_ARCHS} \
93        --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
94        //tmp:tensorflow-lite
95
96   OUT_FILES="${OUT_FILES} bazel-bin/tmp/tensorflow-lite.aar"
97}
98
99function generate_flex_aar {
100  pushd ${TMP_DIR}
101  # Generating the BUILD file.
102  message=(
103    'load("//tensorflow/lite/delegates/flex:build_def.bzl", "tflite_flex_android_library")'
104    'load("//tensorflow/lite/java:aar_with_jni.bzl", "aar_with_jni")'
105    ''
106    'tflite_flex_android_library('
107    '    name = "custom_tensorflowlite_flex",'
108    )
109  message+=('    '$(generate_list_field "models" $MODEL_NAMES))
110  message+=(
111    ')'
112    ''
113    'aar_with_jni('
114    '    name = "tensorflow-lite-select-tf-ops",'
115    '    android_library = ":custom_tensorflowlite_flex",'
116    ')'
117  )
118  printf '%s\n' "${message[@]}" >> BUILD
119
120  cp ${ROOT_DIR}/tensorflow/lite/java/AndroidManifest.xml .
121  cp ${ROOT_DIR}/tensorflow/lite/java/proguard.flags .
122  popd
123
124  # Build the aar package.
125  bazel ${CACHE_DIR_FLAG} build -c opt --cxxopt='--std=c++17' \
126      --config=monolithic \
127      --fat_apk_cpu=${TARGET_ARCHS} \
128      --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
129      //tmp:tensorflow-lite-select-tf-ops
130
131  OUT_FILES="${OUT_FILES} bazel-bin/tmp/tensorflow-lite-select-tf-ops.aar"
132}
133
134# Check command line flags.
135TARGET_ARCHS=x86,x86_64,arm64-v8a,armeabi-v7a
136# If the environmant variable BAZEL_CACHE_DIR is set, use it as the user root
137# directory of bazel.
138if [ ! -z ${BAZEL_CACHE_DIR} ]; then
139  CACHE_DIR_FLAG="--output_user_root=${BAZEL_CACHE_DIR}/cache"
140fi
141
142if [ "$#" -gt 4 ]; then
143  echo "ERROR: Too many arguments."
144  print_usage
145fi
146
147for i in "$@"
148do
149case $i in
150    --input_models=*)
151      FLAG_MODELS="${i#*=}"
152      shift;;
153    --target_archs=*)
154      TARGET_ARCHS="${i#*=}"
155      shift;;
156    --tflite_custom_ops_srcs=*)
157      FLAG_TFLITE_OPS_SRCS="${i#*=}"
158      shift;;
159    --tflite_custom_ops_deps=*)
160      FLAG_TFLITE_OPS_DEPS="${i#*=}"
161      shift;;
162    *)
163      echo "ERROR: Unrecognized argument: ${i}"
164      print_usage;;
165esac
166done
167
168# Check if users already run configure
169cd $ROOT_DIR
170if [ ! -f "$ROOT_DIR/.tf_configure.bazelrc" ]; then
171  echo "ERROR: Please run ./configure first."
172  exit 1
173else
174  if ! grep -q ANDROID_SDK_HOME "$ROOT_DIR/.tf_configure.bazelrc"; then
175    echo "ERROR: Please run ./configure with Android config."
176    exit 1
177  fi
178fi
179
180# Build the standard aar package of no models provided.
181if [ -z ${FLAG_MODELS} ]; then
182  bazel ${CACHE_DIR_FLAG} build -c opt --cxxopt='--std=c++17' \
183    --config=monolithic \
184    --fat_apk_cpu=${TARGET_ARCHS} \
185    --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \
186    //tensorflow/lite/java:tensorflow-lite
187
188  print_output bazel-bin/tensorflow/lite/java/tensorflow-lite.aar
189  exit 0
190fi
191
192# Prepare the tmp directory.
193TMP_DIR="${ROOT_DIR}/tmp/"
194rm -rf ${TMP_DIR} && mkdir -p ${TMP_DIR}
195
196# Copy models to tmp directory.
197MODEL_NAMES=""
198for model in $(echo ${FLAG_MODELS} | sed "s/,/ /g")
199do
200  cp ${model} ${TMP_DIR}
201  MODEL_NAMES="${MODEL_NAMES},$(basename ${model})"
202done
203
204# Copy srcs of additional tflite ops to tmp directory.
205TFLITE_OPS_SRCS=""
206for src_file in $(echo ${FLAG_TFLITE_OPS_SRCS} | sed "s/,/ /g")
207do
208  cp ${src_file} ${TMP_DIR}
209  TFLITE_OPS_SRCS="${TFLITE_OPS_SRCS},$(basename ${src_file})"
210done
211
212# Build the custom aar package.
213generate_tflite_aar
214
215# Build flex aar if one of the models contain flex ops.
216bazel ${CACHE_DIR_FLAG} build -c opt --config=monolithic \
217  //tensorflow/lite/tools:list_flex_ops_no_kernel_main
218bazel-bin/tensorflow/lite/tools/list_flex_ops_no_kernel_main \
219  --graphs=${FLAG_MODELS} > ${TMP_DIR}/ops_list.txt
220if [[ `cat ${TMP_DIR}/ops_list.txt` != "[]" ]]; then
221  generate_flex_aar
222fi
223
224# List the output files.
225rm -rf ${TMP_DIR}
226print_output ${OUT_FILES}
227