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);
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();
120   void flush() { OS->flush(); }
121 
122   const std::string &currentContext() const { return CurrentContext; }
123 
124   /// Check if there is at least an observation for `currentContext()`.
125   bool hasObservationInProgress() const {
126     return hasAnyObservationForContext(CurrentContext);
127   }
128 
129   /// Check if there is at least an observation for the context `Ctx`.
130   bool hasAnyObservationForContext(StringRef Ctx) const {
131     return ObservationIDs.contains(Ctx);
132   }
133 
134   template <typename T> void logReward(T Value) {
135     logRewardImpl(reinterpret_cast<const char *>(&Value));
136   }
137 
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