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 /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have 151 /// corresponding indices) with any MLModelRunner implementations 152 /// corresponding to the model being trained/logged. 153 Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs, 154 const TensorSpec &RewardSpec, bool IncludeReward); 155 156 ~Logger(); 157 158 void logFloatReward(float Value); 159 void logInt32Reward(int32_t Value); 160 void logInt64Reward(int64_t Value); 161 162 void logFloatFinalReward(float Value); 163 void logInt32FinalReward(int32_t Value); 164 void logInt64FinalReward(int64_t Value); 165 166 void logFloatValue(size_t FeatureID, const float *Value); 167 void logInt32Value(size_t FeatureID, const int32_t *Value); 168 void logInt64Value(size_t FeatureID, const int64_t *Value); 169 170 void logSpecifiedTensorValue(size_t FeatureID, const char *RawData); 171 172 // Warning! For int32_t, the return is set up for int64_t, so the caller needs 173 // to piecemeal cast their int32_t values. 174 // FIXME: let's drop int32_t support. While it's supported by evaluator, it's 175 // not supported by the tensorflow::SequenceExample proto. For small values, 176 // we can consider using bytes. 177 char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID); 178 179 // Flush the content of the log to the stream, clearing the stored data in the 180 // process. 181 void flush(std::string *Str); 182 void flush(raw_ostream &OS); 183 184 // Flush a set of logs that are produced from the same module, e.g. 185 // per-function regalloc traces, as a google::protobuf::Struct message. 186 static void flushLogs(raw_ostream &OS, 187 const StringMap<std::unique_ptr<Logger>> &Loggers); 188 189 private: 190 std::vector<LoggedFeatureSpec> FeatureSpecs; 191 TensorSpec RewardSpec; 192 const bool IncludeReward; 193 std::unique_ptr<LoggerDataImpl> LoggerData; 194 }; 195 196 class TFModelEvaluator final { 197 public: 198 /// The result of a model evaluation. Handles the lifetime of the output 199 /// tensors, which means that their values need to be used before 200 /// the EvaluationResult's dtor is called. 201 class EvaluationResult { 202 public: 203 EvaluationResult(const EvaluationResult &) = delete; 204 EvaluationResult &operator=(const EvaluationResult &Other) = delete; 205 206 EvaluationResult(EvaluationResult &&Other); 207 EvaluationResult &operator=(EvaluationResult &&Other); 208 209 ~EvaluationResult(); 210 211 /// Get a (const) pointer to the first element of the tensor at Index. 212 template <typename T> T *getTensorValue(size_t Index) { 213 return static_cast<T *>(getUntypedTensorValue(Index)); 214 } 215 216 template <typename T> const T *getTensorValue(size_t Index) const { 217 return static_cast<T *>(getUntypedTensorValue(Index)); 218 } 219 220 /// Get a (const) pointer to the untyped data of the tensor. 221 void *getUntypedTensorValue(size_t Index); 222 const void *getUntypedTensorValue(size_t Index) const; 223 224 private: 225 friend class TFModelEvaluator; 226 EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl); 227 std::unique_ptr<EvaluationResultImpl> Impl; 228 }; 229 230 TFModelEvaluator(StringRef SavedModelPath, 231 const std::vector<TensorSpec> &InputSpecs, 232 const std::vector<TensorSpec> &OutputSpecs, 233 const char *Tags = "serve"); 234 TFModelEvaluator(StringRef SavedModelPath, 235 const std::vector<TensorSpec> &InputSpecs, 236 function_ref<TensorSpec(size_t)> GetOutputSpecs, 237 size_t OutputSpecsSize, const char *Tags = "serve"); 238 239 ~TFModelEvaluator(); 240 TFModelEvaluator(const TFModelEvaluator &) = delete; 241 TFModelEvaluator(TFModelEvaluator &&) = delete; 242 243 /// Evaluate the model, assuming it is valid. Returns None if the evaluation 244 /// fails or the model is invalid, or an EvaluationResult otherwise. The 245 /// inputs are assumed to have been already provided via getInput(). When 246 /// returning None, it also invalidates this object. 247 Optional<EvaluationResult> evaluate(); 248 249 /// Provides access to the input vector. 250 template <typename T> T *getInput(size_t Index) { 251 return static_cast<T *>(getUntypedInput(Index)); 252 } 253 254 /// Returns true if the tensorflow model was loaded successfully, false 255 /// otherwise. 256 bool isValid() const { return !!Impl; } 257 258 /// Untyped access to input. 259 void *getUntypedInput(size_t Index); 260 261 private: 262 std::unique_ptr<TFModelEvaluatorImpl> Impl; 263 }; 264 265 /// List of supported types, as a pair: 266 /// - C++ type 267 /// - enum name (implementation-specific) 268 #define TFUTILS_SUPPORTED_TYPES(M) \ 269 M(float, TF_FLOAT) \ 270 M(double, TF_DOUBLE) \ 271 M(int8_t, TF_INT8) \ 272 M(uint8_t, TF_UINT8) \ 273 M(int16_t, TF_INT16) \ 274 M(uint16_t, TF_UINT16) \ 275 M(int32_t, TF_INT32) \ 276 M(uint32_t, TF_UINT32) \ 277 M(int64_t, TF_INT64) \ 278 M(uint64_t, TF_UINT64) 279 280 #define TFUTILS_GETDATATYPE_DEF(T, E) \ 281 template <> int TensorSpec::getDataType<T>(); 282 283 TFUTILS_SUPPORTED_TYPES(TFUTILS_GETDATATYPE_DEF) 284 285 #undef TFUTILS_GETDATATYPE_DEF 286 } // namespace llvm 287 288 #endif // LLVM_HAVE_TF_API 289 #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H 290