1 //===- TFUtils.cpp - tensorflow evaluation utilities ----------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file implements utilities for interfacing with tensorflow C APIs.
11 //
12 //===----------------------------------------------------------------------===//
13 #include "llvm/Config/config.h"
14 #if defined(LLVM_HAVE_TF_API)
15 
16 #include "llvm/Analysis/Utils/TFUtils.h"
17 #include "llvm/ADT/Twine.h"
18 #include "llvm/Support/Debug.h"
19 #include "llvm/Support/ManagedStatic.h"
20 #include "llvm/Support/raw_ostream.h"
21 
22 #include "tensorflow/c/c_api.h"
23 #include "tensorflow/c/c_api_experimental.h"
24 
25 #include <cassert>
26 
27 using namespace llvm;
28 
29 namespace {
30 
31 using TFGraphPtr = std::unique_ptr<TF_Graph, decltype(&TF_DeleteGraph)>;
32 using TFSessionOptionsPtr =
33     std::unique_ptr<TF_SessionOptions, decltype(&TF_DeleteSessionOptions)>;
34 using TFStatusPtr = std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)>;
35 
36 struct TFInitializer {
37   TFInitializer() {
38     assert(!IsInitialized && "TFInitialized should be called only once");
39     int Argc = 1;
40     const char *Name = "";
41     const char **NamePtr = &Name;
42     TF_InitMain(Name, &Argc, const_cast<char ***>(&NamePtr));
43     IsInitialized = true;
44   }
45   bool IsInitialized = false;
46 };
47 
48 llvm::ManagedStatic<TFInitializer> TFLibInitializer;
49 
50 bool ensureInitTF() { return TFLibInitializer->IsInitialized; }
51 
52 TFGraphPtr createTFGraph() {
53   return TFGraphPtr(TF_NewGraph(), &TF_DeleteGraph);
54 }
55 
56 TFStatusPtr createTFStatus() {
57   return TFStatusPtr(TF_NewStatus(), &TF_DeleteStatus);
58 }
59 
60 TFSessionOptionsPtr createTFSessionOptions() {
61   return TFSessionOptionsPtr(TF_NewSessionOptions(), &TF_DeleteSessionOptions);
62 }
63 } // namespace
64 
65 namespace llvm {
66 class EvaluationResultImpl {
67 public:
68   EvaluationResultImpl(size_t OutputSize)
69       : OutputSize(OutputSize), Output(OutputSize){};
70 
71   ~EvaluationResultImpl() {
72     for (auto *P : Output)
73       if (P)
74         TF_DeleteTensor(P);
75   }
76 
77   EvaluationResultImpl(const EvaluationResultImpl &) = delete;
78   EvaluationResultImpl(EvaluationResultImpl &&Other) = delete;
79   std::vector<TF_Tensor *> &getOutput() { return Output; }
80 
81 private:
82   const size_t OutputSize;
83   std::vector<TF_Tensor *> Output;
84 };
85 
86 class TFModelEvaluatorImpl {
87 public:
88   TFModelEvaluatorImpl(StringRef SavedModelPath,
89                        const std::vector<TensorSpec> &InputSpecs,
90                        const std::vector<TensorSpec> &OutputSpecs,
91                        const char *Tags);
92 
93   bool isValid() const { return IsValid; }
94   size_t OutputSize() const { return OutputFeed.size(); }
95 
96   void evaluate(TF_Tensor **Output, TF_Status *Status) {
97     TF_SessionRun(Session, nullptr, InputFeed.data(), Input.data(),
98                   Input.size(), OutputFeed.data(), Output, OutputFeed.size(),
99                   nullptr, 0, nullptr, Status);
100   }
101 
102   void initInput(size_t Index, TF_DataType Type,
103                  const std::vector<int64_t> &Dimensions);
104   const std::vector<TF_Tensor *> &getInput() const { return Input; }
105 
106   ~TFModelEvaluatorImpl();
107 
108 private:
109   /// The objects necessary for carrying out an evaluation of the SavedModel.
110   /// They are expensive to set up, and we maintain them accross all the
111   /// evaluations of the model.
112   TF_Session *Session = nullptr;
113   TFGraphPtr Graph;
114   TFSessionOptionsPtr Options;
115 
116   /// The specification of the input nodes.
117   std::vector<TF_Output> InputFeed;
118 
119   /// The input tensors. They must match by index of the corresponding InputFeed
120   /// value. We set up the tensors once and just mutate theirs scalars before
121   /// each evaluation. The input tensors keep their value after an evaluation.
122   std::vector<TF_Tensor *> Input;
123 
124   /// The specification of the output nodes. When evaluating, the tensors in the
125   /// output tensor vector must match by index the corresponding element in the
126   /// OutputFeed.
127   std::vector<TF_Output> OutputFeed;
128 
129   void invalidate() { IsValid = false; }
130 
131   bool IsValid = true;
132 
133   /// Reusable utility for ensuring we can bind the requested Name to a node in
134   /// the SavedModel Graph.
135   bool checkReportAndInvalidate(const TF_Output &Output,
136                                 const TensorSpec &OutputSpec);
137 };
138 } // namespace llvm
139 
140 TFModelEvaluatorImpl::TFModelEvaluatorImpl(
141     StringRef SavedModelPath, const std::vector<TensorSpec> &InputSpecs,
142     const std::vector<TensorSpec> &OutputSpecs, const char *Tags)
143     : Graph(createTFGraph()), Options(createTFSessionOptions()),
144       InputFeed(InputSpecs.size()), Input(InputSpecs.size()),
145       OutputFeed(OutputSpecs.size()) {
146   if (!ensureInitTF()) {
147     errs() << "Tensorflow should have been initialized";
148     return;
149   }
150   auto Status = createTFStatus();
151 
152   Session = TF_LoadSessionFromSavedModel(Options.get(), nullptr,
153                                          SavedModelPath.str().c_str(), &Tags, 1,
154                                          Graph.get(), nullptr, Status.get());
155   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
156     errs() << TF_Message(Status.get());
157     invalidate();
158   }
159   for (size_t I = 0; I < InputSpecs.size(); ++I) {
160     auto &InputSpec = InputSpecs[I];
161     InputFeed[I] = {
162         TF_GraphOperationByName(Graph.get(), (InputSpec.name()).c_str()),
163         InputSpec.port()};
164     if (!checkReportAndInvalidate(InputFeed[I], InputSpec))
165       return;
166     initInput(I, static_cast<TF_DataType>(InputSpec.typeIndex()),
167               InputSpec.shape());
168   }
169   for (size_t I = 0; I < OutputSpecs.size(); ++I) {
170     auto &OutputSpec = OutputSpecs[I];
171     OutputFeed[I] = {
172         TF_GraphOperationByName(Graph.get(), (OutputSpec.name()).c_str()),
173         OutputSpec.port()};
174     if (!checkReportAndInvalidate(OutputFeed[I], OutputSpec))
175       return;
176   }
177 }
178 
179 TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
180                                    const std::vector<TensorSpec> &InputSpecs,
181                                    const std::vector<TensorSpec> &OutputSpecs,
182                                    const char *Tags)
183     : Impl(new TFModelEvaluatorImpl(SavedModelPath, InputSpecs, OutputSpecs,
184                                     Tags)) {
185   if (!Impl->isValid())
186     Impl.reset();
187 }
188 
189 TFModelEvaluatorImpl::~TFModelEvaluatorImpl() {
190   for (auto *T : Input) {
191     TF_DeleteTensor(T);
192   }
193   if (Session == nullptr)
194     return;
195   auto Status = createTFStatus();
196   TF_DeleteSession(Session, Status.get());
197   Session = nullptr;
198   if (TF_GetCode(Status.get()) != TF_Code::TF_OK)
199     errs() << "Could not delete TF session";
200 }
201 
202 bool TFModelEvaluatorImpl::checkReportAndInvalidate(
203     const TF_Output &Output, const TensorSpec &OutputSpec) {
204   if (Output.oper)
205     return true;
206   errs() << "Could not find TF_Output named: " + OutputSpec.name();
207   IsValid = false;
208   return IsValid;
209 }
210 
211 Optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
212   if (!isValid())
213     return None;
214   std::unique_ptr<EvaluationResultImpl> Ret =
215       std::make_unique<EvaluationResultImpl>(Impl->OutputSize());
216   auto Status = createTFStatus();
217   Impl->evaluate(Ret->getOutput().data(), Status.get());
218   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
219     errs() << TF_Message(Status.get());
220     Impl.reset();
221     return None;
222   }
223   return EvaluationResult(std::move(Ret));
224 }
225 
226 void TFModelEvaluatorImpl::initInput(size_t Index, TF_DataType Type,
227                                      const std::vector<int64_t> &Dimensions) {
228   int64_t TotalSize = TF_DataTypeSize(Type);
229   for (auto &D : Dimensions)
230     TotalSize *= D;
231 
232   Input[Index] =
233       TF_AllocateTensor(Type, Dimensions.data(), Dimensions.size(), TotalSize);
234   std::memset(TF_TensorData(Input[Index]), 0, TotalSize);
235 }
236 
237 void *TFModelEvaluator::getUntypedInput(size_t Index) {
238   return TF_TensorData(Impl->getInput()[Index]);
239 }
240 
241 TFModelEvaluator::EvaluationResult::EvaluationResult(
242     std::unique_ptr<EvaluationResultImpl> Impl)
243     : Impl(std::move(Impl)) {}
244 
245 TFModelEvaluator::EvaluationResult::EvaluationResult(EvaluationResult &&Other)
246     : Impl(std::move(Other.Impl)) {}
247 
248 void *TFModelEvaluator::EvaluationResult::getUntypedTensorValue(size_t Index) {
249   return TF_TensorData(Impl->getOutput()[Index]);
250 }
251 
252 template <> int TensorSpec::getDataType<float>() { return TF_FLOAT; }
253 
254 template <> int TensorSpec::getDataType<double>() { return TF_DOUBLE; }
255 
256 template <> int TensorSpec::getDataType<int8_t>() { return TF_INT8; }
257 
258 template <> int TensorSpec::getDataType<uint8_t>() { return TF_UINT8; }
259 
260 template <> int TensorSpec::getDataType<int16_t>() { return TF_INT16; }
261 
262 template <> int TensorSpec::getDataType<uint16_t>() { return TF_UINT16; }
263 
264 template <> int TensorSpec::getDataType<int32_t>() { return TF_INT32; }
265 
266 template <> int TensorSpec::getDataType<uint32_t>() { return TF_UINT32; }
267 
268 template <> int TensorSpec::getDataType<int64_t>() { return TF_INT64; }
269 
270 template <> int TensorSpec::getDataType<uint64_t>() { return TF_UINT64; }
271 
272 TFModelEvaluator::EvaluationResult::~EvaluationResult() {}
273 TFModelEvaluator::~TFModelEvaluator() {}
274 #endif // defined(LLVM_HAVE_TF_API)
275