Add a way to build TFLite PIP package with Flex support

usage)
CUSTOM_BAZEL_FLAGS=--define=tflite_pip_with_flex=true \
  tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
PiperOrigin-RevId: 313695767
Change-Id: I92f90ac1fc120bf9a1196bbe64ccac4dbb9e55d8

tensorflow/lite/python/interpreter_wrapper/BUILD

@@ -61,6 +61,13 @@ cc_library(
     ],
 )
 
+config_setting(
+    name = "tflite_pip_with_flex",
+    define_values = {
+        "tflite_pip_with_flex": "true",
+    },
+)
+
 pybind_extension(
     name = "_pywrap_tensorflow_interpreter_wrapper",
     srcs = [
@@ -71,10 +78,13 @@ pybind_extension(
     module_name = "_pywrap_tensorflow_interpreter_wrapper",
     deps = [
         ":interpreter_wrapper_lib",
+        "@pybind11",
+        "//third_party/python_runtime:headers",
         "//tensorflow/lite:framework_lib",
         "//tensorflow/lite/experimental/tflite_api_dispatcher",
         "//tensorflow/python:pybind11_lib",
-        "//third_party/python_runtime:headers",
-        "@pybind11",
-    ],
+    ] + select({
+        ":tflite_pip_with_flex": ["//tensorflow/lite/delegates/flex:delegate"],
+        "//conditions:default": [],
+    }),
 )

tensorflow/lite/tools/pip_package/README.md

@@ -93,6 +93,30 @@ CI_DOCKER_EXTRA_PARAMS="-e CI_BUILD_PYTHON=python3 -e CROSSTOOL_PYTHON_INCLUDE_P
   tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh aarch64
 ```
 
+## Enable TF OP support (Flex delegate)
+
+If you want to use TF ops with Python API, you need to enable flex support.
+You can build TFLite interpreter with flex ops support by providing
+"--define=tflite_pip_with_flex=true" to Bazel.
+
+Here are some examples.
+
+### Native build with Flex for your workstation
+
+```sh
+CUSTOM_BAZEL_FLAGS=--define=tflite_pip_with_flex=true \
+  tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
+```
+
+### Cross build with Flex for armhf Python 3.5
+
+```sh
+CI_DOCKER_EXTRA_PARAMS="-e CUSTOM_BAZEL_FLAGS=--define=tflite_pip_with_flex=true \
+  -e CI_BUILD_PYTHON=python3 -e CROSSTOOL_PYTHON_INCLUDE_PATH=/usr/include/python3.5" \
+  tensorflow/tools/ci_build/ci_build.sh PI-PYTHON3 \
+  tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh armhf
+```
+
 ## Usage
 
 Note, unlike tensorflow this will be installed to a tflite_runtime namespace.

tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh

@@ -25,6 +25,21 @@ export PACKAGE_VERSION="${TENSORFLOW_VERSION}${VERSION_SUFFIX}"
 BUILD_DIR="${SCRIPT_DIR}/gen/tflite_pip/${PYTHON}"
 TENSORFLOW_TARGET=$1
 
+# Fix container image for cross build.
+if [ ! -z "${CI_BUILD_HOME}" ] && [ `pwd` = "/workspace" ]; then
+  # Fix for curl build problem in 32-bit, see https://stackoverflow.com/questions/35181744/size-of-array-curl-rule-01-is-negative
+  if [ "${TENSORFLOW_TARGET}" = "armhf" ]; then
+    sudo sed -i 's/define CURL_SIZEOF_LONG 8/define CURL_SIZEOF_LONG 4/g' /usr/include/curl/curlbuild.h
+    sudo sed -i 's/define CURL_SIZEOF_CURL_OFF_T 8/define CURL_SIZEOF_CURL_OFF_T 4/g' /usr/include/curl/curlbuild.h
+  fi
+
+  # The system-installed OpenSSL headers get pulled in by the latest BoringSSL
+  # release on this configuration, so move them before we build:
+  if [ -d /usr/include/openssl ]; then
+    sudo mv /usr/include/openssl /usr/include/openssl.original
+  fi
+fi
+
 # Build source tree.
 rm -rf "${BUILD_DIR}" && mkdir -p "${BUILD_DIR}/tflite_runtime"
 cp -r "${TENSORFLOW_LITE_DIR}/tools/pip_package/debian" \
@@ -40,14 +55,16 @@ echo "__git_version__ = '$(git -C "${TENSORFLOW_DIR}" describe)'" >> "${BUILD_DI
 # Build python interpreter_wrapper.
 cd "${BUILD_DIR}"
 case "${TENSORFLOW_TARGET}" in
-  rpi|armhf)
+  armhf)
     BAZEL_FLAGS="--config=elinux_armhf
       --copt=-march=armv7-a --copt=-mfpu=neon-vfpv4
       --copt=-O3 --copt=-fno-tree-pre --copt=-fpermissive
+      --define tensorflow_mkldnn_contraction_kernel=0
       --define=raspberry_pi_with_neon=true"
     ;;
   aarch64)
     BAZEL_FLAGS="--config=elinux_aarch64
+      --define tensorflow_mkldnn_contraction_kernel=0
       --copt=-O3"
     ;;
   *)
@@ -58,14 +75,15 @@ esac
 # include path for Python 3.x builds to work.
 export CROSSTOOL_PYTHON_INCLUDE_PATH
 
-bazel build -c opt -s --config=monolithic ${BAZEL_FLAGS} //tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper
+bazel build -c opt -s --config=monolithic --config=noaws --config=nogcp --config=nohdfs --config=nonccl \
+  ${BAZEL_FLAGS} ${CUSTOM_BAZEL_FLAGS} //tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper
 cp "${TENSORFLOW_DIR}/bazel-bin/tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper.so" \
    "${BUILD_DIR}/tflite_runtime"
 
 # Build python wheel.
 cd "${BUILD_DIR}"
 case "${TENSORFLOW_TARGET}" in
-  rpi|armhf)
+  armhf)
     ${PYTHON} setup_with_bazel.py bdist --plat-name=linux-armv7l \
                        bdist_wheel --plat-name=linux-armv7l
     ;;
@@ -111,7 +129,7 @@ EOF
 fi
 
 case "${TENSORFLOW_TARGET}" in
-  rpi|armhf)
+  armhf)
     dpkg-buildpackage -b -rfakeroot -us -uc -tc -d -a armhf
     ;;
   aarch64)

tensorflow/tools/ci_build/Dockerfile.pi

@@ -22,3 +22,6 @@ RUN /install/install_pi_toolchain.sh
 
 # Set up the master bazelrc configuration file.
 COPY install/.bazelrc /etc/bazel.bazelrc
+
+# XLA is not needed for PI
+ENV TF_ENABLE_XLA=0

tensorflow/tools/ci_build/Dockerfile.pi-python3

@@ -21,3 +21,6 @@ RUN /install/install_pi_python3_toolchain.sh
 
 # Set up the master bazelrc configuration file.
 COPY install/.bazelrc /etc/bazel.bazelrc
+
+# XLA is not needed for PI
+ENV TF_ENABLE_XLA=0

tensorflow/tools/ci_build/Dockerfile.pi-python37

@@ -21,3 +21,6 @@ RUN /install/install_pi_python37_toolchain.sh
 
 # Set up the master bazelrc configuration file.
 COPY install/.bazelrc /etc/bazel.bazelrc
+
+# XLA is not needed for PI
+ENV TF_ENABLE_XLA=0