[MLGO] Use TFLite in 'development mode'

[mtrofin]

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!

[phosek]

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 {

[mtrofin]

actual patch

[mtrofin]

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]

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).

[ebrevdo]

(my comments were minor)

[mtrofin]

feedback

[mtrofin]

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]

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]

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]

a fix for when building both release and development modes together

[phosek]

LGTM

llvm/CMakeLists.txt

 set( CMAKE_LIBRARY_OUTPUT_DIRECTORY ${LLVM_LIBRARY_DIR} )
 set( CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${LLVM_LIBRARY_DIR} )
 
 if(LLVM_INCLUDE_TESTS)
   include(GetErrcMessages)
   get_errc_messages(LLVM_LIT_ERRC_MESSAGES)
 endif()
 
-# For up-to-date instructions for installing the Tensorflow dependency, refer to
+# For up-to-date instructions for installing the TFLite dependency, refer to
 # the bot setup script: https://github.com/google/ml-compiler-opt/blob/main/buildbot/buildbot_init.sh
-# In this case, the latest C API library is available for download from
-# https://www.tensorflow.org/install/lang_c.
-# We will expose the conditional compilation variable,
-# LLVM_HAVE_TF_API, through llvm-config.h, so that a user of the LLVM library may
-# also leverage the dependency.
 set(TENSORFLOW_C_LIB_PATH "" CACHE PATH "Path to TensorFlow C library install")
-if (TENSORFLOW_C_LIB_PATH)
+set(LLVM_HAVE_TFLITE "" CACHE BOOL "Use tflite")
+if (LLVM_HAVE_TFLITE)
+  find_package(protobuf REQUIRED)
+  find_package(tensorflow-lite REQUIRED)
+  set(LLVM_HAVE_TF_API "ON" CACHE BOOL "Full Tensorflow API available")
+  set(LLVM_PROTOBUF_OUT_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/protobuf_gen)
+  include_directories(${LLVM_PROTOBUF_OUT_DIR})
+elseif (TENSORFLOW_C_LIB_PATH)
   find_library(tensorflow_c_api tensorflow PATHS ${TENSORFLOW_C_LIB_PATH}/lib NO_DEFAULT_PATH REQUIRED)
   # Currently, the protobuf headers are distributed with the pip package that corresponds to the version
   # of the C API library.
   find_library(tensorflow_fx tensorflow_framework PATHS ${TENSORFLOW_C_LIB_PATH}/lib NO_DEFAULT_PATH REQUIRED)
   set(LLVM_HAVE_TF_API "ON" CACHE BOOL "Full Tensorflow API available")
   include_directories(${TENSORFLOW_C_LIB_PATH}/include)
   if (NOT TF_PROTO_HEADERS)
     message(STATUS "TF_PROTO_HEADERS not defined. Looking for tensorflow pip package.")

 #
 set(TENSORFLOW_AOT_PATH "" CACHE PATH "Path to TensorFlow pip install dir")
 
 if (NOT TENSORFLOW_AOT_PATH STREQUAL "")
   set(LLVM_HAVE_TF_AOT "ON" CACHE BOOL "Tensorflow AOT available")
   set(TENSORFLOW_AOT_COMPILER
     "${TENSORFLOW_AOT_PATH}/../../../../bin/saved_model_cli"
     CACHE PATH "Path to the Tensorflow AOT compiler")
-  include_directories(${TENSORFLOW_AOT_PATH}/include)
+  # This needs to happen to avoid clashing protobuf codegen when building both AOT and development mode.
+  # We plan to replace protobuf with a simpler alternative, so this will go away.
+  file(COPY ${TENSORFLOW_AOT_PATH}/include DESTINATION ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/tensorflow
+    PATTERN "google/*" EXCLUDE
+    PATTERN "*.pb.h" EXCLUDE) 
+  include_directories(${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/tensorflow/include)
   add_subdirectory(${TENSORFLOW_AOT_PATH}/xla_aot_runtime_src
     ${CMAKE_ARCHIVE_OUTPUT_DIRECTORY}/tf_runtime)
   install(TARGETS tf_xla_runtime EXPORT LLVMExports
     ARCHIVE DESTINATION lib${LLVM_LIBDIR_SUFFIX} COMPONENT tf_xla_runtime)
   set_property(GLOBAL APPEND PROPERTY LLVM_EXPORTS tf_xla_runtime)
   # Once we add more modules, we should handle this more automatically.
   if (DEFINED LLVM_OVERRIDE_MODEL_HEADER_INLINERSIZEMODEL)
     set(LLVM_INLINER_MODEL_PATH "none")

llvm/cmake/modules/TensorFlowCompile.cmake

else()
     tf_compile(${LLVM_ML_MODELS_ABSOLUTE} ${tag_set} ${signature_def_key} ${fname} ${cpp_class} ${hdr_file} ${obj_file})
   endif()
 
   set(GeneratedMLSources ${GeneratedMLSources} ${GENERATED_OBJS} ${GENERATED_HEADERS} PARENT_SCOPE)
   set(MLDeps ${MLDeps} tf_xla_runtime PARENT_SCOPE)
   set(MLLinkDeps ${MLLinkDeps} tf_xla_runtime PARENT_SCOPE)
   add_definitions(-DLLVM_HAVE_TF_AOT_${fname_allcaps})
 endfunction()
+
+function(build_proto)
+  foreach (P ${ARGV})
+    set(PB_SRCS ${PB_SRCS} ${LLVM_PROTOBUF_OUT_DIR}/${P}.pb.cc)
+    set(PB_HDRS ${PB_HDRS} ${LLVM_PROTOBUF_OUT_DIR}/${P}.pb.h)
+    set(PBS ${PBS} ${TENSORFLOW_SRC_DIR}/${P}.proto)
+  endforeach()
+  add_custom_command(OUTPUT ${PB_SRCS} ${PB_HDRS}
+    COMMAND protobuf::protoc 
+    ARGS --proto_path=${TENSORFLOW_SRC_DIR} --cpp_out=${LLVM_PROTOBUF_OUT_DIR} ${PBS})
+  set_source_files_properties(${PB_SRCS} PROPERTIES
+    GENERATED 1)
+  set(GeneratedMLSources ${GeneratedMLSources} ${PB_SRCS} PARENT_SCOPE)
+  set(MLDeps ${MLDeps} ${MLDeps} PARENT_SCOPE)
+endfunction()

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

 
 /* Define if zlib compression is available */
 #cmakedefine01 LLVM_ENABLE_ZLIB
 
 /* Define if zstd compression is available */
 #cmakedefine01 LLVM_ENABLE_ZSTD
 
 /* Define if LLVM was built with a dependency to the libtensorflow dynamic library */
 #cmakedefine LLVM_HAVE_TF_API
 
+/* Define if LLVM is using tflite instead of libtensorflow */
+#cmakedefine LLVM_HAVE_TFLITE
+
 /* Define to 1 if you have the <sysexits.h> header file. */
 #cmakedefine HAVE_SYSEXITS_H ${HAVE_SYSEXITS_H}
 
 /* Define if the xar_open() function is supported on this platform. */
 #cmakedefine LLVM_HAVE_LIBXAR ${LLVM_HAVE_LIBXAR}
 
 /* Define if building libLLVM shared library */
 #cmakedefine LLVM_BUILD_LLVM_DYLIB

llvm/lib/Analysis/CMakeLists.txt

tf_find_and_compile(
       serve
       action
       InlinerSizeModel
       llvm::InlinerSizeModel
     )
   endif()
 
   if (DEFINED LLVM_HAVE_TF_API)
+    if (DEFINED LLVM_HAVE_TFLITE)
+      build_proto(
+        tensorflow/core/protobuf/error_codes
+        tensorflow/core/example/feature
+        tensorflow/core/example/example)
+      list(APPEND MLLinkDeps 
+        tensorflow-lite::tensorflow-lite
+        protobuf::libprotobuf)
+    else()
       list(APPEND MLLinkDeps ${tensorflow_c_api} ${tensorflow_fx})
     endif()
   endif()
+endif()
 
 add_llvm_component_library(LLVMAnalysis
   AliasAnalysis.cpp
   AliasAnalysisEvaluator.cpp
   AliasAnalysisSummary.cpp
   AliasSetTracker.cpp
   Analysis.cpp
   AssumeBundleQueries.cpp
add_llvm_component_library(LLVMAnalysis
   ScalarEvolutionAliasAnalysis.cpp
   ScalarEvolutionDivision.cpp
   ScalarEvolutionNormalization.cpp
   StackLifetime.cpp
   StackSafetyAnalysis.cpp
   SyncDependenceAnalysis.cpp
   SyntheticCountsUtils.cpp
   TFUtils.cpp
+  TFLiteUtils.cpp
   TargetLibraryInfo.cpp
   TargetTransformInfo.cpp
   TensorSpec.cpp
   Trace.cpp
   TrainingLogger.cpp
   TypeBasedAliasAnalysis.cpp
   TypeMetadataUtils.cpp
   ScopedNoAliasAA.cpp

llvm/lib/Analysis/TFLiteUtils.cpp

+//===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements utilities for interfacing with tensorflow C APIs.
+//
+//===----------------------------------------------------------------------===//
+#include "llvm/Config/config.h"
+#if defined(LLVM_HAVE_TFLITE)
+
+#include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/Support/Base64.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/model.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/op_resolver.h"
+
+#include <cassert>
+#include <numeric>
+
+using namespace llvm;
+
+namespace llvm {
+class EvaluationResultImpl {
+public:
+  EvaluationResultImpl(const std::vector<const TfLiteTensor *> &Outputs)
+      : Outputs(Outputs){};
+
+  const TfLiteTensor *getOutput(size_t I) { return Outputs[I]; }
+
+  EvaluationResultImpl(const EvaluationResultImpl &) = delete;
+  EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
+
+private:
+  const std::vector<const TfLiteTensor *> Outputs;
+};
+
+class TFModelEvaluatorImpl {
+public:
+  TFModelEvaluatorImpl(StringRef SavedModelPath,
+                       const std::vector<TensorSpec> &InputSpecs,
+                       function_ref<TensorSpec(size_t)> GetOutputSpecs,
+                       size_t OutputSpecsSize, const char *Tags);
+
+  bool isValid() const { return IsValid; }
+  size_t outputSize() const { return Output.size(); }
+
+  std::unique_ptr<EvaluationResultImpl> evaluate() {
+    Interpreter->Invoke();
+    return std::make_unique<EvaluationResultImpl>(Output);
+  }
+
+  const std::vector<TfLiteTensor *> &getInput() const { return Input; }
+
+  ~TFModelEvaluatorImpl();
+
+private:
+  std::unique_ptr<tflite::FlatBufferModel> Model;
+
+  /// The objects necessary for carrying out an evaluation of the SavedModel.
+  /// They are expensive to set up, and we maintain them accross all the
+  /// evaluations of the model.
+  std::unique_ptr<tflite::Interpreter> Interpreter;
+
+  /// The input tensors. We set up the tensors once and just mutate theirs
+  /// scalars before each evaluation. The input tensors keep their value after
+  /// an evaluation.
+  std::vector<TfLiteTensor *> Input;
+
+  /// The output nodes.
+  std::vector<const TfLiteTensor *> Output;
+
+  void invalidate() { IsValid = false; }
+
+  bool IsValid = true;
+
+  /// Reusable utility for ensuring we can bind the requested Name to a node in
+  /// the SavedModel Graph.
+  bool checkReportAndInvalidate(const TfLiteTensor *Tensor,
+                                const TensorSpec &Spec);
+};
+
+} // namespace llvm
+
+TFModelEvaluatorImpl::TFModelEvaluatorImpl(
+    StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
+    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
+    const char *Tags = "serve")
+    : Input(InputSpecs.size()), Output(OutputSpecsSize) {
+  // FIXME: make ErrorReporter a member (may also need subclassing
+  // StatefulErrorReporter) to easily get the latest error status, for
+  // debugging.
+  tflite::StderrReporter ErrorReporter;
+  SmallVector<char, 128> TFLitePathBuff;
+  llvm::sys::path::append(TFLitePathBuff, SavedModelPath, "model.tflite");
+  StringRef TFLitePath(TFLitePathBuff.data(), TFLitePathBuff.size());
+  Model = tflite::FlatBufferModel::BuildFromFile(TFLitePath.str().c_str(),
+                                                 &ErrorReporter);
+  if (!Model) {
+    invalidate();
+    return;
+  }
+
+  tflite::ops::builtin::BuiltinOpResolver Resolver;
+  tflite::InterpreterBuilder Builder(*Model, Resolver);
+  Builder(&Interpreter);
+
+  if (!Interpreter ||
+      Interpreter->AllocateTensors() != TfLiteStatus::kTfLiteOk) {
+    invalidate();
+    return;
+  }
+  // Known inputs and outputs
+  StringMap<int> InputsMap;
+  StringMap<int> OutputsMap;
+  for (size_t I = 0; I < Interpreter->inputs().size(); ++I)
+    InputsMap[Interpreter->GetInputName(I)] = I;
+  for (size_t I = 0; I < Interpreter->outputs().size(); ++I)
+    OutputsMap[Interpreter->GetOutputName(I)] = I;
+
+  for (size_t I = 0; I < InputSpecs.size(); ++I) {
+    auto &InputSpec = InputSpecs[I];
+    auto MapI = InputsMap.find(InputSpec.name() + ":" +
+                               std::to_string(InputSpec.port()));
+    if (MapI == InputsMap.end()) {
+      Input[I] = nullptr;
+      continue;
+    }
+    Input[I] = Interpreter->tensor(MapI->second);
+    if (!checkReportAndInvalidate(Input[I], InputSpec))
+      return;
+    std::memset(Input[I]->data.data, 0,
+                InputSpecs[I].getTotalTensorBufferSize());
+  }
+
+  for (size_t I = 0; I < OutputSpecsSize; ++I) {
+    auto OutputSpec = GetOutputSpecs(I);
+    Output[I] = Interpreter->output_tensor(
+        OutputsMap[OutputSpec.name() + ":" +
+                   std::to_string(OutputSpec.port())]);
+    if (!checkReportAndInvalidate(Output[I], OutputSpec))
+      return;
+  }
+}
+
+TFModelEvaluator::TFModelEvaluator(
+    StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
+    function_ref<TensorSpec(size_t)> GetOutputSpecs, size_t OutputSpecsSize,
+    const char *Tags)
+    : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, GetOutputSpecs,
+                                    OutputSpecsSize, Tags)) {
+  if (!Impl->isValid())
+    Impl.reset();
+}
+
+TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
+                                   const std::vector<TensorSpec> &InputSpecs,
+                                   const std::vector<TensorSpec> &OutputSpecs,
+                                   const char *Tags)
+    : TFModelEvaluator(
+          SavedModelPath, InputSpecs, [&](size_t I) { return OutputSpecs[I]; },
+          OutputSpecs.size(), Tags) {}
+
+TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {}
+
+bool TFModelEvaluatorImpl::checkReportAndInvalidate(const TfLiteTensor *Tensor,
+                                                    const TensorSpec &Spec) {
+  if (!Tensor) {
+    errs() << "Could not find TF_Output named: " + Spec.name();
+    IsValid = false;
+  }
+  if (Spec.getTotalTensorBufferSize() != Tensor->bytes)
+    IsValid = false;
+
+  // If the total sizes match, there could still be a mismatch in the shape.
+  // We ignore that for now.
+
+  return IsValid;
+}
+
+Optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
+  if (!isValid())
+    return None;
+  return EvaluationResult(Impl->evaluate());
+}
+
+void *TFModelEvaluator::getUntypedInput(size_t Index) {
+  TfLiteTensor *T = Impl->getInput()[Index];
+  if (!T)
+    return nullptr;
+  return T->data.data;
+}
+
+TFModelEvaluator::EvaluationResult::EvaluationResult(
+    std::unique_ptr<EvaluationResultImpl> Impl)
+    : Impl(std::move(Impl)) {}
+
+TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
+    : Impl(std::move(Other.Impl)) {}
+
+TFModelEvaluator::EvaluationResult &
+TFModelEvaluator::EvaluationResult::operator=(EvaluationResult &&Other) {
+  Impl = std::move(Other.Impl);
+  return *this;
+}
+
+void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
+  return Impl->getOutput(Index)->data.data;
+}
+
+const void *
+TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
+  return Impl->getOutput(Index)->data.data;
+}
+
+TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
+TFModelEvaluator::~TFModelEvaluator() {}
+
+#endif // defined(LLVM_HAVE_TF_API)

llvm/lib/Analysis/TFUtils.cpp

 //===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for interfacing with tensorflow C APIs.
 //
 //===----------------------------------------------------------------------===//
 #include "llvm/Config/config.h"
-#if defined(LLVM_HAVE_TF_API)
+#if defined(LLVM_HAVE_TF_API) && !defined(LLVM_HAVE_TFLITE)
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/Utils/TFUtils.h"
 #include "llvm/Support/Base64.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"

llvm/lib/Analysis/models/saved-model-to-tflite.py

+"""Convert a saved model to tflite model.
+
+Usage: python3 saved-model-to-tflite.py <mlgo saved_model_dir> <tflite dest_dir>
+
+The <tflite dest_dir> will contain:
+  model.tflite: this is the converted saved model
+  output_spec.json: the output spec, copied from the saved_model dir.
+"""
+
+import tensorflow as tf
+import os
+import sys
+from tf_agents.policies import greedy_policy
+
+
+def main(argv):
+  assert len(argv) == 3
+  sm_dir = argv[1]
+  tfl_dir = argv[2]
+  tf.io.gfile.makedirs(tfl_dir)
+  tfl_path = os.path.join(tfl_dir, 'model.tflite')
+  converter = tf.lite.TFLiteConverter.from_saved_model(sm_dir)
+  converter.target_spec.supported_ops = [
+    tf.lite.OpsSet.TFLITE_BUILTINS,
+  ]
+  tfl_model = converter.convert()
+  with tf.io.gfile.GFile(tfl_path, 'wb') as f:
+    f.write(tfl_model)
+  
+  json_file = 'output_spec.json'
+  src_json = os.path.join(sm_dir, json_file)
+  if tf.io.gfile.exists(src_json):
+    tf.io.gfile.copy(src_json,
+                     os.path.join(tfl_dir, json_file))
+
+if __name__ == '__main__':
+  main(sys.argv)

llvm/test/CodeGen/MLRegalloc/dev-mode-log-2-fcts.ll

 ; REQUIRES: have_tf_api
 ; REQUIRES: x86_64-linux
 ;
 ; Check that we can log more than 1 function.
 ;
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
 ; RUN:   -regalloc-training-log=%t1 -tfutils-text-log < %s
 ; RUN: sed -i 's/ \+/ /g' %t1
 ; RUN: sed -i 's/\\n key:/\n key:/g' %t1
 ; RUN: sed -i 's/\\n feature/\n feature/g' %t1
 ; RUN: sed -i 's/\\n/ /g' %t1
 ; RUN: FileCheck --input-file %t1 %s
 
-; RUN: rm -rf %t && mkdir %t
-; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t
+; RUN: rm -rf %t %t_savedmodel
+; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t_savedmodel
+; RUN: %python %S/../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
 ; RUN:   -regalloc-training-log=%t2 -tfutils-text-log -regalloc-model=%t < %s
 ; RUN: sed -i 's/ \+/ /g' %t2
 ; RUN: sed -i 's/\\n key:/\n key:/g' %t2
 ; RUN: sed -i 's/\\n feature/\n feature/g' %t2
 ; RUN: sed -i 's/\\n/ /g' %t2
 ; RUN: FileCheck --input-file %t2 %s
 

llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll

 ; REQUIRES: have_tf_api
 ; REQUIRES: x86_64-linux
 ;
 ; Check that we log correctly, both with a learned policy, and the default policy
 ;
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
 ; RUN:   -regalloc-training-log=%t1 -tfutils-text-log < %S/Inputs/input.ll
 ; RUN: sed -i 's/ \+/ /g' %t1
 ; RUN: sed -i 's/\\n key:/\n key:/g' %t1
 ; RUN: sed -i 's/\\n feature/\n feature/g' %t1
 ; RUN: sed -i 's/\\n/ /g' %t1
 ; RUN: FileCheck --input-file %t1 %s --check-prefixes=CHECK,NOML
 ; RUN: diff %t1 %S/Inputs/reference-log-noml.txt
 
-; RUN: rm -rf %t && mkdir %t
-; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t
+; RUN: rm -rf %t_savedmodel %t
+; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t_savedmodel
+; RUN: %python %S/../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
 ; RUN:   -regalloc-training-log=%t2 -tfutils-text-log -regalloc-model=%t < %S/Inputs/input.ll
 ; RUN: sed -i 's/ \+/ /g' %t2
 ; RUN: sed -i 's/\\n key:/\n key:/g' %t2
 ; RUN: sed -i 's/\\n feature/\n feature/g' %t2
 ; RUN: sed -i 's/\\n/ /g' %t2
 ; RUN: FileCheck --input-file %t2 %s --check-prefixes=CHECK,ML
 

llvm/test/CodeGen/MLRegalloc/dev-rel-equivalence.ll

 ; REQUIRES: have_tf_api
 ; REQUIRES: llvm_raevict_model_autogenerated
 ; REQUIRES: x86_64-linux
 ;
 ; Check that the same model (==the autogenerated one) produces the same output
 ; regardless of how it's evaluated, which is different from the default
 ;
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=default \
 ; RUN:   %S/Inputs/input.ll -o %t.default
 
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=release \
 ; RUN:   %S/Inputs/input.ll -o %t.release
 
-; RUN: rm -rf %t && mkdir %t
+; RUN: rm -rf %t %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t
 ; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
 ; RUN:   -regalloc-model=%t %S/Inputs/input.ll -o %t.development
 ; RUN: diff %t.release %t.development
 ; RUN: not diff %t.release %t.default

llvm/test/Transforms/Inline/ML/bounds-checks-rewards.ll

 ; Test behavior when inlining policy grows size out of control.
 ; In all cases, the end result is the same: mandatory inlinings must happen.
 ; Also in all cases, we don't record the mandatory inlining (there's nothing to
 ; learn from it).
 ; However, when we discover we 'trip' over the artificially-low size increase 
 ; factor, we penalize the 'bad' decision.
 ; REQUIRES: have_tf_api
 ;
 ; Generate mock model
 ; RUN: rm -rf %t
-; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t
+; RUN: rm -rf %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ;
 ; When the bounds are very wide ("no bounds"), all inlinings happen.
 ; RUN: opt -passes=scc-oz-module-inliner -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-model-under-training=%t -training-log=- -tfutils-text-log -enable-ml-inliner=development -ml-advisor-size-increase-threshold=10.0 -S < %s 2>&1 | FileCheck %s --check-prefix=CHECK --check-prefix=NOBOUNDS
 ;
 ; When the bounds are very restrictive, the first inlining happens but it's
 ; considered as "bad" (since it trips over the bounds) and its reward is a
 ; penalty. However, the mandatory inlining, which is considered next, happens.
 ; No other inlinings happend.

llvm/test/Transforms/Inline/ML/development-training-log.ll

 ; Test that we can produce a log if we have or do not have a model, in development mode.
 ; REQUIRES: have_tf_api
 ; Generate mock model
-; RUN: rm -rf %t
-; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t
+; RUN: rm -rf %t_savedmodel %t
+; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ;
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s 
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -ml-inliner-output-spec-override=%S/Inputs/test_output_spec.json -S < %s | FileCheck %s --check-prefixes=EXTRA-OUTPUTS,CHECK
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-ir2native-model=%S/../../../../unittests/Analysis/Inputs/ir2native_x86_64_model -S < %s | FileCheck %s
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -ml-inliner-model-under-training=%t -S < %s | FileCheck %s --check-prefix=NOREWARD
 ; RUN: opt -enable-ml-inliner=development -passes=scc-oz-module-inliner -training-log=- -tfutils-text-log -S < %s | FileCheck %s --check-prefix=NOREWARD
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-pc-linux-gnu"

llvm/test/Transforms/Inline/ML/ml-test-development-mode.ll

 ; The default inliner doesn't elide @adder, it believes it's too costly to inline 
 ; adder into switcher. The ML inliner carries out that inlining, resulting in
 ; a smaller result (part of it is that adder gets elided).
 ;
 ; This test uses Inputs/test-module.ll, as it shares it with a similar test
 ; for the 'release' mode.
 ;
 ; REQUIRES: have_tf_api
-; RUN: rm -rf %t && mkdir %t
-; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t
+; RUN: rm -rf %t
+; RUN: rm -rf %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/gen-inline-oz-test-model.py %t_savedmodel
+; RUN: %python %S/../../../../lib/Analysis/models/saved-model-to-tflite.py %t_savedmodel %t
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=default -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=DEFAULT
 ; RUN: opt -passes=scc-oz-module-inliner -enable-ml-inliner=development -ml-inliner-model-under-training=%t -S < %S/Inputs/test-module.ll 2>&1 | FileCheck %S/Inputs/test-module.ll --check-prefix=CHECK

llvm/unittests/Analysis/TFUtilsTest.cpp

TEST(TFUtilsTest, EvalError) {
   // dimension (1, 214)
   const static int64_t KnownSize = 213;
   std::vector<TensorSpec> InputSpecs{TensorSpec::createSpec<int32_t>(
       "serving_default_input_1", {1, KnownSize})};
   std::vector<TensorSpec> OutputSpecs{
       TensorSpec::createSpec<float>("StatefulPartitionedCall", {1})};
 
   TFModelEvaluator Evaluator(getModelPath(), InputSpecs, OutputSpecs);
-  EXPECT_TRUE(Evaluator.isValid());
-
-  int32_t *V = Evaluator.getInput<int32_t>(0);
-  // Fill it up with 1's, we know the output.
-  for (auto I = 0; I < KnownSize; ++I) {
-    V[I] = 1;
-  }
-  auto ER = Evaluator.evaluate();
-  EXPECT_FALSE(ER.hasValue());
   EXPECT_FALSE(Evaluator.isValid());
 }
 
 TEST(TFUtilsTest, UnsupportedFeature) {
   const static int64_t KnownSize = 214;
   std::vector<TensorSpec> InputSpecs{
       TensorSpec::createSpec<int32_t>("serving_default_input_1",
                                       {1, KnownSize}),

utils/bazel/llvm_configs/llvm-config.h.cmake

 #cmakedefine01 LLVM_ENABLE_ZLIB
 
 /* Define if zstd compression is available */
 #cmakedefine01 LLVM_ENABLE_ZSTD
 
 /* Define if LLVM was built with a dependency to the libtensorflow dynamic library */
 #cmakedefine LLVM_HAVE_TF_API
 
+/* Define if LLVM is using tflite instead of libtensorflow */
+#cmakedefine LLVM_HAVE_TFLITE
+
 /* Define to 1 if you have the <sysexits.h> header file. */
 #cmakedefine HAVE_SYSEXITS_H ${HAVE_SYSEXITS_H}
 
 /* Define if the xar_open() function is supported on this platform. */
 #cmakedefine LLVM_HAVE_LIBXAR ${LLVM_HAVE_LIBXAR}
 
 /* Define if building libLLVM shared library */
 #cmakedefine LLVM_BUILD_LLVM_DYLIB