1 //===- TFUtils.cpp - tensorflow 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 tensorflow C APIs. 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) 250