1 //===- TrainingLogger.h - mlgo feature/reward logging ----------*- 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 // The design goals of the logger are: 10 // - no dependencies that llvm doesn't already have. 11 // - support streaming, so that we don't need to buffer data during compilation 12 // - 0-decoding tensor values. Tensor values are potentially very large buffers 13 // of scalars. Because of their potentially large size, avoiding 14 // serialization/deserialization overhead is preferred. 15 // 16 // The simple logger produces an output of the form (each line item on its line) 17 // - header: a json object describing the data that will follow. 18 // - context: e.g. function name, for regalloc, or "default" for module-wide 19 // optimizations like the inliner. This is the context to which the subsequent 20 // data corresponds. 21 // - observation number. 22 // - tensor values - raw bytes of the tensors, in the order given in the header. 23 // The values are in succession, i.e. no separator is found between successive 24 // tensor values. At the end, there is a new line character. 25 // - [score] - this is optional, and is present if it was present in the header. 26 // Currently, for final rewards, we output "0" scores after each observation, 27 // except for the last one. 28 // <repeat> 29 // The file should be read as binary, but the reason we use newlines is mostly 30 // ease of debugging: the log can be opened in a text editor and, while tensor 31 // values are inscrutable, at least the sequence of data can be easily observed. 32 // Of course, the buffer of tensor values could contain '\n' bytes. A reader 33 // should use the header information to know how much data to read for the 34 // tensor values, and not use line information for that. 35 // 36 // An example reader, used for test, is available at 37 // Analysis/models/log_reader.py 38 // 39 // Example: 40 // {"features":[list of TensorSpecs], "score":<a tensor spec>} 41 // {"context": "aFunction"} 42 // {"observation": 0} 43 // <bytes> 44 // {"outcome": 0} 45 // <bytes for the tensor corresponding to the "score" spec in the header> 46 // {"observation": 1} 47 // ... 48 // {"context": "anotherFunction"} 49 // {"observation": 0} 50 // ... 51 // 52 53 #ifndef LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H 54 #define LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H 55 56 #include "llvm/Config/llvm-config.h" 57 58 #include "llvm/ADT/StringMap.h" 59 #include "llvm/Analysis/TensorSpec.h" 60 #include "llvm/IR/LLVMContext.h" 61 #include "llvm/Support/JSON.h" 62 63 #include <memory> 64 #include <optional> 65 #include <vector> 66 67 namespace llvm { 68 69 /// Logging utility - given an ordered specification of features, and assuming 70 /// a scalar reward, allow logging feature values and rewards. 71 /// The assumption is that, for an event to be logged (i.e. a set of feature 72 /// values and a reward), the user calls the log* API for each feature exactly 73 /// once, providing the index matching the position in the feature spec list 74 /// provided at construction. The example assumes the first feature's element 75 /// type is float, the second is int64, and the reward is float: 76 /// 77 /// event 0: 78 /// logFloatValue(0, ...) 79 /// logInt64Value(1, ...) 80 /// ... 81 /// logFloatReward(...) 82 /// event 1: 83 /// logFloatValue(0, ...) 84 /// logInt64Value(1, ...) 85 /// ... 86 /// logFloatReward(...) 87 /// 88 /// At the end, call print to generate the log. 89 /// Alternatively, don't call logReward at the end of each event, just 90 /// log{Float|Int32|Int64}FinalReward at the end. 91 class Logger final { 92 std::unique_ptr<raw_ostream> OS; 93 const std::vector<TensorSpec> FeatureSpecs; 94 const TensorSpec RewardSpec; 95 const bool IncludeReward; 96 StringMap<size_t> ObservationIDs; 97 std::string CurrentContext; 98 99 void writeHeader(std::optional<TensorSpec> AdviceSpec); writeTensor(const TensorSpec & Spec,const char * RawData)100 void writeTensor(const TensorSpec &Spec, const char *RawData) { 101 OS->write(RawData, Spec.getTotalTensorBufferSize()); 102 } 103 void logRewardImpl(const char *RawData); 104 105 public: 106 /// Construct a Logger. If IncludeReward is false, then logReward or 107 /// logFinalReward shouldn't be called, and the reward feature won't be 108 /// printed out. 109 /// NOTE: the FeatureSpecs are expected to be in the same order (i.e. have 110 /// corresponding indices) with any MLModelRunner implementations 111 /// corresponding to the model being trained/logged. 112 Logger(std::unique_ptr<raw_ostream> OS, 113 const std::vector<TensorSpec> &FeatureSpecs, 114 const TensorSpec &RewardSpec, bool IncludeReward, 115 std::optional<TensorSpec> AdviceSpec = std::nullopt); 116 117 void switchContext(StringRef Name); 118 void startObservation(); 119 void endObservation(); flush()120 void flush() { OS->flush(); } 121 currentContext()122 const std::string ¤tContext() const { return CurrentContext; } 123 124 /// Check if there is at least an observation for `currentContext()`. hasObservationInProgress()125 bool hasObservationInProgress() const { 126 return hasAnyObservationForContext(CurrentContext); 127 } 128 129 /// Check if there is at least an observation for the context `Ctx`. hasAnyObservationForContext(StringRef Ctx)130 bool hasAnyObservationForContext(StringRef Ctx) const { 131 return ObservationIDs.contains(Ctx); 132 } 133 logReward(T Value)134 template <typename T> void logReward(T Value) { 135 logRewardImpl(reinterpret_cast<const char *>(&Value)); 136 } 137 logTensorValue(size_t FeatureID,const char * RawData)138 void logTensorValue(size_t FeatureID, const char *RawData) { 139 writeTensor(FeatureSpecs[FeatureID], RawData); 140 } 141 }; 142 143 } // namespace llvm 144 #endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H 145