llvm/lib/Analysis/TFLiteUtils.cpp

   1 //===- TFUtils.cpp - TFLite-based evaluation utilities --------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements utilities for interfacing with TFLite.
  10 //
  11 //===----------------------------------------------------------------------===//
  12 #include "llvm/Config/config.h"
  13 #if defined(LLVM_HAVE_TFLITE)
  14
  15 #include "llvm/ADT/Twine.h"
  16 #include "llvm/Analysis/Utils/TFUtils.h"
  17 #include "llvm/Support/Base64.h"
  18 #include "llvm/Support/CommandLine.h"
  19 #include "llvm/Support/Debug.h"
  20 #include "llvm/Support/JSON.h"
  21 #include "llvm/Support/MemoryBuffer.h"
  22 #include "llvm/Support/Path.h"
  23 #include "llvm/Support/raw_ostream.h"
  24
  25 #include "tensorflow/lite/interpreter.h"
  26 #include "tensorflow/lite/kernels/register.h"
  27 #include "tensorflow/lite/model.h"
  28 #include "tensorflow/lite/model_builder.h"
  29 #include "tensorflow/lite/op_resolver.h"
  30 #include "tensorflow/lite/logger.h"
  31
  32 #include <cassert>
  33 #include <numeric>
  34 #include <optional>
  35
  36 using namespace llvm;
  37
  38 namespace llvm {
  39 class EvaluationResultImpl {
  40 public:
  41   EvaluationResultImpl(const std::vector<const TfLiteTensor *> &Outputs)
  42       : Outputs(Outputs){};
  43
  44   const TfLiteTensor *getOutput(size_t I) { return Outputs[I]; }
  45
  46   EvaluationResultImpl(const EvaluationResultImpl &) = delete;
  47   EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
  48
  49 private:
  50   const std::vector<const TfLiteTensor *> Outputs;
  51 };
  52
  53 class TFModelEvaluatorImpl {
  54 public:
  55   TFModelEvaluatorImpl(StringRef SavedModelPath,
  56                        const std::vector<TensorSpec> &InputSpecs,
  57                        const std::vector<TensorSpec> &OutputSpecs,
  58                        const char *Tags);
  59
  60   bool isValid() const { return IsValid; }
  61   size_t outputSize() const { return Output.size(); }
  62
  63   std::unique_ptr<EvaluationResultImpl> evaluate() {
  64     Interpreter->Invoke();
  65     return std::make_unique<EvaluationResultImpl>(Output);
  66   }
  67
  68   const std::vector<TfLiteTensor *> &getInput() const { return Input; }
  69
  70   ~TFModelEvaluatorImpl();
  71
  72 private:
  73   std::unique_ptr<tflite::FlatBufferModel> Model;
  74
  75   /// The objects necessary for carrying out an evaluation of the SavedModel.
  76   /// They are expensive to set up, and we maintain them accross all the
  77   /// evaluations of the model.
  78   std::unique_ptr<tflite::Interpreter> Interpreter;
  79
  80   /// The input tensors. We set up the tensors once and just mutate theirs
  81   /// scalars before each evaluation. The input tensors keep their value after
  82   /// an evaluation.
  83   std::vector<TfLiteTensor *> Input;
  84
  85   /// The output nodes.
  86   std::vector<const TfLiteTensor *> Output;
  87
  88   void invalidate() { IsValid = false; }
  89
  90   bool IsValid = true;
  91
  92   /// Reusable utility for ensuring we can bind the requested Name to a node in
  93   /// the SavedModel Graph.
  94   bool checkReportAndInvalidate(const TfLiteTensor *Tensor,
  95                                 const TensorSpec &Spec);
  96 };
  97
  98 } // namespace llvm
  99
 100 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
 101     StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
 102     const std::vector<TensorSpec> &OutputSpecs, const char *Tags = "serve")
 103     : Input(InputSpecs.size()), Output(OutputSpecs.size()) {
 104   // INFO and DEBUG messages could be numerous and not particularly interesting
 105   tflite::LoggerOptions::SetMinimumLogSeverity(tflite::TFLITE_LOG_WARNING);
 106   // FIXME: make ErrorReporter a member (may also need subclassing
 107   // StatefulErrorReporter) to easily get the latest error status, for
 108   // debugging.
 109   tflite::StderrReporter ErrorReporter;
 110   SmallVector<char, 128> TFLitePathBuff;
 111   llvm::sys::path::append(TFLitePathBuff, SavedModelPath, "model.tflite");
 112   StringRef TFLitePath(TFLitePathBuff.data(), TFLitePathBuff.size());
 113   Model = tflite::FlatBufferModel::BuildFromFile(TFLitePath.str().c_str(),
 114                                                  &ErrorReporter);
 115   if (!Model) {
 116     invalidate();
 117     return;
 118   }
 119
 120   tflite::ops::builtin::BuiltinOpResolver Resolver;
 121   tflite::InterpreterBuilder Builder(*Model, Resolver);
 122   Builder(&Interpreter);
 123
 124   if (!Interpreter) {
 125     invalidate();
 126     return;
 127   }
 128
 129   // We assume the input buffers are valid for the lifetime of the interpreter.
 130   // By default, tflite allocates memory in an arena and will periodically take
 131   // away memory and reallocate it in a different location after evaluations in
 132   // order to improve utilization of the buffers owned in the arena. So, we
 133   // explicitly mark our input buffers as persistent to avoid this behavior.
 134   for (size_t I = 0; I < Interpreter->inputs().size(); ++I)
 135     Interpreter->tensor(I)->allocation_type =
 136         TfLiteAllocationType::kTfLiteArenaRwPersistent;
 137
 138   if (Interpreter->AllocateTensors() != TfLiteStatus::kTfLiteOk) {
 139     invalidate();
 140     return;
 141   }
 142   // Known inputs and outputs
 143   StringMap<int> InputsMap;
 144   StringMap<int> OutputsMap;
 145   for (size_t I = 0; I < Interpreter->inputs().size(); ++I)
 146     InputsMap[Interpreter->GetInputName(I)] = I;
 147   for (size_t I = 0; I < Interpreter->outputs().size(); ++I)
 148     OutputsMap[Interpreter->GetOutputName(I)] = I;
 149
 150   size_t NumberFeaturesPassed = 0;
 151   for (size_t I = 0; I < InputSpecs.size(); ++I) {
 152     auto &InputSpec = InputSpecs[I];
 153     auto MapI = InputsMap.find(InputSpec.name() + ":" +
 154                                std::to_string(InputSpec.port()));
 155     if (MapI == InputsMap.end()) {
 156       Input[I] = nullptr;
 157       continue;
 158     }
 159     Input[I] = Interpreter->tensor(MapI->second);
 160     if (!checkReportAndInvalidate(Input[I], InputSpec))
 161       return;
 162     std::memset(Input[I]->data.data, 0,
 163                 InputSpecs[I].getTotalTensorBufferSize());
 164     ++NumberFeaturesPassed;
 165   }
 166
 167   if (NumberFeaturesPassed < Interpreter->inputs().size()) {
 168     // we haven't passed all the required features to the model, throw an error.
 169     errs() << "Required feature(s) have not been passed to the ML model";
 170     invalidate();
 171     return;
 172   }
 173
 174   for (size_t I = 0; I < OutputSpecs.size(); ++I) {
 175     const auto &OutputSpec = OutputSpecs[I];
 176     Output[I] = Interpreter->output_tensor(
 177         OutputsMap[OutputSpec.name() + ":" +
 178                    std::to_string(OutputSpec.port())]);
 179     if (!checkReportAndInvalidate(Output[I], OutputSpec))
 180       return;
 181   }
 182 }
 183
 184 TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
 185                                    const std::vector<TensorSpec> &InputSpecs,
 186                                    const std::vector<TensorSpec> &OutputSpecs,
 187                                    const char *Tags)
 188     : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs,
 189                                     Tags)) {
 190   if (!Impl->isValid())
 191     Impl.reset();
 192 }
 193
 194 TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {}
 195
 196 bool TFModelEvaluatorImpl::checkReportAndInvalidate(const TfLiteTensor *Tensor,
 197                                                     const TensorSpec &Spec) {
 198   if (!Tensor) {
 199     errs() << "Could not find TF_Output named: " + Spec.name();
 200     IsValid = false;
 201   }
 202   if (Spec.getTotalTensorBufferSize() != Tensor->bytes)
 203     IsValid = false;
 204
 205   // If the total sizes match, there could still be a mismatch in the shape.
 206   // We ignore that for now.
 207
 208   return IsValid;
 209 }
 210
 211 std::optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
 212   if (!isValid())
 213     return std::nullopt;
 214   return EvaluationResult(Impl->evaluate());
 215 }
 216
 217 void *TFModelEvaluator::getUntypedInput(size_t Index) {
 218   TfLiteTensor *T = Impl->getInput()[Index];
 219   if (!T)
 220     return nullptr;
 221   return T->data.data;
 222 }
 223
 224 TFModelEvaluator::EvaluationResult::EvaluationResult(
 225     std::unique_ptr<EvaluationResultImpl> Impl)
 226     : Impl(std::move(Impl)) {}
 227
 228 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
 229     : Impl(std::move(Other.Impl)) {}
 230
 231 TFModelEvaluator::EvaluationResult &
 232 TFModelEvaluator::EvaluationResult::operator=(EvaluationResult &&Other) {
 233   Impl = std::move(Other.Impl);
 234   return *this;
 235 }
 236
 237 void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
 238   return Impl->getOutput(Index)->data.data;
 239 }
 240
 241 const void *
 242 TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) const {
 243   return Impl->getOutput(Index)->data.data;
 244 }
 245
 246 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
 247 TFModelEvaluator::~TFModelEvaluator() {}
 248
 249 #endif // defined(LLVM_HAVE_TFLITE)