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 (std::optional<TensorSpec> AdviceSpec); |
100 | void writeTensor(const TensorSpec &Spec, const char *RawData) { |
101 | OS->write(Ptr: RawData, Size: 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 ¤tContext() const { return CurrentContext; } |
123 | |
124 | /// Check if there is at least an observation for `currentContext()`. |
125 | bool hasObservationInProgress() const { |
126 | return hasAnyObservationForContext(Ctx: 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(Key: Ctx); |
132 | } |
133 | |
134 | template <typename T> void logReward(T Value) { |
135 | logRewardImpl(RawData: reinterpret_cast<const char *>(&Value)); |
136 | } |
137 | |
138 | void logTensorValue(size_t FeatureID, const char *RawData) { |
139 | writeTensor(Spec: FeatureSpecs[FeatureID], RawData); |
140 | } |
141 | }; |
142 | |
143 | } // namespace llvm |
144 | #endif // LLVM_ANALYSIS_UTILS_TRAININGLOGGER_H |
145 | |