Page MenuHomePhabricator

[MLGO] Use TFLite in 'development mode'
ClosedPublic

Authored by mtrofin on Feb 10 2022, 5:55 PM.

Details

Summary

TLite is a lightweight, statically linkable[1], model evaluator, supporting a subset of what the full tensorflow library does, sufficient for the types of scenarios we envision having. It is also faster.

We still use saved models as "source of truth" - 'release' mode's AOT starts from a saved model; and the ML training side operates in terms of saved models.

Using TFLite solves the following problems compared to using the full TF C API:

  • a compiler-friendly implementation for runtime-loadable (as opposed to AOT-embedded) models: it's statically linked; it can be built via cmake;
  • solves an issue we had when building the compiler with both AOT and full TF C API support, whereby, due to a packaging issue on the TF side, we needed to have the pip package and the TF C API library at the same version. We have no such constraints now.

The main liability is it supporting a subset of what the full TF framework does. We do not expect that to cause an issue, but should that be the case, we can always revert back to using the full framework (after also figuring out a way to address the problems that motivated the move to TFLite).

Details:

This change switches the development mode to TFLite. Models are still expected to be placed in a directory - i.e. the parameters to clang don't change; what changes is the directory content: we still need an output_spec.json file; but instead of the saved_model protobuf and the variables directory, we now just have one file, model.tflite.

The change includes a utility showing how to take a saved model and convert it to TFLite, which it uses for testing.

The full TF implementation can still be built (not side-by-side). We intend to remove it shortly, after patching downstream dependencies. The build behavior, however, prioritizes TFLite - i.e. trying to enable both full TF C API and TFLite will just pick TFLite.

[1] thanks to @petrhosek's changes to TFLite's cmake support and its deps!

Diff Detail

Unit TestsFailed

TimeTest
60,070 msx64 debian > ThreadSanitizer-x86_64.ThreadSanitizer-x86_64::restore_stack.cpp
Script: -- : 'RUN: at line 1'; /var/lib/buildkite-agent/builds/llvm-project/build/./bin/clang --driver-mode=g++ -fsanitize=thread -Wall -m64 -msse4.2 -gline-tables-only -I/var/lib/buildkite-agent/builds/llvm-project/compiler-rt/test/tsan/../ -std=c++11 -I/var/lib/buildkite-agent/builds/llvm-project/compiler-rt/test/tsan/../ -O1 /var/lib/buildkite-agent/builds/llvm-project/compiler-rt/test/tsan/restore_stack.cpp -o /var/lib/buildkite-agent/builds/llvm-project/build/projects/compiler-rt/test/tsan/X86_64Config/Output/restore_stack.cpp.tmp && not /var/lib/buildkite-agent/builds/llvm-project/build/projects/compiler-rt/test/tsan/X86_64Config/Output/restore_stack.cpp.tmp 2>&1 | FileCheck /var/lib/buildkite-agent/builds/llvm-project/compiler-rt/test/tsan/restore_stack.cpp

Event Timeline

mtrofin created this revision.Feb 10 2022, 5:55 PM
mtrofin requested review of this revision.Feb 10 2022, 5:55 PM
Herald added a project: Restricted Project. · View Herald TranscriptFeb 10 2022, 5:55 PM

With the tip-of-tree TensorFlow, it should now be possible to use TFLite as a CMake package. I had to make the following changes to your patch:

diff --git a/llvm/CMakeLists.txt b/llvm/CMakeLists.txt
index 4ce0d9c21900..aa5fef9d1d18 100644
--- a/llvm/CMakeLists.txt
+++ b/llvm/CMakeLists.txt
@@ -861,23 +861,9 @@ include(AddLLVM)
 # also leverage the dependency.
 set(TENSORFLOW_SRC_DIR "" CACHE PATH "Path to TensorFlow source")
 if (TENSORFLOW_SRC_DIR)
-  set(TFLITE_ENABLE_XNNPACK OFF)
-  # find_library(tflite_c tensorflowlite_c PATHS ${TENSORFLOW_SRC_DIR}/build NO_DEFAULT_PATH REQUIRED)
-  # find_library(tflite tensorflow-lite PATHS ${TENSORFLOW_SRC_DIR}/build/tensorflow-lite NO_DEFAULT_PATH REQUIRED)
-  # include_directories(${TENSORFLOW_SRC_DIR}/build/abseil-cpp)
-  # include_directories(${TENSORFLOW_SRC_DIR}/build/eigen)
-  # include_directories(${TENSORFLOW_SRC_DIR}/build/flatbuffers/include)
-  message(STATUS "about to")
-  set(LLVM_EXTERNAL_TFLITE_SOURCE_DIR "${TENSORFLOW_SRC_DIR}/tensorflow/lite")
-  set(LLVM_TOOL_TFLITE_BUILD On)
-  add_llvm_external_project(TFLITE)
-  # install(TARGETS tensorflow-lite EXPORT LLVMExports
-  #   ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX} COMPONENT tensorflow-lite)
-  message(STATUS "done")
   set(LLVM_HAVE_TF_API "ON" CACHE BOOL "Full Tensorflow API available")
-  set(Protobuf_USE_STATIC_LIBS ON)
-  find_package(Protobuf REQUIRED)
-  include_directories(${Protobuf_INCLUDE_DIRS})
+  find_package(protobuf REQUIRED)
+  find_package(tensorflow-lite REQUIRED)
 endif()
 
 # For up-to-date instructions for installing the Tensorflow dependency, refer to
diff --git a/llvm/cmake/modules/TensorFlowCompile.cmake b/llvm/cmake/modules/TensorFlowCompile.cmake
index c1bb03db5e96..9d329cd3e81e 100644
--- a/llvm/cmake/modules/TensorFlowCompile.cmake
+++ b/llvm/cmake/modules/TensorFlowCompile.cmake
@@ -123,10 +123,10 @@ function(build_proto)
     set(PBS ${PBS} ${TENSORFLOW_SRC_DIR}/${P}.proto)
   endforeach()
   add_custom_command(OUTPUT ${PB_SRCS} ${PB_HDRS}
-    COMMAND ${PROTOBUF_PROTOC_EXECUTABLE} --proto_path=${TENSORFLOW_SRC_DIR}
-    --cpp_out=${CMAKE_CURRENT_BINARY_DIR} ${PBS})
+    COMMAND protobuf::protoc
+    ARGS --proto_path=${TENSORFLOW_SRC_DIR} --cpp_out=${CMAKE_CURRENT_BINARY_DIR} ${PBS})
   set_source_files_properties(${PB_SRCS} PROPERTIES
     GENERATED 1)
   set(GeneratedMLSources ${GeneratedMLSources} ${PB_SRCS} PARENT_SCOPE)
   set(MLDeps ${MLDeps} ${GeneratedMLSources} PARENT_SCOPE)
-endfunction()
\ No newline at end of file
+endfunction()
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
index 29a6bcda7dbd..6907cb2f3b58 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -25,15 +25,7 @@ if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
       tensorflow/core/protobuf/error_codes
       tensorflow/core/example/feature
       tensorflow/core/example/example)
-    # list(APPEND MLLinkDeps ${tflite} ${tflite_c} ${Protobuf_LIBRARIES} 
-    #   ${TENSORFLOW_SRC_DIR}/build/_deps/ruy-build/libruy.a
-    #   ${TENSORFLOW_SRC_DIR}/build/_deps/farmhash-build/libfarmhash.a
-    #   ${TENSORFLOW_SRC_DIR}/build/_deps/flatbuffers-build/libflatbuffers.a
-    #   ${TENSORFLOW_SRC_DIR}/build/_deps/fft2d-build/libfft2d_fftsg.a
-    #   ${TENSORFLOW_SRC_DIR}/build/_deps/fft2d-build/libfft2d_fftsg2d.a
-    # )
-    # include_directories(${TENSORFLOW_SRC_DIR})
-    list(APPEND MLLinkDeps tensorflow-lite ${Protobuf_LIBRARIES})
+    list(APPEND MLLinkDeps tensorflow-lite::tensorflow-lite protobuf::libprotobuf)
   endif()
 endif()
 
diff --git a/llvm/lib/Analysis/TFUtils.cpp b/llvm/lib/Analysis/TFUtils.cpp
index fbaa47813ecb..24e86c9ca591 100644
--- a/llvm/lib/Analysis/TFUtils.cpp
+++ b/llvm/lib/Analysis/TFUtils.cpp
@@ -50,6 +50,7 @@ void serialize(const Message &SE, std::string *OutStr) {
     *OutStr = SE.SerializeAsString();
   }
 }
+}
 
 namespace llvm {
 class EvaluationResultImpl {
Herald added a project: Restricted Project. · View Herald TranscriptJul 29 2022, 1:27 PM
mtrofin updated this revision to Diff 451483.Aug 10 2022, 8:26 AM

actual patch

mtrofin edited the summary of this revision. (Show Details)Aug 10 2022, 8:29 AM
mtrofin edited the summary of this revision. (Show Details)
mtrofin added reviewers: gjain, ebrevdo, yundiqian.
mtrofin edited the summary of this revision. (Show Details)Aug 10 2022, 8:31 AM
mtrofin edited the summary of this revision. (Show Details)
mtrofin edited the summary of this revision. (Show Details)
mtrofin updated this revision to Diff 451488.Aug 10 2022, 8:35 AM

removed some spurious changes.

llvm/lib/Analysis/models/gen-regalloc-eviction-test-model.py
49 ↗(On Diff #451483)

To reviewers: I'll send this separately before landing this patch.

llvm/unittests/Analysis/TFUtilsTest.cpp
92

This was removed because with TFLite we can eagerly make the determination that the input is invalid. We don't need to wait to evaluate.

ebrevdo requested changes to this revision.Aug 10 2022, 1:37 PM

Some small items.

llvm/CMakeLists.txt
950–957

Should this be LLVM_USE_TFLITE?

957

do you plan to remove this in future?

llvm/include/llvm/Config/llvm-config.h.cmake
102

plan to remove this in future?

llvm/lib/Analysis/TFLiteUtils.cpp
104

Suggest subclassing tflite::StatefulErrorReporter like here so that you can check error status by looking at reporter.message() after the initialization and after running inference. you can even add a bool like has_error which will be cheaper (no copy).

This revision now requires changes to proceed.Aug 10 2022, 1:37 PM
ebrevdo accepted this revision.Aug 11 2022, 8:55 AM

(my comments were minor)

This revision is now accepted and ready to land.Aug 11 2022, 8:55 AM
mtrofin updated this revision to Diff 453756.Aug 18 2022, 12:36 PM
mtrofin marked 4 inline comments as done.

feedback

mtrofin added inline comments.Aug 18 2022, 2:24 PM
llvm/CMakeLists.txt
950–957

We used the "have" pattern before though. The idea is that if the dep is present we enable the code. I think "use" is used to say "enable this" (deps not an issue).

957

Yes. Keeping it for now until we sort out downstream (so a few days after this lands). After that, this and the old TFUtils.cpp get deleted.

llvm/include/llvm/Config/llvm-config.h.cmake
102

yes.

llvm/lib/Analysis/TFLiteUtils.cpp
104

adding a FIXME to do that - want to get the core functionality up and going with the minimal change (and this is large-ish already).

phosek added inline comments.Aug 23 2022, 11:58 PM
llvm/CMakeLists.txt
956

Could this be scoped more narrowly? That is, only add the include directory to the targets that actually need it with target_include_directories.

mtrofin marked an inline comment as done.Aug 24 2022, 1:26 PM
mtrofin added inline comments.
llvm/CMakeLists.txt
956

I tried, it seems quite complicated - and we want to subsequently get rid of the protobuf dependency anyway, so wonder if it's worth the trouble for the interim

mtrofin updated this revision to Diff 455396.Aug 24 2022, 3:14 PM
mtrofin marked an inline comment as done.

a fix for when building both release and development modes together

phosek accepted this revision.Aug 24 2022, 3:54 PM

LGTM