From 6d11baf02b33c3ec89178f94627ecec5025dac7e Mon Sep 17 00:00:00 2001 From: Mircea Trofin Date: Thu, 19 Jan 2023 17:37:08 -0800 Subject: [PATCH] [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 --- llvm/include/llvm/Analysis/Utils/TrainingLogger.h | 67 ++++----- llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp | 42 +++--- llvm/lib/Analysis/TrainingLogger.cpp | 160 +++++---------------- llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp | 107 +++++++------- llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp | 106 +++++++------- llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll | 4 +- 6 files changed, 187 insertions(+), 299 deletions(-) diff --git a/llvm/include/llvm/Analysis/Utils/TrainingLogger.h b/llvm/include/llvm/Analysis/Utils/TrainingLogger.h index 7ee25bb..b7db58f 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 +#include #include 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 OS; const std::vector FeatureSpecs; const TensorSpec RewardSpec; const bool IncludeReward; - std::vector> FeatureStorage; - std::vector> 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 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 @@ public: /// 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 &FeatureSpecs, - const TensorSpec &RewardSpec, bool IncludeReward) - : FeatureSpecs(FeatureSpecs), RewardSpec(RewardSpec), - IncludeReward(IncludeReward) {} - - template void logReward(T Value) { - logRewardImpl(reinterpret_cast(&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 OS, + const std::vector &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 void logReward(T Value) { + logRewardImpl(reinterpret_cast(&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> &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 63aa4ca..a91d2ff 100644 --- a/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp +++ b/llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp @@ -102,9 +102,6 @@ public: 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 @@ public: 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(DecisionName, {1})); + std::error_code EC; + auto OS = std::make_unique(TrainingLog, EC); + if (EC) + dbgs() << (EC.message() + ":" + TrainingLog); L = std::make_unique( - FT, TensorSpec::createSpec(RewardName, {1}), + std::move(OS), FT, TensorSpec::createSpec(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(CurrentFeature); - L->logInt64Value(CurrentFeature, &F); - } + for (; CurrentFeature < NumberOfFeatures; ++CurrentFeature) + L->logTensorValue(CurrentFeature, + reinterpret_cast( + ModelRunner.getTensorUntyped(CurrentFeature))); if (MUTR) for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I) { const char *RawData = reinterpret_cast(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(&Event.DefaultDecision)); + L->logTensorValue(DecisionPos, + reinterpret_cast(&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 ModelRunner, @@ -342,11 +341,6 @@ DevelopmentModeMLInlineAdvisor::DevelopmentModeMLInlineAdvisor( assert(IsDoingInference || isLogging()); } -DevelopmentModeMLInlineAdvisor::~DevelopmentModeMLInlineAdvisor() { - if (isLogging()) - Logger->print(); -} - std::optional DevelopmentModeMLInlineAdvisor::getNativeSizeEstimate(const Function &F) const { if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested()) diff --git a/llvm/lib/Analysis/TrainingLogger.cpp b/llvm/lib/Analysis/TrainingLogger.cpp index 72ec14a..dcee8d4 100644 --- a/llvm/lib/Analysis/TrainingLogger.cpp +++ b/llvm/lib/Analysis/TrainingLogger.cpp @@ -32,8 +32,8 @@ static cl::opt 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(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(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()); \ - 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()); - logSpecifiedTensorValue(FeatureID, reinterpret_cast(Value)); -} - -void Logger::logInt64Value(size_t FeatureID, const int64_t *Value) { - assert(FeatureSpecs[FeatureID].isElementType()); - logSpecifiedTensorValue(FeatureID, reinterpret_cast(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(NewObservationID)); + }); + *OS << "\n"; } -void Logger::logInt32Value(size_t FeatureID, const int32_t *Value) { - assert(FeatureSpecs[FeatureID].isElementType()); - logSpecifiedTensorValue(FeatureID, reinterpret_cast(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()) - for (size_t I = 0; I < Spec.getElementCount(); ++I) - (reinterpret_cast(Buff))[I] = - static_cast((reinterpret_cast(RawData))[I]); - else if (Spec.isElementType() || Spec.isElementType()) - std::memcpy(Buff, RawData, - Spec.getElementCount() * Spec.getElementByteSize()); - else - llvm_unreachable("Unsupported tensor type"); -} - -char *Logger::addEntryAndGetFloatOrInt64Buffer(size_t FeatureID) { - return reinterpret_cast(addNewTensor(FeatureID)); +void Logger::logRewardImpl(const char *RawData) { + assert(IncludeReward); + json::OStream JOS(*OS); + JOS.object([&]() { + JOS.attribute("outcome", static_cast( + ObservationIDs.find(CurrentContext)->second)); + }); + *OS << "\n"; + writeTensor(RewardSpec, RawData); + *OS << "\n"; } -void Logger::flushLogs(raw_ostream &OS, - const StringMap> &Loggers) { - bool IsFirst = true; - for (const auto &NamedLogger : Loggers) { - NamedLogger.second->flush(OS, IsFirst, NamedLogger.first()); - IsFirst = false; - } +Logger::Logger(std::unique_ptr OS, + const std::vector &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 21a305e..5cc8ad3 100644 --- a/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp @@ -456,19 +456,20 @@ public: 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 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(GetReward()); } private: @@ -481,8 +482,22 @@ private: 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(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 @@ private: M.getContext().emitError(EC.message() + ":" + TrainingLog); return false; } - Logger::flushLogs(*OS, LogMap); + std::vector LFS = InputFeatures; + if (auto *MUTR = dyn_cast(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(std::move(OS), LFS, Reward, + /*IncludeReward*/ true); return false; } std::unique_ptr 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(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 LFS = InputFeatures; - if (auto *MUTR = dyn_cast(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(LFS, Reward, /*IncludeReward*/ true))); - assert(I.second); - Log = I.first->second.get(); - } + if (Log) + Log->switchContext(MF.getName()); return std::make_unique( MF, RA, Runner.get(), getAnalysis(), - getAnalysis(), Log); + getAnalysis(), Log.get()); } std::unique_ptr Runner; - StringMap> LogMap; + std::unique_ptr 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(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( + Log->logTensorValue(CurrentFeature, + reinterpret_cast( getRunner().getTensorUntyped(CurrentFeature))); } if (auto *MUTR = dyn_cast(&getRunner())) for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I, ++CurrentFeature) - Log->logSpecifiedTensorValue( + Log->logTensorValue( CurrentFeature, reinterpret_cast(MUTR->getUntypedExtraOutputValue(I))); // The output is right after the features and the extra outputs - Log->logInt64Value(CurrentFeature, &Ret); + Log->logTensorValue(CurrentFeature, reinterpret_cast(&Ret)); + Log->endObservation(); return Ret; } diff --git a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp index 0d44048..320a184 100644 --- a/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegallocPriorityAdvisor.cpp @@ -177,19 +177,20 @@ public: 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 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(GetReward()); } private: @@ -200,7 +201,22 @@ private: } // 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(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 @@ private: M.getContext().emitError(EC.message() + ":" + TrainingLog); return false; } - Logger::flushLogs(*OS, LogMap); + std::vector LFS = InputFeatures; + if (auto *MUTR = dyn_cast(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(std::move(OS), LFS, Reward, + /*IncludeReward*/ true); return false; } std::unique_ptr 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(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 LFS = InputFeatures; - if (auto *MUTR = dyn_cast(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(LFS, Reward, /*IncludeReward*/ true))); - assert(I.second); - Log = I.first->second.get(); + if (Log) { + Log->switchContext(MF.getName()); } return std::make_unique( - MF, RA, &getAnalysis(), Runner.get(), Log); + MF, RA, &getAnalysis(), Runner.get(), Log.get()); } std::unique_ptr Runner; - StringMap> LogMap; + std::unique_ptr 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(0.0); + + Log->startObservation(); size_t CurrentFeature = 0; for (; CurrentFeature < InputFeatures.size(); ++CurrentFeature) { - Log->logSpecifiedTensorValue( - CurrentFeature, reinterpret_cast( + Log->logTensorValue(CurrentFeature, + reinterpret_cast( getRunner().getTensorUntyped(CurrentFeature))); } if (auto *MUTR = dyn_cast(&getRunner())) { for (size_t I = 0; I < MUTR->extraOutputsForLoggingSpecs().size(); ++I, ++CurrentFeature) - Log->logSpecifiedTensorValue( + Log->logTensorValue( CurrentFeature, reinterpret_cast(MUTR->getUntypedExtraOutputValue(I))); } float Ret = static_cast(Prio); - Log->logFloatValue(CurrentFeature, &Ret); + Log->logTensorValue(CurrentFeature, reinterpret_cast(&Ret)); + Log->endObservation(); return static_cast(Prio); } diff --git a/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll b/llvm/test/CodeGen/MLRegalloc/dev-mode-logging.ll index d60f563..16ff33f 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 -- 2.7.4