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 
14 #include "llvm/Analysis/Utils/TFUtils.h"
15 #include "llvm/ADT/Twine.h"
16 #include "llvm/Support/Debug.h"
17 #include "llvm/Support/ManagedStatic.h"
18 #include "llvm/Support/raw_ostream.h"
19 
20 #include "tensorflow/c/c_api_experimental.h"
21 
22 #include <cassert>
23 
24 using namespace llvm;
25 
26 namespace {
27 
28 struct TFInitializer {
29   TFInitializer() {
30     assert(!IsInitialized && "TFInitialized should be called only once");
31     int Argc = 1;
32     const char *Name = "";
33     const char **NamePtr = &Name;
34     TF_InitMain(Name, &Argc, const_cast<char ***>(&NamePtr));
35     IsInitialized = true;
36   }
37   bool IsInitialized = false;
38 };
39 
40 llvm::ManagedStatic<TFInitializer> TFLibInitializer;
41 
42 bool ensureInitTF() { return TFLibInitializer->IsInitialized; }
43 
44 TFModelEvaluator::TFGraphPtr createTFGraph() {
45   return TFModelEvaluator::TFGraphPtr(TF_NewGraph(), &TF_DeleteGraph);
46 }
47 
48 TFModelEvaluator::TFStatusPtr createTFStatus() {
49   return TFModelEvaluator::TFStatusPtr(TF_NewStatus(), &TF_DeleteStatus);
50 }
51 
52 TFModelEvaluator::TFSessionOptionsPtr createTFSessionOptions() {
53   return TFModelEvaluator::TFSessionOptionsPtr(TF_NewSessionOptions(),
54                                                &TF_DeleteSessionOptions);
55 }
56 } // namespace
57 
58 TFModelEvaluator::TFModelEvaluator(StringRef SavedModelPath,
59                                    const std::vector<std::string> &InputNames,
60                                    const std::vector<std::string> &OutputNames,
61                                    const char *Tags)
62     : Graph(createTFGraph()), Options(createTFSessionOptions()),
63       InputFeed(InputNames.size()), Input(InputNames.size()),
64       OutputFeed(OutputNames.size()) {
65   if (!ensureInitTF()) {
66     errs() << "Tensorflow should have been initialized";
67     return;
68   }
69   auto Status = createTFStatus();
70 
71   Session = TF_LoadSessionFromSavedModel(Options.get(), nullptr,
72                                          SavedModelPath.str().c_str(), &Tags, 1,
73                                          Graph.get(), nullptr, Status.get());
74   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
75     errs() << TF_Message(Status.get());
76     deleteSession();
77   }
78   for (size_t I = 0; I < InputNames.size(); ++I) {
79     InputFeed[I] = {
80         TF_GraphOperationByName(Graph.get(), (InputNames[I]).c_str()), 0};
81     if (!checkReportAndReset(InputFeed[I], InputNames[I]))
82       return;
83   }
84   for (size_t I = 0; I < OutputNames.size(); ++I) {
85     OutputFeed[I] = {
86         TF_GraphOperationByName(Graph.get(), (OutputNames[I]).c_str()), 0};
87     if (!checkReportAndReset(OutputFeed[I], OutputNames[I]))
88       return;
89   }
90 }
91 
92 TFModelEvaluator::~TFModelEvaluator() {
93   for (auto *T : Input) {
94     TF_DeleteTensor(T);
95   }
96   deleteSession();
97 }
98 
99 bool TFModelEvaluator::checkReportAndReset(const TF_Output &Output,
100                                            StringRef Name) {
101   if (Output.oper)
102     return true;
103   errs() << "Could not find TF_Output named: " + Name;
104   deleteSession();
105   return false;
106 }
107 
108 void TFModelEvaluator::deleteSession() {
109   if (Session == nullptr)
110     return;
111   auto Status = createTFStatus();
112   TF_DeleteSession(Session, Status.get());
113   Session = nullptr;
114   if (TF_GetCode(Status.get()) != TF_Code::TF_OK)
115     errs() << "Could not delete TF session";
116 }
117 
118 Optional<TFModelEvaluator::EvaluationResult> TFModelEvaluator::evaluate() {
119   if (!isValid())
120     return None;
121   EvaluationResult Ret(OutputFeed.size());
122   auto Status = createTFStatus();
123   TF_SessionRun(Session, nullptr, InputFeed.data(), Input.data(), Input.size(),
124                 OutputFeed.data(), Ret.Output.data(), Ret.Output.size(),
125                 nullptr, 0, nullptr, Status.get());
126   if (TF_GetCode(Status.get()) != TF_Code::TF_OK) {
127     errs() << TF_Message(Status.get());
128     deleteSession();
129     return None;
130   }
131   return Ret;
132 }
133 
134 void TFModelEvaluator::initInput(int Index, TF_DataType Type,
135                                  const std::vector<int64_t> &Dimensions) {
136   int64_t TotalSize = TF_DataTypeSize(Type);
137   for (auto &D : Dimensions)
138     TotalSize *= D;
139 
140   Input[Index] =
141       TF_AllocateTensor(Type, Dimensions.data(), Dimensions.size(), TotalSize);
142   std::memset(TF_TensorData(Input[Index]), 0, TotalSize);
143 }