[llvm] 6d11baf - [mlgo] Stream the training data
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 20 07:01:18 PST 2023
Author: Mircea Trofin
Date: 2023-01-20T07:01:08-08:00
New Revision: 6d11baf02b33c3ec89178f94627ecec5025dac7e
URL: https://github.com/llvm/llvm-project/commit/6d11baf02b33c3ec89178f94627ecec5025dac7e
DIFF: https://github.com/llvm/llvm-project/commit/6d11baf02b33c3ec89178f94627ecec5025dac7e.diff
LOG: [mlgo] Stream the training data
This leverages the new logging format in that we don't need to buffer
the training data, we can just write it out.
Differential Revision: https://reviews.llvm.org/D142168
Added:
Modified:
llvm/include/llvm/Analysis/Utils/TrainingLogger.h
llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
llvm/lib/Analysis/TrainingLogger.cpp
llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
Removed:
################################################################################
diff --git a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
index 7ee25bb44df51..b7db58f67436a 100644
--- a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
+++ b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h
@@ -61,6 +61,7 @@
#include "llvm/Support/JSON.h"
#include <memory>
+#include <optional>
#include <vector>
namespace llvm {
@@ -88,19 +89,18 @@ namespace llvm {
/// 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;
- std::vector<std::unique_ptr<char[]>> FeatureStorage;
- std::vector<std::unique_ptr<char[]>> RewardStorage;
- raw_ostream &dumpHeader(raw_ostream &OS) const;
- raw_ostream &startContext(raw_ostream &OS, StringRef Name) const;
- raw_ostream &startObservation(raw_ostream &OS, size_t Nr) const;
- raw_ostream &writeOutcome(raw_ostream &OS, size_t CurrentObservationID) const;
- char *addNewTensor(size_t FeatureID);
- size_t getNrRecords() const;
+ StringMap<size_t> ObservationIDs;
+ std::string CurrentContext;
- void logRewardImpl(const char *Value, size_t Size);
+ 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
@@ -109,44 +109,27 @@ class Logger final {
/// 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(const std::vector<TensorSpec> &FeatureSpecs,
- const TensorSpec &RewardSpec, bool IncludeReward)
- : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
- IncludeReward(IncludeReward) {}
-
- template <typename T> void logReward(T Value) {
- logRewardImpl(reinterpret_cast<const char *>(&Value), sizeof(T));
- }
- void logFloatReward(float Value);
- void logInt32Reward(int32_t Value);
- void logInt64Reward(int64_t Value);
-
- void logFloatFinalReward(float Value);
- void logInt32FinalReward(int32_t Value);
- void logInt64FinalReward(int64_t Value);
+ Logger(std::unique_ptr<raw_ostream> OS,
+ const std::vector<TensorSpec> &FeatureSpecs,
+ const TensorSpec &RewardSpec, bool IncludeReward);
- void logFloatValue(size_t FeatureID, const float *Value);
- void logInt32Value(size_t FeatureID, const int32_t *Value);
- void logInt64Value(size_t FeatureID, const int64_t *Value);
+ void switchContext(StringRef Name);
+ void startObservation();
+ void endObservation();
- void logSpecifiedTensorValue(size_t FeatureID, const char *RawData);
+ const std::string ¤tContext() const { return CurrentContext; }
- // Warning! For int32_t, the return is set up for int64_t, so the caller needs
- // to piecemeal cast their int32_t values.
- // FIXME: let's drop int32_t support. While it's supported by evaluator, it's
- // not supported by the tensorflow::SequenceExample proto. For small values,
- // we can consider using bytes.
- char *addEntryAndGetFloatOrInt64Buffer(size_t FeatureID);
+ bool hasObservationInProgress() const {
+ return ObservationIDs.find(CurrentContext) != ObservationIDs.end();
+ }
- // Flush the content of the log to the stream, clearing the stored data in the
- // process.
- raw_ostream &flush(raw_ostream &OS, bool WithHeader = true,
- StringRef Context = "default") const;
+ template <typename T> void logReward(T Value) {
+ logRewardImpl(reinterpret_cast<const char *>(&Value));
+ }
- // Flush a set of logs that are produced from the same module, e.g.
- // per-function regalloc traces.
- static void flushLogs(raw_ostream &OS,
- const StringMap<std::unique_ptr<Logger>> &Loggers);
+ void logTensorValue(size_t FeatureID, const char *RawData) {
+ writeTensor(FeatureSpecs[FeatureID], RawData);
+ }
};
} // namespace llvm
diff --git a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
index 63aa4ca3301f6..a91d2ffe60424 100644
--- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
+++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp
@@ -102,9 +102,6 @@ class TrainingLogger final {
void logInlineEvent(const InlineEvent &Event,
const MLModelRunner &ModelRunner);
- /// Print the stored tensors.
- void print();
-
private:
StringRef LogFileName;
const ModelUnderTrainingRunner *const MUTR;
@@ -150,7 +147,6 @@ class DevelopmentModeMLInlineAdvisor : public MLInlineAdvisor {
size_t getTotalSizeEstimate();
- virtual ~DevelopmentModeMLInlineAdvisor();
void updateNativeSizeEstimate(int64_t Change) {
*CurrentNativeSize += Change;
}
@@ -288,45 +284,48 @@ TrainingLogger::TrainingLogger(StringRef LogFileName,
DecisionPos = FT.size();
FT.push_back(TensorSpec::createSpec<int64_t>(DecisionName, {1}));
+ std::error_code EC;
+ auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
+ if (EC)
+ dbgs() << (EC.message() + ":" + TrainingLog);
L = std::make_unique<Logger>(
- FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
+ std::move(OS), FT, TensorSpec::createSpec<int64_t>(RewardName, {1}),
InlineSizeEstimatorAnalysis::isEvaluatorRequested());
+ L->switchContext("");
}
/// Log one inlining event.
void TrainingLogger::logInlineEvent(const InlineEvent &Event,
const MLModelRunner &ModelRunner) {
+ L->startObservation();
size_t CurrentFeature = 0;
- for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) {
- int64_t F = *ModelRunner.getTensor<int64_t>(CurrentFeature);
- L->logInt64Value(CurrentFeature, &F);
- }
+ for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature)
+ L->logTensorValue(CurrentFeature,
+ reinterpret_cast<const char *>(
+ ModelRunner.getTensorUntyped(CurrentFeature)));
if (MUTR)
for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) {
const char *RawData =
reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I));
- L->logSpecifiedTensorValue(CurrentFeature, RawData);
+ L->logTensorValue(CurrentFeature, RawData);
++CurrentFeature;
}
assert(CurrentFeature == DefaultDecisionPos);
- L->logInt64Value(DefaultDecisionPos, &Event.DefaultDecision);
- L->logInt64Value(DecisionPos, &Event.AdvisedDecision);
+ L->logTensorValue(DefaultDecisionPos,
+ reinterpret_cast<const char *>(&Event.DefaultDecision));
+ L->logTensorValue(DecisionPos,
+ reinterpret_cast<const char *>(&Event.AdvisedDecision));
+ L->endObservation();
if (InlineSizeEstimatorAnalysis::isEvaluatorRequested())
- L->logInt64Reward(Event.Reward);
+ L->logReward(Event.Reward);
// For debugging / later use
Effects.push_back(Event.Effect);
}
-void TrainingLogger::print() {
- std::error_code EC;
- raw_fd_ostream OutFile(LogFileName, EC);
- L->flush(OutFile);
-}
-
DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
Module &M, ModuleAnalysisManager &MAM,
std::unique_ptr<MLModelRunner> ModelRunner,
@@ -342,11 +341,6 @@ DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor(
assert(IsDoingInference || isLogging());
}
-DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() {
- if (isLogging())
- Logger->print();
-}
-
std::optional<size_t>
DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const {
if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
diff --git a/llvm/lib/Analysis/TrainingLogger.cpp b/llvm/lib/Analysis/TrainingLogger.cpp
index 72ec14ae7c9bd..dcee8d40c53d7 100644
--- a/llvm/lib/Analysis/TrainingLogger.cpp
+++ b/llvm/lib/Analysis/TrainingLogger.cpp
@@ -32,8 +32,8 @@ static cl::opt<bool>
UseSimpleLogger("tfutils-use-simplelogger", cl::init(true), cl::Hidden,
cl::desc("Output simple (non-protobuf) log."));
-raw_ostream &Logger::dumpHeader(raw_ostream &OS) const {
- json::OStream JOS(OS);
+void Logger::writeHeader() {
+ json::OStream JOS(*OS);
JOS.object([&]() {
JOS.attributeArray("features", [&]() {
for (const auto &TS : FeatureSpecs)
@@ -45,140 +45,44 @@ raw_ostream &Logger::dumpHeader(raw_ostream &OS) const {
JOS.attributeEnd();
}
});
- OS << "\n";
- return OS;
+ *OS << "\n";
}
-raw_ostream &Logger::startContext(raw_ostream &OS, StringRef Name) const {
- json::OStream JOS(OS);
+void Logger::switchContext(StringRef Name) {
+ CurrentContext = Name.str();
+ json::OStream JOS(*OS);
JOS.object([&]() { JOS.attribute("context", Name); });
- OS << "\n";
- return OS;
+ *OS << "\n";
}
-raw_ostream &Logger::startObservation(raw_ostream &OS, size_t Nr) const {
- json::OStream JOS(OS);
- JOS.object([&]() { JOS.attribute("observation", static_cast<int64_t>(Nr)); });
- OS << "\n";
- return OS;
-}
-
-raw_ostream &Logger::writeOutcome(raw_ostream &OS,
- size_t CurrentObservationID) const {
- if (IncludeReward) {
- OS << "\n";
- json::OStream JOS(OS);
- JOS.object([&]() {
- JOS.attribute("outcome", static_cast<int64_t>(CurrentObservationID));
- });
- OS << "\n";
- OS.write(RewardStorage[CurrentObservationID].get(),
- RewardSpec.getTotalTensorBufferSize());
- }
- OS << "\n";
- return OS;
-}
-
-char *Logger::addNewTensor(size_t FeatureID) {
- return FeatureStorage
- .emplace_back(
- new char[FeatureSpecs[FeatureID].getTotalTensorBufferSize()])
- .get();
-}
-
-size_t Logger::getNrRecords() const {
- assert(FeatureStorage.size() % FeatureSpecs.size() == 0);
- return FeatureStorage.size() / FeatureSpecs.size();
-}
-
-void Logger::logRewardImpl(const char *Value, size_t Size) {
- std::memcpy(RewardStorage.emplace_back(new char[Size]).get(), Value, Size);
-}
-
-raw_ostream &Logger::flush(raw_ostream &OS, bool WithHeader,
- StringRef Context) const {
- if (WithHeader)
- dumpHeader(OS);
- startContext(OS, Context);
- size_t CurrentObservationID = 0;
- for (size_t I = 0; I < FeatureStorage.size(); ++I) {
- size_t TensorID = I % FeatureSpecs.size();
- if (TensorID == 0) {
- CurrentObservationID = I / FeatureSpecs.size();
- startObservation(OS, CurrentObservationID);
- }
- OS.write(FeatureStorage[I].get(),
- FeatureSpecs[TensorID].getTotalTensorBufferSize());
- if (TensorID == FeatureSpecs.size() - 1) {
- writeOutcome(OS, CurrentObservationID);
- }
- }
- return OS;
-}
-
-#define LOG_REWARD(NAME, TYPE) \
- void Logger::log##NAME##Reward(TYPE Value) { \
- assert(IncludeReward); \
- (void)IncludeReward; \
- logReward(Value); \
- }
-
-LOG_REWARD(Float, float)
-LOG_REWARD(Int32, int32_t)
-LOG_REWARD(Int64, int64_t)
-#undef LOG_REWARD
-
-#define LOG_FINAL_REWARD(NAME, TYPE) \
- void Logger::log##NAME##FinalReward(TYPE Value) { \
- assert(RewardSpec.isElementType<TYPE>()); \
- for (size_t I = 1; I < getNrRecords(); ++I) \
- log##NAME##Reward(0); \
- log##NAME##Reward(Value); \
- }
-
-LOG_FINAL_REWARD(Float, float)
-LOG_FINAL_REWARD(Int32, int32_t)
-LOG_FINAL_REWARD(Int64, int64_t)
-#undef LOG_FINAL_REWARD
-
-void Logger::logFloatValue(size_t FeatureID, const float *Value) {
- assert(FeatureSpecs[FeatureID].isElementType<float>());
- logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
-
-void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) {
- assert(FeatureSpecs[FeatureID].isElementType<int64_t>());
- logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
+void Logger::startObservation() {
+ auto I = ObservationIDs.insert({CurrentContext, 0});
+ size_t NewObservationID = I.second ? 0 : ++I.first->second;
+ json::OStream JOS(*OS);
+ JOS.object([&]() {
+ JOS.attribute("observation", static_cast<int64_t>(NewObservationID));
+ });
+ *OS << "\n";
}
-void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) {
- assert(FeatureSpecs[FeatureID].isElementType<int32_t>());
- logSpecifiedTensorValue(FeatureID, reinterpret_cast<const char *>(Value));
-}
+void Logger::endObservation() { *OS << "\n"; }
-void Logger::logSpecifiedTensorValue(size_t FeatureID, const char *RawData) {
- const auto &Spec = FeatureSpecs[FeatureID];
- char *Buff = addEntryAndGetFloatOrInt64Buffer(FeatureID);
- if (Spec.isElementType<int32_t>())
- for (size_t I = 0; I < Spec.getElementCount(); ++I)
- (reinterpret_cast<int64_t *>(Buff))[I] =
- static_cast<int64_t>((reinterpret_cast<const int32_t *>(RawData))[I]);
- else if (Spec.isElementType<int64_t>() || Spec.isElementType<float>())
- std::memcpy(Buff, RawData,
- Spec.getElementCount() * Spec.getElementByteSize());
- else
- llvm_unreachable("Unsupported tensor type");
-}
-
-char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) {
- return reinterpret_cast<char *>(addNewTensor(FeatureID));
+void Logger::logRewardImpl(const char *RawData) {
+ assert(IncludeReward);
+ json::OStream JOS(*OS);
+ JOS.object([&]() {
+ JOS.attribute("outcome", static_cast<int64_t>(
+ ObservationIDs.find(CurrentContext)->second));
+ });
+ *OS << "\n";
+ writeTensor(RewardSpec, RawData);
+ *OS << "\n";
}
-void Logger::flushLogs(raw_ostream &OS,
- const StringMap<std::unique_ptr<Logger>> &Loggers) {
- bool IsFirst = true;
- for (const auto &NamedLogger : Loggers) {
- NamedLogger.second->flush(OS, IsFirst, NamedLogger.first());
- IsFirst = false;
- }
+Logger::Logger(std::unique_ptr<raw_ostream> OS,
+ const std::vector<TensorSpec> &FeatureSpecs,
+ const TensorSpec &RewardSpec, bool IncludeReward)
+ : OS(std::move(OS)), FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec),
+ IncludeReward(IncludeReward) {
+ writeHeader();
}
diff --git a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
index 21a305e25e648..5cc8ad3d609e0 100644
--- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp
@@ -456,19 +456,20 @@ class DevelopmentModeEvictionAdvisorAnalysis final
return R->getAdvisorMode() == AdvisorMode::Development;
}
- /// get the logger for the given function, or nullptr if we didn't collect
- /// one. This is used to inject the score by the RegAllocScoring pass.
- Logger *getLogger(const MachineFunction &MF) const {
- auto I = LogMap.find(MF.getName());
- if (I == LogMap.end())
- return nullptr;
- return I->second.get();
- }
-
void logRewardIfNeeded(const MachineFunction &MF,
llvm::function_ref<float()> GetReward) override {
- if (auto *Log = this->getLogger(MF))
- Log->logFloatFinalReward(GetReward());
+ if (!Log)
+ return;
+ // The function pass manager would run all the function passes for a
+ // function, so we assume the last context belongs to this function. If
+ // this invariant ever changes, we can implement at that time switching
+ // contexts. At this point, it'd be an error
+ if (Log->currentContext() != MF.getName()) {
+ MF.getFunction().getContext().emitError(
+ "The training log context shouldn't have had changed.");
+ }
+ if (Log->hasObservationInProgress())
+ Log->logReward<float>(GetReward());
}
private:
@@ -481,8 +482,22 @@ class DevelopmentModeEvictionAdvisorAnalysis final
RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
}
- // Save all the logs (when requested).
- bool doFinalization(Module &M) override {
+ bool doInitialization(Module &M) override {
+ LLVMContext &Ctx = M.getContext();
+ if (ModelUnderTraining.empty() && TrainingLog.empty()) {
+ Ctx.emitError("Regalloc development mode should be requested with at "
+ "least logging enabled and/or a training model");
+ return false;
+ }
+ if (ModelUnderTraining.empty())
+ Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
+ else
+ Runner = ModelUnderTrainingRunner::createAndEnsureValid(
+ Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
+ if (!Runner) {
+ Ctx.emitError("Regalloc: could not set up the model runner");
+ return false;
+ }
if (TrainingLog.empty())
return false;
std::error_code EC;
@@ -491,52 +506,32 @@ class DevelopmentModeEvictionAdvisorAnalysis final
M.getContext().emitError(EC.message() + ":" + TrainingLog);
return false;
}
- Logger::flushLogs(*OS, LogMap);
+ std::vector<TensorSpec> LFS = InputFeatures;
+ if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
+ append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
+ // We always log the output; in particular, if we're not evaluating, we
+ // don't have an output spec json file. That's why we handle the
+ // 'normal' output separately.
+ LFS.push_back(Output);
+
+ Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+ /*IncludeReward*/ true);
return false;
}
std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
- LLVMContext &Ctx = MF.getFunction().getContext();
- if (ModelUnderTraining.empty() && TrainingLog.empty()) {
- Ctx.emitError("Regalloc development mode should be requested with at "
- "least logging enabled and/or a training model");
+ if (!Runner)
return nullptr;
- }
- if (!Runner) {
- if (ModelUnderTraining.empty())
- Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
- else
- Runner = ModelUnderTrainingRunner::createAndEnsureValid(
- Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
- if (!Runner) {
- Ctx.emitError("Regalloc: could not set up the model runner");
- return nullptr;
- }
- }
-
- Logger *Log = nullptr;
- if (!TrainingLog.empty()) {
- std::vector<TensorSpec> LFS = InputFeatures;
- if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
- append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
- // We always log the output; in particular, if we're not evaluating, we
- // don't have an output spec json file. That's why we handle the
- // 'normal' output separately.
- LFS.push_back(Output);
- auto I = LogMap.insert(std::make_pair(
- MF.getFunction().getName(),
- std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
- assert(I.second);
- Log = I.first->second.get();
- }
+ if (Log)
+ Log->switchContext(MF.getName());
return std::make_unique<DevelopmentModeEvictAdvisor>(
MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
- getAnalysis<MachineLoopInfo>(), Log);
+ getAnalysis<MachineLoopInfo>(), Log.get());
}
std::unique_ptr<MLModelRunner> Runner;
- StringMap<std::unique_ptr<Logger>> LogMap;
+ std::unique_ptr<Logger> Log;
};
#endif //#ifdef LLVM_HAVE_TFLITE
@@ -1092,23 +1087,31 @@ int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
}
if (TrainingLog.empty())
return Ret;
+ // TODO(mtrofin): when we support optional rewards, this can go away. In the
+ // meantime, we log the "pretend" reward (0) for the previous observation
+ // before starting a new one.
+ if (Log->hasObservationInProgress())
+ Log->logReward<float>(0.0);
+
+ Log->startObservation();
size_t CurrentFeature = 0;
size_t FeatureCount = EnableDevelopmentFeatures
? FeatureIDs::FeaturesWithDevelopmentCount
: FeatureIDs::FeatureCount;
for (; CurrentFeature < FeatureCount; ++CurrentFeature) {
- Log->logSpecifiedTensorValue(
- CurrentFeature, reinterpret_cast<const char *>(
+ Log->logTensorValue(CurrentFeature,
+ reinterpret_cast<const char *>(
getRunner().getTensorUntyped(CurrentFeature)));
}
if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
++I, ++CurrentFeature)
- Log->logSpecifiedTensorValue(
+ Log->logTensorValue(
CurrentFeature,
reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
// The output is right after the features and the extra outputs
- Log->logInt64Value(CurrentFeature, &Ret);
+ Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+ Log->endObservation();
return Ret;
}
diff --git a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
index 0d440482ed588..320a184bdcc51 100644
--- a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp
@@ -177,19 +177,20 @@ class DevelopmentModePriorityAdvisorAnalysis final
return R->getAdvisorMode() == AdvisorMode::Development;
}
- /// get the logger for the given function, or nullptr if we didn't collect
- /// one. This is used to inject the score by the RegAllocScoring pass.
- Logger *getLogger(const MachineFunction &MF) const {
- auto I = LogMap.find(MF.getName());
- if (I == LogMap.end())
- return nullptr;
- return I->second.get();
- }
-
void logRewardIfNeeded(const MachineFunction &MF,
llvm::function_ref<float()> GetReward) override {
- if (auto *Log = this->getLogger(MF))
- Log->logFloatFinalReward(GetReward());
+ if (!Log)
+ return;
+ // The function pass manager would run all the function passes for a
+ // function, so we assume the last context belongs to this function. If
+ // this invariant ever changes, we can implement at that time switching
+ // contexts. At this point, it'd be an error
+ if (Log->currentContext() != MF.getName()) {
+ MF.getFunction().getContext().emitError(
+ "The training log context shouldn't have had changed.");
+ }
+ if (Log->hasObservationInProgress())
+ Log->logReward<float>(GetReward());
}
private:
@@ -200,7 +201,22 @@ class DevelopmentModePriorityAdvisorAnalysis final
}
// Save all the logs (when requested).
- bool doFinalization(Module &M) override {
+ bool doInitialization(Module &M) override {
+ LLVMContext &Ctx = M.getContext();
+ if (ModelUnderTraining.empty() && TrainingLog.empty()) {
+ Ctx.emitError("Regalloc development mode should be requested with at "
+ "least logging enabled and/or a training model");
+ return false;
+ }
+ if (ModelUnderTraining.empty())
+ Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
+ else
+ Runner = ModelUnderTrainingRunner::createAndEnsureValid(
+ Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
+ if (!Runner) {
+ Ctx.emitError("Regalloc: could not set up the model runner");
+ return false;
+ }
if (TrainingLog.empty())
return false;
std::error_code EC;
@@ -209,53 +225,33 @@ class DevelopmentModePriorityAdvisorAnalysis final
M.getContext().emitError(EC.message() + ":" + TrainingLog);
return false;
}
- Logger::flushLogs(*OS, LogMap);
+ std::vector<TensorSpec> LFS = InputFeatures;
+ if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
+ append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
+ // We always log the output; in particular, if we're not evaluating, we
+ // don't have an output spec json file. That's why we handle the
+ // 'normal' output separately.
+ LFS.push_back(Output);
+
+ Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
+ /*IncludeReward*/ true);
return false;
}
std::unique_ptr<RegAllocPriorityAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
-
- LLVMContext &Ctx = MF.getFunction().getContext();
- if (ModelUnderTraining.empty() && TrainingLog.empty()) {
- Ctx.emitError("Regalloc development mode should be requested with at "
- "least logging enabled and/or a training model");
+ if (!Runner)
return nullptr;
- }
- if (!Runner) {
- if (ModelUnderTraining.empty())
- Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
- else
- Runner = ModelUnderTrainingRunner::createAndEnsureValid(
- Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
- if (!Runner) {
- Ctx.emitError("Regalloc: could not set up the model runner");
- return nullptr;
- }
- }
-
- Logger *Log = nullptr;
- if (!TrainingLog.empty()) {
- std::vector<TensorSpec> LFS = InputFeatures;
- if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
- append_range(LFS, MUTR->extraOutputsForLoggingSpecs());
- // We always log the output; in particular, if we're not evaluating, we
- // don't have an output spec json file. That's why we handle the
- // 'normal' output separately.
- LFS.push_back(Output);
- auto I = LogMap.insert(std::make_pair(
- MF.getFunction().getName(),
- std::make_unique<Logger>(LFS, Reward, /*IncludeReward*/ true)));
- assert(I.second);
- Log = I.first->second.get();
+ if (Log) {
+ Log->switchContext(MF.getName());
}
return std::make_unique<DevelopmentModePriorityAdvisor>(
- MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log);
+ MF, RA, &getAnalysis<SlotIndexes>(), Runner.get(), Log.get());
}
std::unique_ptr<MLModelRunner> Runner;
- StringMap<std::unique_ptr<Logger>> LogMap;
+ std::unique_ptr<Logger> Log;
};
#endif //#ifdef LLVM_HAVE_TFLITE
@@ -307,23 +303,31 @@ DevelopmentModePriorityAdvisor::getPriority(const LiveInterval &LI) const {
if (TrainingLog.empty())
return Prio;
+ // TODO(mtrofin): when we support optional rewards, this can go away. In the
+ // meantime, we log the "pretend" reward (0) for the previous observation
+ // before starting a new one.
+ if (Log->hasObservationInProgress())
+ Log->logReward<float>(0.0);
+
+ Log->startObservation();
size_t CurrentFeature = 0;
for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) {
- Log->logSpecifiedTensorValue(
- CurrentFeature, reinterpret_cast<const char *>(
+ Log->logTensorValue(CurrentFeature,
+ reinterpret_cast<const char *>(
getRunner().getTensorUntyped(CurrentFeature)));
}
if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner())) {
for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size();
++I, ++CurrentFeature)
- Log->logSpecifiedTensorValue(
+ Log->logTensorValue(
CurrentFeature,
reinterpret_cast<const char *>(MUTR->getUntypedExtraOutputValue(I)));
}
float Ret = static_cast<float>(Prio);
- Log->logFloatValue(CurrentFeature, &Ret);
+ Log->logTensorValue(CurrentFeature, reinterpret_cast<const char *>(&Ret));
+ Log->endObservation();
return static_cast<unsigned>(Prio);
}
diff --git a/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll b/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
index d60f563c2741f..16ff33ffb45fd 100644
--- a/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
+++ b/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll
@@ -36,13 +36,13 @@
; NOML: reward: 36.64
-; CHECK-TWO-FCTS: context: SyFgetsCopy
+; CHECK-TWO-FCTS: context: SyFgets
; CHECK-TWO-FCTS-NEXT: observation: 0
; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
; CHECK-TWO-FCTS: index_to_evict: 12
; CHECK-TWO-FCTS: observation: 16
; CHECK-TWO-FCTS: reward: 36.64
-; CHECK-TWO-FCTS: context: SyFgets
+; CHECK-TWO-FCTS: context: SyFgetsCopy
; CHECK-TWO-FCTS-NEXT: observation: 0
; CHECK-TWO-FCTS-NEXT: mask: 0,0,0,0,0,0,0,0,0,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1
; CHECK-TWO-FCTS: index_to_evict: 12
More information about the llvm-commits
mailing list