1 //===- TrainingLogger.cpp - mlgo feature/reward logging -------------------===//
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 // This file implements logging infrastructure for extracting features and
10 // rewards for mlgo policy training.
11 //
12 //===----------------------------------------------------------------------===//
13 #include "llvm/Analysis/TensorSpec.h"
14 #include "llvm/Config/config.h"
15
16 #include "llvm/ADT/Twine.h"
17 #include "llvm/Analysis/Utils/TrainingLogger.h"
18 #include "llvm/Support/CommandLine.h"
19 #include "llvm/Support/Debug.h"
20 #include "llvm/Support/JSON.h"
21 #include "llvm/Support/MemoryBuffer.h"
22 #include "llvm/Support/Path.h"
23 #include "llvm/Support/raw_ostream.h"
24
25 #include <cassert>
26 #include <numeric>
27
28 using namespace llvm;
29
30 // FIXME(mtrofin): remove the flag altogether
31 static cl::opt<bool>
32 UseSimpleLogger("tfutils-use-simplelogger", cl::init(true), cl::Hidden,
33 cl::desc("Output simple (non-protobuf) log."));
34
writeHeader()35 void Logger::writeHeader() {
36 json::OStream JOS(*OS);
37 JOS.object([&]() {
38 JOS.attributeArray("features", [&]() {
39 for (const auto &TS : FeatureSpecs)
40 TS.toJSON(JOS);
41 });
42 if (IncludeReward) {
43 JOS.attributeBegin("score");
44 RewardSpec.toJSON(JOS);
45 JOS.attributeEnd();
46 }
47 });
48 *OS << "\n";
49 }
50
switchContext(StringRef Name)51 void Logger::switchContext(StringRef Name) {
52 CurrentContext = Name.str();
53 json::OStream JOS(*OS);
54 JOS.object([&]() { JOS.attribute("context", Name); });
55 *OS << "\n";
56 }
57
startObservation()58 void Logger::startObservation() {
59 auto I = ObservationIDs.insert({CurrentContext, 0});
60 size_t NewObservationID = I.second ? 0 : ++I.first->second;
61 json::OStream JOS(*OS);
62 JOS.object([&]() {
63 JOS.attribute("observation", static_cast<int64_t>(NewObservationID));
64 });
65 *OS << "\n";
66 }
67
endObservation()68 void Logger::endObservation() { *OS << "\n"; }
69
logRewardImpl(const char * RawData)70 void Logger::logRewardImpl(const char *RawData) {
71 assert(IncludeReward);
72 json::OStream JOS(*OS);
73 JOS.object([&]() {
74 JOS.attribute("outcome", static_cast<int64_t>(
75 ObservationIDs.find(CurrentContext)->second));
76 });
77 *OS << "\n";
78 writeTensor(RewardSpec, RawData);
79 *OS << "\n";
80 }
81
Logger(std::unique_ptr<raw_ostream> OS,const std::vector<TensorSpec> & FeatureSpecs,const TensorSpec & RewardSpec,bool IncludeReward)82 Logger::Logger(std::unique_ptr<raw_ostream> OS,
83 const std::vector<TensorSpec> &FeatureSpecs,
84 const TensorSpec &RewardSpec, bool IncludeReward)
85 : OS(std::move(OS)), FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
86 IncludeReward(IncludeReward) {
87 writeHeader();
88 }
89