TFUtils.h 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. #pragma once
  2. #ifdef __GNUC__
  3. #pragma GCC diagnostic push
  4. #pragma GCC diagnostic ignored "-Wunused-parameter"
  5. #endif
  6. //===- TFUtils.h - utilities for tensorflow C API ---------------*- C++ -*-===//
  7. //
  8. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  9. // See https://llvm.org/LICENSE.txt for license information.
  10. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  11. //
  12. //===----------------------------------------------------------------------===//
  13. //
  14. #ifndef LLVM_ANALYSIS_UTILS_TFUTILS_H
  15. #define LLVM_ANALYSIS_UTILS_TFUTILS_H
  16. #include "llvm/Config/llvm-config.h"
  17. #ifdef LLVM_HAVE_TF_API
  18. #include "llvm/ADT/StringMap.h"
  19. #include "llvm/IR/LLVMContext.h"
  20. #include "llvm/Support/JSON.h"
  21. #include <memory>
  22. #include <vector>
  23. namespace llvm {
  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. /// TensorSpec encapsulates the specification of a tensor: its dimensions, or
  41. /// "shape" (row-major), its type (see TensorSpec::getDataType specializations
  42. /// for supported types), its name and port (see "TensorFlow: Large-Scale
  43. /// Machine Learning on Heterogeneous Distributed Systems", section 4.2, para 2:
  44. /// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45166.pdf)
  45. ///
  46. /// TensorSpec is used to set up a TFModelEvaluator by describing the expected
  47. /// inputs and outputs.
  48. class TensorSpec final {
  49. public:
  50. template <typename T>
  51. static TensorSpec createSpec(const std::string &Name,
  52. const std::vector<int64_t> &Shape,
  53. int Port = 0) {
  54. return TensorSpec(Name, Port, getDataType<T>(), Shape);
  55. }
  56. const std::string &name() const { return Name; }
  57. int port() const { return Port; }
  58. int typeIndex() const { return TypeIndex; }
  59. const std::vector<int64_t> &shape() const { return Shape; }
  60. bool operator==(const TensorSpec &Other) const {
  61. return Name == Other.Name && Port == Other.Port &&
  62. TypeIndex == Other.TypeIndex && Shape == Other.Shape;
  63. }
  64. bool operator!=(const TensorSpec &Other) const { return !(*this == Other); }
  65. /// Get the number of elements in a tensor with this shape.
  66. size_t getElementCount() const { return ElementCount; }
  67. /// Get the size, in bytes, of one element.
  68. size_t getElementByteSize() const;
  69. template <typename T> bool isElementType() const {
  70. return getDataType<T>() == TypeIndex;
  71. }
  72. private:
  73. TensorSpec(const std::string &Name, int Port, int TypeIndex,
  74. const std::vector<int64_t> &Shape);
  75. template <typename T> static int getDataType() {
  76. llvm_unreachable("Undefined tensor type");
  77. }
  78. std::string Name;
  79. int Port = 0;
  80. int TypeIndex = 0;
  81. std::vector<int64_t> Shape;
  82. size_t ElementCount = 0;
  83. };
  84. /// Construct a TensorSpec from a JSON dictionary of the form:
  85. /// { "name": <string>,
  86. /// "port": <int>,
  87. /// "type": <string. Use LLVM's types, e.g. float, double, int64_t>,
  88. /// "shape": <array of ints> }
  89. /// For the "type" field, see the C++ primitive types used in
  90. /// TFUTILS_SUPPORTED_TYPES.
  91. Optional<TensorSpec> getTensorSpecFromJSON(LLVMContext &Ctx,
  92. const json::Value &Value);
  93. struct LoggedFeatureSpec {
  94. TensorSpec Spec;
  95. Optional<std::string> LoggingName;
  96. const std::string &getLoggingName() const {
  97. return LoggingName ? *LoggingName : Spec.name();
  98. }
  99. };
  100. /// Load the output specs. If SpecFileOverride is not empty, that path is used.
  101. /// Otherwise, the file is assumed to be called 'output_spec.json' and be found
  102. /// under ModelPath (the model directory).
  103. /// The first output tensor name must match ExpectedDecisionName.
  104. /// In case of error, the return is None and the error is logged.
  105. Optional<std::vector<LoggedFeatureSpec>>
  106. loadOutputSpecs(LLVMContext &Ctx, StringRef ExpectedDecisionName,
  107. StringRef ModelPath, StringRef SpecFileOverride = StringRef());
  108. /// Logging utility - given an ordered specification of features, and assuming
  109. /// a scalar reward, allow logging feature values and rewards, and then print
  110. /// as tf.train.SequenceExample text protobuf.
  111. /// The assumption is that, for an event to be logged (i.e. a set of feature
  112. /// values and a reward), the user calls the log* API for each feature exactly
  113. /// once, providing the index matching the position in the feature spec list
  114. /// provided at construction. The example assumes the first feature's element
  115. /// type is float, the second is int64, and the reward is float:
  116. ///
  117. /// event 0:
  118. /// logFloatValue(0, ...)
  119. /// logInt64Value(1, ...)
  120. /// ...
  121. /// logFloatReward(...)
  122. /// event 1:
  123. /// logFloatValue(0, ...)
  124. /// logInt64Value(1, ...)
  125. /// ...
  126. /// logFloatReward(...)
  127. ///
  128. /// At the end, call print to generate the protobuf.
  129. /// Alternatively, don't call logReward at the end of each event, just
  130. /// log{Float|Int32|Int64}FinalReward at the end.
  131. class LoggerDataImpl;
  132. class Logger final {
  133. public:
  134. /// Construct a Logger. If IncludeReward is false, then logReward or
  135. /// logFinalReward shouldn't be called, and the reward feature won't be
  136. /// printed out.
  137. /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have
  138. /// corresponding indices) with any MLModelRunner implementations
  139. /// corresponding to the model being trained/logged.
  140. Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs,
  141. const TensorSpec &RewardSpec, bool IncludeReward);
  142. ~Logger();
  143. void logFloatReward(float Value);
  144. void logInt32Reward(int32_t Value);
  145. void logInt64Reward(int64_t Value);
  146. void logFloatFinalReward(float Value);
  147. void logInt32FinalReward(int32_t Value);
  148. void logInt64FinalReward(int64_t Value);
  149. void logFloatValue(size_t FeatureID, const float *Value);
  150. void logInt32Value(size_t FeatureID, const int32_t *Value);
  151. void logInt64Value(size_t FeatureID, const int64_t *Value);
  152. void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
  153. // Warning! For int32_t, the return is set up for int64_t, so the caller needs
  154. // to piecemeal cast their int32_t values.
  155. // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
  156. // not supported by the tensorflow::SequenceExample proto. For small values,
  157. // we can consider using bytes.
  158. char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
  159. // Flush the content of the log to the stream, clearing the stored data in the
  160. // process.
  161. void flush(std::string *Str);
  162. void flush(raw_ostream &OS);
  163. // Flush a set of logs that are produced from the same module, e.g.
  164. // per-function regalloc traces, as a google::protobuf::Struct message.
  165. static void flushLogs(raw_ostream &OS,
  166. const StringMap<std::unique_ptr<Logger>> &Loggers);
  167. private:
  168. std::vector<LoggedFeatureSpec> FeatureSpecs;
  169. TensorSpec RewardSpec;
  170. const bool IncludeReward;
  171. std::unique_ptr<LoggerDataImpl> LoggerData;
  172. };
  173. class TFModelEvaluator final {
  174. public:
  175. /// The result of a model evaluation. Handles the lifetime of the output
  176. /// tensors, which means that their values need to be used before
  177. /// the EvaluationResult's dtor is called.
  178. class EvaluationResult {
  179. public:
  180. EvaluationResult(const EvaluationResult &) = delete;
  181. EvaluationResult &operator=(const EvaluationResult &Other) = delete;
  182. EvaluationResult(EvaluationResult &&Other);
  183. EvaluationResult &operator=(EvaluationResult &&Other);
  184. ~EvaluationResult();
  185. /// Get a (const) pointer to the first element of the tensor at Index.
  186. template <typename T> T *getTensorValue(size_t Index) {
  187. return static_cast<T *>(getUntypedTensorValue(Index));
  188. }
  189. template <typename T> const T *getTensorValue(size_t Index) const {
  190. return static_cast<T *>(getUntypedTensorValue(Index));
  191. }
  192. /// Get a (const) pointer to the untyped data of the tensor.
  193. void *getUntypedTensorValue(size_t Index);
  194. const void *getUntypedTensorValue(size_t Index) const;
  195. private:
  196. friend class TFModelEvaluator;
  197. EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl);
  198. std::unique_ptr<EvaluationResultImpl> Impl;
  199. };
  200. TFModelEvaluator(StringRef SavedModelPath,
  201. const std::vector<TensorSpec> &InputSpecs,
  202. const std::vector<TensorSpec> &OutputSpecs,
  203. const char *Tags = "serve");
  204. TFModelEvaluator(StringRef SavedModelPath,
  205. const std::vector<TensorSpec> &InputSpecs,
  206. function_ref<TensorSpec(size_t)> GetOutputSpecs,
  207. size_t OutputSpecsSize, const char *Tags = "serve");
  208. ~TFModelEvaluator();
  209. TFModelEvaluator(const TFModelEvaluator &) = delete;
  210. TFModelEvaluator(TFModelEvaluator &&) = delete;
  211. /// Evaluate the model, assuming it is valid. Returns None if the evaluation
  212. /// fails or the model is invalid, or an EvaluationResult otherwise. The
  213. /// inputs are assumed to have been already provided via getInput(). When
  214. /// returning None, it also invalidates this object.
  215. Optional<EvaluationResult> evaluate();
  216. /// Provides access to the input vector.
  217. template <typename T> T *getInput(size_t Index) {
  218. return static_cast<T *>(getUntypedInput(Index));
  219. }
  220. /// Returns true if the tensorflow model was loaded successfully, false
  221. /// otherwise.
  222. bool isValid() const { return !!Impl; }
  223. /// Untyped access to input.
  224. void *getUntypedInput(size_t Index);
  225. private:
  226. std::unique_ptr<TFModelEvaluatorImpl> Impl;
  227. };
  228. /// List of supported types, as a pair:
  229. /// - C++ type
  230. /// - enum name (implementation-specific)
  231. #define TFUTILS_SUPPORTED_TYPES(M) \
  232. M(float, TF_FLOAT) \
  233. M(double, TF_DOUBLE) \
  234. M(int8_t, TF_INT8) \
  235. M(uint8_t, TF_UINT8) \
  236. M(int16_t, TF_INT16) \
  237. M(uint16_t, TF_UINT16) \
  238. M(int32_t, TF_INT32) \
  239. M(uint32_t, TF_UINT32) \
  240. M(int64_t, TF_INT64) \
  241. M(uint64_t, TF_UINT64)
  242. #define TFUTILS_GETDATATYPE_DEF(T, E) \
  243. template <> int TensorSpec::getDataType<T>();
  244. TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_DEF)
  245. #undef TFUTILS_GETDATATYPE_DEF
  246. } // namespace llvm
  247. #endif // LLVM_HAVE_TF_API
  248. #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H
  249. #ifdef __GNUC__
  250. #pragma GCC diagnostic pop
  251. #endif