Feature/flexops

Author: jishminor

CMakeLists.txt

@@ -18,7 +18,10 @@ set(TRITON_COMMON_REPO_TAG "main" CACHE STRING "Tag for triton-inference-server/
 
 # TFLite Options
 set(TFLITE_TAG "v2.4.1" CACHE STRING "Version of TFLite to build")
+set(BAZEL_BINARY "/usr/bin/bazel" CACHE STRING "Location of bazel binary on host")
+option(TFLITE_BAZEL_BUILD "Build tflite bazel" OFF)
 option(TFLITE_ENABLE_RUY "Use RUY library in TFLite build" ON)
+option(TFLITE_ENABLE_FLEX_OPS "Build tflite with flex ops support (must build with bazel)" OFF)
 
 # ArmNN Options
 set(ARMNN_TAG "v21.08" CACHE STRING "Tag for ArmNN repo")
@@ -44,6 +47,11 @@ if (TRITON_ENABLE_MALI_GPU)
   set(ACL_ENABLE_CL "1")
 endif()
 
+# Ensure bazel build set if enabling flex ops delegate
+if (TFLITE_ENABLE_FLEX_OPS AND (NOT(TFLITE_BAZEL_BUILD)))
+  message(FATAL_ERROR "Please set TFLITE_BAZEL_BUILD if you wish to build flex ops delegate")
+endif()
+
 if (NOT DEFINED ACL_ARCH)
   if (${TARGET_ARCH} MATCHES "armv7")
     set(ACL_ARCH "armv7a")
@@ -95,15 +103,23 @@ FetchContent_Declare(
   GIT_TAG ${TRITON_BACKEND_REPO_TAG}
   GIT_SHALLOW ON
 )
-FetchContent_Declare(
-  tensorflow-lite
-  GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git
-  GIT_TAG ${TFLITE_TAG}
-  GIT_SHALLOW ON
-  PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/xnnpack_commit.patch || true # patch updates the commit hash for xnnpack for tensorflow v2.4.1
-  SOURCE_SUBDIR tensorflow/lite
-)
-FetchContent_MakeAvailable(repo-common repo-core repo-backend tensorflow-lite)
+
+set(MAKE_AVAILABLE_LIST repo-common repo-core repo-backend)
+
+if (NOT TFLITE_BAZEL_BUILD)
+  FetchContent_Declare(
+    tensorflow-lite
+    GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git
+    GIT_TAG ${TFLITE_TAG}
+    GIT_SHALLOW ON
+    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/xnnpack_commit.patch || true # patch updates the commit hash for xnnpack for tensorflow v2.4.1
+    SOURCE_SUBDIR tensorflow/lite
+  )
+
+  LIST(APPEND MAKE_AVAILABLE_LIST tensorflow-lite)
+endif()
+
+FetchContent_MakeAvailable(${MAKE_AVAILABLE_LIST})
 
 #
 # Shared library implementing the Triton Backend API
@@ -112,6 +128,50 @@ configure_file(src/libtriton_armnn_tflite.ldscript libtriton_armnn_tflite.ldscri
 
 include(ExternalProject)
 
+set(TFLITE_LOCATION ${CMAKE_CURRENT_BINARY_DIR}/external/tensorflow_lite)
+
+if (TFLITE_BAZEL_BUILD)
+  SET(BAZEL_BUILD_FLAGS 
+      "--define=tflite_with_xnnpack=true"
+      "--define=with_select_tf_ops=true"
+      "--config=monolithic"
+      "-c" "opt"
+  )
+
+  if (TFLITE_ENABLE_RUY)
+    LIST(APPEND BAZEL_BUILD_FLAGS "--define=tflite_with_ruy=true")
+  else()
+    LIST(APPEND BAZEL_BUILD_FLAGS "--define=tflite_with_ruy=false")
+  endif()
+
+  if (TFLITE_ENABLE_FLEX_OPS)
+    LIST(APPEND BAZEL_BUILD_FLAGS "--define=with_select_tf_ops=true")
+  else()
+    LIST(APPEND BAZEL_BUILD_FLAGS "--define=with_select_tf_ops=false")
+  endif()
+
+  if (CMAKE_BUILD_TYPE STREQUAL "Debug")
+    LIST(APPEND BAZEL_BUILD_FLAGS "-c" "dbg" "--copt=-g")
+  else()
+    LIST(APPEND BAZEL_BUILD_FLAGS "--strip=always")
+  endif()
+  
+  # Build tensorflow lite shared lib using bazel (required for flex ops delegate)
+  ExternalProject_Add(
+    tensorflow-lite
+    PREFIX ${TFLITE_LOCATION}
+    GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git
+    GIT_TAG ${TFLITE_TAG}
+    GIT_SHALLOW ON
+    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/xnnpack_commit.patch || true # patch updates the commit hash for xnnpack for tensorflow v2.4.1
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ${BAZEL_BINARY} build ${BAZEL_BUILD_FLAGS} //tensorflow/lite:libtensorflowlite.so
+    BUILD_IN_SOURCE ON
+    INSTALL_COMMAND ""
+  )
+endif()
+
+
 if (ARMNN_DELEGATE_ENABLE)
   #
   # Build ArmNN dependencies using custom command
@@ -142,15 +202,26 @@ if (ARMNN_DELEGATE_ENABLE)
     INSTALL_COMMAND ""
   )
 
+  if (TFLITE_BAZEL_BUILD)
+    SET(ARMNN_PATCH_COMMAND "")
+    SET(TENSORFLOW_ROOT ${TFLITE_LOCATION}/src/tensorflow-lite)
+    SET(TFLITE_LIB_ROOT ${TFLITE_LOCATION}/src/tensorflow-lite/bazel-bin)
+  else()
+    # Never fail as patch is always applied after repo clone
+    SET(ARMNN_PATCH_COMMAND "git" "apply" "${CMAKE_CURRENT_SOURCE_DIR}/patches/tflite_armnn_delegate.patch" "||" "true")
+    SET(TENSORFLOW_ROOT ${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-src)
+    SET(TFLITE_LIB_ROOT ${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-build)
+  endif()
+
   # ArmNN and its tflite delegate
   ExternalProject_Add(
     armnn
     PREFIX ${ARMNN_LOCATION}
     GIT_REPOSITORY https://review.mlplatform.org/ml/armnn
     GIT_TAG ${ARMNN_TAG}
     GIT_SHALLOW ON
-    PATCH_COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/patches/tflite_armnn_delegate.patch || true # Never fail as patch is always applied after repo clone
-    CMAKE_ARGS -DCMAKE_INSTALL_PREFIX:PATH=${ARMNN_LOCATION} -DBUILD_ARMNN_TFLITE_DELEGATE=1 -DTENSORFLOW_ROOT=${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-src -DTFLITE_LIB_ROOT=${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-build -DARMCOMPUTE_ROOT=${ACL_LOCATION}/src/acl -DARMCOMPUTENEON=${ACL_ENABLE_NEON} -DARMCOMPUTECL=${ACL_ENABLE_CL} -DBUILD_UNIT_TESTS=0 -DCMAKE_BUILD_TYPE=${ARMNN_BUILD_TYPE} -DBUILD_ARMNN_SERIALIZER=0 -DARMNNREF=1 -DFLATBUFFERS_ROOT=${FLATBUFFERS_LOCATION}/flatbuffers-${FLATBUFFERS_VERSION}/install
+    PATCH_COMMAND ${ARMNN_PATCH_COMMAND}
+    CMAKE_ARGS -DCMAKE_INSTALL_PREFIX:PATH=${ARMNN_LOCATION} -DBUILD_ARMNN_TFLITE_DELEGATE=1 -DTENSORFLOW_ROOT=${TENSORFLOW_ROOT} -DTFLITE_LIB_ROOT=${TFLITE_LIB_ROOT} -DARMCOMPUTE_ROOT=${ACL_LOCATION}/src/acl -DARMCOMPUTENEON=${ACL_ENABLE_NEON} -DARMCOMPUTECL=${ACL_ENABLE_CL} -DBUILD_UNIT_TESTS=0 -DCMAKE_BUILD_TYPE=${ARMNN_BUILD_TYPE} -DBUILD_ARMNN_SERIALIZER=0 -DARMNNREF=1 -DFLATBUFFERS_ROOT=${FLATBUFFERS_LOCATION}/flatbuffers-${FLATBUFFERS_VERSION}/install
     DEPENDS flatbuffers-armnn acl tensorflow-lite
   )
 endif()
@@ -174,15 +245,23 @@ add_library(
   TritonArmNNTFLiteBackend::triton-armnn-tflite-backend ALIAS triton-armnn-tflite-backend
 )
 
-target_include_directories(
-  triton-armnn-tflite-backend
-  PRIVATE
+SET(BACKEND_INCLUDE_DIRS 
     ${CMAKE_CURRENT_SOURCE_DIR}/src
-    ${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-src
+    ${TENSORFLOW_ROOT} # for tflite headers
     ${ARMNN_LOCATION}/include # for armnn headers
     ${ARMNN_LOCATION}/src/armnn/delegate/include # for delegate headers
 )
 
+if (TFLITE_BAZEL_BUILD)
+  LIST(APPEND BACKEND_INCLUDE_DIRS ${TENSORFLOW_ROOT}/bazel-tensorflow-lite/external/flatbuffers/include)
+endif()
+
+target_include_directories(
+  triton-armnn-tflite-backend
+  PRIVATE
+    ${BACKEND_INCLUDE_DIRS}
+)
+
 target_compile_features(triton-armnn-tflite-backend PRIVATE cxx_std_11)
 target_compile_options(
   triton-armnn-tflite-backend PRIVATE
@@ -211,15 +290,30 @@ set_target_properties(
     LINK_FLAGS "-Wl,--no-as-needed,--version-script libtriton_armnn_tflite.ldscript"
 )
 
+SET(BACKEND_LINK_LIBS 
+    triton-core-serverapi 
+    triton-core-backendapi 
+    triton-core-serverstub 
+    triton-backend-utils 
+    ${CMAKE_DL_LIBS}
+)
+
+if (TFLITE_BAZEL_BUILD)
+  # Link the tensorflow lite library
+  target_link_libraries(
+    triton-armnn-tflite-backend
+    PRIVATE
+      "-L${TFLITE_LOCATION}/src/tensorflow-lite/bazel-bin/tensorflow/lite" # From bazel tflite build
+      -ltensorflowlite
+  )
+else()
+  LIST(APPEND BACKEND_LINK_LIBS tensorflow-lite)
+endif()
+
 target_link_libraries(
   triton-armnn-tflite-backend
   PRIVATE
-    triton-core-serverapi  # from repo-core
-    triton-core-backendapi # from repo-core
-    triton-core-serverstub # from repo-core
-    triton-backend-utils   # from repo-backend
-    tensorflow-lite
-    ${CMAKE_DL_LIBS}
+    ${BACKEND_LINK_LIBS}
 )
 
 if (ARMNN_DELEGATE_ENABLE)
@@ -264,10 +358,19 @@ if (ARMNN_DELEGATE_ENABLE)
   )
 endif()
 
+if (TFLITE_BAZEL_BUILD)
+  # Install tflite library
+  install(
+    FILES
+      ${TFLITE_LOCATION}/src/tensorflow-lite/bazel-bin/tensorflow/lite/libtensorflowlite.so
+    DESTINATION ${CMAKE_INSTALL_PREFIX}/backends/armnn_tflite
+  )
+endif()
+
 # Install Tensorflow license
 install(
   FILES
-  ${CMAKE_CURRENT_BINARY_DIR}/_deps/tensorflow-lite-src/LICENSE
+  ${TENSORFLOW_ROOT}/LICENSE
   RENAME tensorflow.LICENSE
   DESTINATION ${CMAKE_INSTALL_PREFIX}/backends/armnn_tflite
 )

Dockerfile

@@ -6,7 +6,7 @@ FROM ubuntu:${UBUNTU_VERSION} as armnn_tflite_backend
 ARG TRITON_REPO_TAG=main
 
 # Cmake Version options
-ARG CMAKE_VERSION=3.19
+ARG CMAKE_VERSION=3.21.1
 
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -19,22 +19,38 @@ RUN apt-get update && \
     curl \
     autoconf \
     libtool \
+    python3-dev \
+    python3-pip \
+    python3-numpy \
     build-essential \
     libssl-dev \
+    zlib1g-dev \
+    default-jdk \
+    libtool \
+    zip \
+    unzip \
     xxd \
     rapidjson-dev \
-    unzip
+    software-properties-common \
+    unzip && \
+    wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | \
+    gpg --dearmor - |  \
+    tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null && \
+    apt-add-repository 'deb https://apt.kitware.com/ubuntu/ focal main' && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+    cmake-data=${CMAKE_VERSION}-0kitware1ubuntu20.04.1 cmake=${CMAKE_VERSION}-0kitware1ubuntu20.04.1 && \
+    pip3 install -U pip wheel && \
+    rm -rf /var/lib/apt/lists/*
 
-# Install cmake from source
-RUN build=1 && \
-    mkdir /temp && \
-    cd /temp && \
-    wget https://cmake.org/files/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}.$build.tar.gz && \
-    tar -xzvf cmake-${CMAKE_VERSION}.$build.tar.gz && \
-    cd cmake-${CMAKE_VERSION}.$build/ && \
-    ./bootstrap --parallel=$(nproc) && \
-    make -j$(nproc) && \
-    make install
+# Install Bazel from source
+RUN wget -O bazel-3.1.0-dist.zip https://github.com/bazelbuild/bazel/releases/download/3.1.0/bazel-3.1.0-dist.zip && \
+    unzip -d bazel bazel-3.1.0-dist.zip && \
+    rm bazel-3.1.0-dist.zip && \
+    cd bazel && \
+    ln -s /usr/bin/python3 /usr/bin/python && \
+    env EXTRA_BAZEL_ARGS="--host_javabase=@local_jdk//:jdk" bash ./compile.sh && \
+    cp output/bazel /usr/bin/bazel
 
 # Build ArmNN TFLite Backend
 WORKDIR /opt/armnn_tflite_backend
@@ -49,6 +65,5 @@ RUN mkdir build && \
     -DTRITON_ENABLE_GPU=OFF \
     -DTRITON_ENABLE_MALI_GPU=ON \
     -DTFLITE_ENABLE_RUY=ON \
-    -DJOBS=$(nproc) \
-    && \
+    -DJOBS=$(nproc) && \
     make -j$(nproc) install