LLVM  14.0.0git
TFUtils.h
Go to the documentation of this file.
1 //===- TFUtils.h - utilities for tensorflow C API ---------------*- C++ -*-===//
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 #ifndef LLVM_ANALYSIS_UTILS_TFUTILS_H
10 #define LLVM_ANALYSIS_UTILS_TFUTILS_H
11 
12 #include "llvm/Config/llvm-config.h"
13 
14 #ifdef LLVM_HAVE_TF_API
15 #include "llvm/ADT/StringMap.h"
16 #include "llvm/IR/LLVMContext.h"
17 #include "llvm/Support/JSON.h"
18 
19 #include <memory>
20 #include <vector>
21 
22 namespace llvm {
23 
24 /// Load a SavedModel, find the given inputs and outputs, and setup storage
25 /// for input tensors. The user is responsible for correctly dimensioning the
26 /// input tensors and setting their values before calling evaluate().
27 /// To initialize:
28 /// - construct the object
29 /// - initialize the input tensors using initInput. Indices must correspond to
30 /// indices in the InputNames used at construction.
31 /// To use:
32 /// - set input values by using getInput to get each input tensor, and then
33 /// setting internal scalars, for all dimensions (tensors are row-major:
34 /// https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205)
35 /// - call evaluate. The input tensors' values are not consumed after this, and
36 /// may still be read.
37 /// - use the outputs in the output vector
38 class TFModelEvaluatorImpl;
39 class EvaluationResultImpl;
40 
41 /// TensorSpec encapsulates the specification of a tensor: its dimensions, or
42 /// "shape" (row-major), its type (see TensorSpec::getDataType specializations
43 /// for supported types), its name and port (see "TensorFlow: Large-Scale
44 /// Machine Learning on Heterogeneous Distributed Systems", section 4.2, para 2:
45 /// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45166.pdf)
46 ///
47 /// TensorSpec is used to set up a TFModelEvaluator by describing the expected
48 /// inputs and outputs.
49 class TensorSpec final {
50 public:
51  template <typename T>
52  static TensorSpec createSpec(const std::string &Name,
53  const std::vector<int64_t> &Shape,
54  int Port = 0) {
55  return TensorSpec(Name, Port, getDataType<T>(), Shape);
56  }
57 
58  const std::string &name() const { return Name; }
59  int port() const { return Port; }
60  int typeIndex() const { return TypeIndex; }
61  const std::vector<int64_t> &shape() const { return Shape; }
62 
63  bool operator==(const TensorSpec &Other) const {
64  return Name == Other.Name && Port == Other.Port &&
65  TypeIndex == Other.TypeIndex && Shape == Other.Shape;
66  }
67 
68  bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
69 
70  /// Get the number of elements in a tensor with this shape.
71  size_t getElementCount() const { return ElementCount; }
72  /// Get the size, in bytes, of one element.
73  size_t getElementByteSize() const;
74 
75  template <typename T> bool isElementType() const {
76  return getDataType<T>() == TypeIndex;
77  }
78 
79 private:
80  TensorSpec(const std::string &Name, int Port, int TypeIndex,
81  const std::vector<int64_t> &Shape);
82 
83  template <typename T> static int getDataType() {
84  llvm_unreachable("Undefined tensor type");
85  }
86 
87  std::string Name;
88  int Port = 0;
89  int TypeIndex = 0;
90  std::vector<int64_t> Shape;
91  size_t ElementCount = 0;
92 };
93 
94 /// Construct a TensorSpec from a JSON dictionary of the form:
95 /// { "name": <string>,
96 /// "port": <int>,
97 /// "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
98 /// "shape": <array of ints> }
99 /// For the "type" field, see the C++ primitive types used in
100 /// TFUTILS_SUPPORTED_TYPES.
101 Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
102  const json::Value &Value);
103 
104 struct LoggedFeatureSpec {
105  TensorSpec Spec;
106  Optional<std::string> LoggingName;
107  const std::string &getLoggingName() const {
108  return LoggingName ? *LoggingName : Spec.name();
109  }
110 };
111 
112 /// Load the output specs. If SpecFileOverride is not empty, that path is used.
113 /// Otherwise, the file is assumed to be called 'output_spec.json' and be found
114 /// under ModelPath (the model directory).
115 /// The first output tensor name must match ExpectedDecisionName.
116 /// In case of error, the return is None and the error is logged.
117 Optional<std::vector<LoggedFeatureSpec>>
118 loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
119  StringRef ModelPath, StringRef SpecFileOverride = StringRef());
120 
121 /// Logging utility - given an ordered specification of features, and assuming
122 /// a scalar reward, allow logging feature values and rewards, and then print
123 /// as tf.train.SequenceExample text protobuf.
124 /// The assumption is that, for an event to be logged (i.e. a set of feature
125 /// values and a reward), the user calls the log* API for each feature exactly
126 /// once, providing the index matching the position in the feature spec list
127 /// provided at construction. The example assumes the first feature's element
128 /// type is float, the second is int64, and the reward is float:
129 ///
130 /// event 0:
131 /// logFloatValue(0, ...)
132 /// logInt64Value(1, ...)
133 /// ...
134 /// logFloatReward(...)
135 /// event 1:
136 /// logFloatValue(0, ...)
137 /// logInt64Value(1, ...)
138 /// ...
139 /// logFloatReward(...)
140 ///
141 /// At the end, call print to generate the protobuf.
142 /// Alternatively, don't call logReward at the end of each event, just
143 /// log{Float|Int32|Int64}FinalReward at the end.
144 class LoggerDataImpl;
145 class Logger final {
146 public:
147  /// Construct a Logger. If IncludeReward is false, then logReward or
148  /// logFinalReward shouldn't be called, and the reward feature won't be
149  /// printed out.
150  Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
151  const TensorSpec &RewardSpec, bool IncludeReward);
152 
153  ~Logger();
154 
155  void logFloatReward(float Value);
156  void logInt32Reward(int32_t Value);
157  void logInt64Reward(int64_t Value);
158 
159  void logFloatFinalReward(float Value);
160  void logInt32FinalReward(int32_t Value);
161  void logInt64FinalReward(int64_t Value);
162 
163  void logFloatValue(size_t FeatureID, const float *Value);
164  void logInt32Value(size_t FeatureID, const int32_t *Value);
165  void logInt64Value(size_t FeatureID, const int64_t *Value);
166 
167  void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
168 
169  // Warning! For int32_t, the return is set up for int64_t, so the caller needs
170  // to piecemeal cast their int32_t values.
171  // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
172  // not supported by the tensorflow::SequenceExample proto. For small values,
173  // we can consider using bytes.
174  char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
175 
176  // Flush the content of the log to the stream, clearing the stored data in the
177  // process.
178  void flush(raw_ostream &OS);
179 
180 private:
181  std::vector<LoggedFeatureSpec> FeatureSpecs;
182  TensorSpec RewardSpec;
183  const bool IncludeReward;
184  std::unique_ptr<LoggerDataImpl> LoggerData;
185 };
186 
187 class TFModelEvaluator final {
188 public:
189  /// The result of a model evaluation. Handles the lifetime of the output
190  /// tensors, which means that their values need to be used before
191  /// the EvaluationResult's dtor is called.
192  class EvaluationResult {
193  public:
194  EvaluationResult(const EvaluationResult &) = delete;
195  EvaluationResult &operator=(const EvaluationResult &Other) = delete;
196 
197  EvaluationResult(EvaluationResult &&Other);
198  EvaluationResult &operator=(EvaluationResult &&Other);
199 
200  ~EvaluationResult();
201 
202  /// Get a (const) pointer to the first element of the tensor at Index.
203  template <typename T> T *getTensorValue(size_t Index) {
204  return static_cast<T *>(getUntypedTensorValue(Index));
205  }
206 
207  template <typename T> const T *getTensorValue(size_t Index) const {
208  return static_cast<T *>(getUntypedTensorValue(Index));
209  }
210 
211  /// Get a (const) pointer to the untyped data of the tensor.
212  void *getUntypedTensorValue(size_t Index);
213  const void *getUntypedTensorValue(size_t Index) const;
214 
215  private:
216  friend class TFModelEvaluator;
217  EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
218  std::unique_ptr<EvaluationResultImpl> Impl;
219  };
220 
221  TFModelEvaluator(StringRef SavedModelPath,
222  const std::vector<TensorSpec> &InputSpecs,
223  const std::vector<TensorSpec> &OutputSpecs,
224  const char *Tags = "serve");
225  TFModelEvaluator(StringRef SavedModelPath,
226  const std::vector<TensorSpec> &InputSpecs,
227  function_ref<TensorSpec(size_t)> GetOutputSpecs,
228  size_t OutputSpecsSize, const char *Tags = "serve");
229 
230  ~TFModelEvaluator();
231  TFModelEvaluator(const TFModelEvaluator &) = delete;
232  TFModelEvaluator(TFModelEvaluator &&) = delete;
233 
234  /// Evaluate the model, assuming it is valid. Returns None if the evaluation
235  /// fails or the model is invalid, or an EvaluationResult otherwise. The
236  /// inputs are assumed to have been already provided via getInput(). When
237  /// returning None, it also invalidates this object.
238  Optional<EvaluationResult> evaluate();
239 
240  /// Provides access to the input vector.
241  template <typename T> T *getInput(size_t Index) {
242  return static_cast<T *>(getUntypedInput(Index));
243  }
244 
245  /// Returns true if the tensorflow model was loaded successfully, false
246  /// otherwise.
247  bool isValid() const { return !!Impl; }
248 
249 private:
250  void *getUntypedInput(size_t Index);
251  std::unique_ptr<TFModelEvaluatorImpl> Impl;
252 };
253 
254 /// List of supported types, as a pair:
255 /// - C++ type
256 /// - enum name (implementation-specific)
257 #define TFUTILS_SUPPORTED_TYPES(M) \
258  M(float, TF_FLOAT) \
259  M(double, TF_DOUBLE) \
260  M(int8_t, TF_INT8) \
261  M(uint8_t, TF_UINT8) \
262  M(int16_t, TF_INT16) \
263  M(uint16_t, TF_UINT16) \
264  M(int32_t, TF_INT32) \
265  M(uint32_t, TF_UINT32) \
266  M(int64_t, TF_INT64) \
267  M(uint64_t, TF_UINT64)
268 
269 #define TFUTILS_GETDATATYPE_DEF(T, E) \
270  template <> int TensorSpec::getDataType<T>();
271 
272 TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_DEF)
273 
274 #undef TFUTILS_GETDATATYPE_DEF
275 } // namespace llvm
276 
277 #endif // LLVM_HAVE_TF_API
278 #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H
llvm
This is an optimization pass for GlobalISel generic memory operations.
Definition: AllocatorList.h:23
llvm::operator!=
bool operator!=(uint64_t V1, const APInt &V2)
Definition: APInt.h:1981
T
#define T
Definition: Mips16ISelLowering.cpp:341
JSON.h
StringMap.h
Index
uint32_t Index
Definition: ELFObjHandler.cpp:84
llvm::operator==
bool operator==(uint64_t V1, const APInt &V2)
Definition: APInt.h:1979
llvm_unreachable
#define llvm_unreachable(msg)
Marks that the current location is not supposed to be reachable.
Definition: ErrorHandling.h:134
name
static const char * name
Definition: SVEIntrinsicOpts.cpp:78
llvm::GraphProgram::Name
Name
Definition: GraphWriter.h:52
llvm::TargetStackID::Value
Value
Definition: TargetFrameLowering.h:27
isValid
static bool isValid(const char C)
Returns true if C is a valid mangled character: <0-9a-zA-Z_>.
Definition: RustDemangle.cpp:217
LLVMContext.h
Other
Optional< std::vector< StOtherPiece > > Other
Definition: ELFYAML.cpp:1195