//===- 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