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/Analysis/TensorSpec.h" 17 #include "llvm/IR/LLVMContext.h" 18 #include "llvm/Support/JSON.h" 19 20 #include <memory> 21 #include <vector> 22 23 namespace llvm { 24 25 /// Load a SavedModel, find the given inputs and outputs, and setup storage 26 /// for input tensors. The user is responsible for correctly dimensioning the 27 /// input tensors and setting their values before calling evaluate(). 28 /// To initialize: 29 /// - construct the object 30 /// - initialize the input tensors using initInput. Indices must correspond to 31 /// indices in the InputNames used at construction. 32 /// To use: 33 /// - set input values by using getInput to get each input tensor, and then 34 /// setting internal scalars, for all dimensions (tensors are row-major: 35 /// https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/c/c_api.h#L205) 36 /// - call evaluate. The input tensors' values are not consumed after this, and 37 /// may still be read. 38 /// - use the outputs in the output vector 39 class TFModelEvaluatorImpl; 40 class EvaluationResultImpl; 41 42 /// Logging utility - given an ordered specification of features, and assuming 43 /// a scalar reward, allow logging feature values and rewards, and then print 44 /// as tf.train.SequenceExample text protobuf. 45 /// The assumption is that, for an event to be logged (i.e. a set of feature 46 /// values and a reward), the user calls the log* API for each feature exactly 47 /// once, providing the index matching the position in the feature spec list 48 /// provided at construction. The example assumes the first feature's element 49 /// type is float, the second is int64, and the reward is float: 50 /// 51 /// event 0: 52 /// logFloatValue(0, ...) 53 /// logInt64Value(1, ...) 54 /// ... 55 /// logFloatReward(...) 56 /// event 1: 57 /// logFloatValue(0, ...) 58 /// logInt64Value(1, ...) 59 /// ... 60 /// logFloatReward(...) 61 /// 62 /// At the end, call print to generate the protobuf. 63 /// Alternatively, don't call logReward at the end of each event, just 64 /// log{Float|Int32|Int64}FinalReward at the end. 65 class LoggerDataImpl; 66 class Logger final { 67 public: 68 /// Construct a Logger. If IncludeReward is false, then logReward or 69 /// logFinalReward shouldn't be called, and the reward feature won't be 70 /// printed out. 71 /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have 72 /// corresponding indices) with any MLModelRunner implementations 73 /// corresponding to the model being trained/logged. 74 Logger(const std::vector<LoggedFeatureSpec> &FeatureSpecs, 75 const TensorSpec &RewardSpec, bool IncludeReward); 76 77 ~Logger(); 78 79 void logFloatReward(float Value); 80 void logInt32Reward(int32_t Value); 81 void logInt64Reward(int64_t Value); 82 83 void logFloatFinalReward(float Value); 84 void logInt32FinalReward(int32_t Value); 85 void logInt64FinalReward(int64_t Value); 86 87 void logFloatValue(size_t FeatureID, const float *Value); 88 void logInt32Value(size_t FeatureID, const int32_t *Value); 89 void logInt64Value(size_t FeatureID, const int64_t *Value); 90 91 void logSpecifiedTensorValue(size_t FeatureID, const char *RawData); 92 93 // Warning! For int32_t, the return is set up for int64_t, so the caller needs 94 // to piecemeal cast their int32_t values. 95 // FIXME: let's drop int32_t support. While it's supported by evaluator, it's 96 // not supported by the tensorflow::SequenceExample proto. For small values, 97 // we can consider using bytes. 98 char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID); 99 100 // Flush the content of the log to the stream, clearing the stored data in the 101 // process. 102 void flush(std::string *Str); 103 void flush(raw_ostream &OS); 104 105 // Flush a set of logs that are produced from the same module, e.g. 106 // per-function regalloc traces, as a google::protobuf::Struct message. 107 static void flushLogs(raw_ostream &OS, 108 const StringMap<std::unique_ptr<Logger>> &Loggers); 109 110 private: 111 std::vector<LoggedFeatureSpec> FeatureSpecs; 112 TensorSpec RewardSpec; 113 const bool IncludeReward; 114 std::unique_ptr<LoggerDataImpl> LoggerData; 115 }; 116 117 class TFModelEvaluator final { 118 public: 119 /// The result of a model evaluation. Handles the lifetime of the output 120 /// tensors, which means that their values need to be used before 121 /// the EvaluationResult's dtor is called. 122 class EvaluationResult { 123 public: 124 EvaluationResult(const EvaluationResult &) = delete; 125 EvaluationResult &operator=(const EvaluationResult &Other) = delete; 126 127 EvaluationResult(EvaluationResult &&Other); 128 EvaluationResult &operator=(EvaluationResult &&Other); 129 130 ~EvaluationResult(); 131 132 /// Get a (const) pointer to the first element of the tensor at Index. 133 template <typename T> T *getTensorValue(size_t Index) { 134 return static_cast<T *>(getUntypedTensorValue(Index)); 135 } 136 137 template <typename T> const T *getTensorValue(size_t Index) const { 138 return static_cast<T *>(getUntypedTensorValue(Index)); 139 } 140 141 /// Get a (const) pointer to the untyped data of the tensor. 142 void *getUntypedTensorValue(size_t Index); 143 const void *getUntypedTensorValue(size_t Index) const; 144 145 private: 146 friend class TFModelEvaluator; 147 EvaluationResult(std::unique_ptr<EvaluationResultImpl> Impl); 148 std::unique_ptr<EvaluationResultImpl> Impl; 149 }; 150 151 TFModelEvaluator(StringRef SavedModelPath, 152 const std::vector<TensorSpec> &InputSpecs, 153 const std::vector<TensorSpec> &OutputSpecs, 154 const char *Tags = "serve"); 155 TFModelEvaluator(StringRef SavedModelPath, 156 const std::vector<TensorSpec> &InputSpecs, 157 function_ref<TensorSpec(size_t)> GetOutputSpecs, 158 size_t OutputSpecsSize, const char *Tags = "serve"); 159 160 ~TFModelEvaluator(); 161 TFModelEvaluator(const TFModelEvaluator &) = delete; 162 TFModelEvaluator(TFModelEvaluator &&) = delete; 163 164 /// Evaluate the model, assuming it is valid. Returns None if the evaluation 165 /// fails or the model is invalid, or an EvaluationResult otherwise. The 166 /// inputs are assumed to have been already provided via getInput(). When 167 /// returning None, it also invalidates this object. 168 Optional<EvaluationResult> evaluate(); 169 170 /// Provides access to the input vector. 171 template <typename T> T *getInput(size_t Index) { 172 return static_cast<T *>(getUntypedInput(Index)); 173 } 174 175 /// Returns true if the tensorflow model was loaded successfully, false 176 /// otherwise. 177 bool isValid() const { return !!Impl; } 178 179 /// Untyped access to input. 180 void *getUntypedInput(size_t Index); 181 182 private: 183 std::unique_ptr<TFModelEvaluatorImpl> Impl; 184 }; 185 186 } // namespace llvm 187 188 #endif // LLVM_HAVE_TF_API 189 #endif // LLVM_ANALYSIS_UTILS_TFUTILS_H 190