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