[llvm] [CodeGen][NewPM] Port RegAllocEvictionAdvisor analysis to NPM (PR #117309)
Akshat Oke via llvm-commits
llvm-commits at lists.llvm.org
Fri Nov 29 01:20:39 PST 2024
https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/117309
>From 0bef4ef6160dcaa5dd06b782326d1662037459cb Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 22 Nov 2024 09:31:50 +0000
Subject: [PATCH 1/2] [CodeGen][NewPM] Port RegAllocEvictionAdvisor analysis to
NPM
---
.../llvm}/CodeGen/RegAllocEvictionAdvisor.h | 69 +++++++-
llvm/include/llvm/InitializePasses.h | 2 +-
llvm/include/llvm/Passes/CodeGenPassBuilder.h | 6 +-
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 167 +++++++++++++-----
llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp | 107 ++++++++---
llvm/lib/CodeGen/RegAllocGreedy.cpp | 9 +-
llvm/lib/CodeGen/RegAllocGreedy.h | 1 -
llvm/lib/CodeGen/RegAllocPriorityAdvisor.h | 2 +-
llvm/lib/Passes/PassBuilder.cpp | 1 +
llvm/lib/Passes/PassRegistry.def | 1 +
10 files changed, 284 insertions(+), 81 deletions(-)
rename llvm/{lib => include/llvm}/CodeGen/RegAllocEvictionAdvisor.h (75%)
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
similarity index 75%
rename from llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
rename to llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 52dd946a685400a..847bf032235c1dd 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -9,11 +9,13 @@
#ifndef LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H
#define LLVM_CODEGEN_REGALLOCEVICTIONADVISOR_H
+#include "llvm/ADT/Any.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/CodeGen/Register.h"
#include "llvm/Config/llvm-config.h"
+#include "llvm/IR/PassManager.h"
#include "llvm/MC/MCRegister.h"
#include "llvm/Pass.h"
@@ -164,12 +166,12 @@ class RegAllocEvictionAdvisor {
///
/// Because we need to offer additional services in 'development' mode, the
/// implementations of this analysis need to implement RTTI support.
-class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
+class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
public:
enum class AdvisorMode : int { Default, Release, Development };
- RegAllocEvictionAdvisorAnalysis(AdvisorMode Mode)
- : ImmutablePass(ID), Mode(Mode){};
+ RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode Mode)
+ : ImmutablePass(ID), Mode(Mode) {};
static char ID;
/// Get an advisor for the given context (i.e. machine function, etc)
@@ -177,7 +179,7 @@ class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
AdvisorMode getAdvisorMode() const { return Mode; }
virtual void logRewardIfNeeded(const MachineFunction &MF,
- llvm::function_ref<float()> GetReward){};
+ llvm::function_ref<float()> GetReward) {};
protected:
// This analysis preserves everything, and subclasses may have additional
@@ -191,13 +193,66 @@ class RegAllocEvictionAdvisorAnalysis : public ImmutablePass {
const AdvisorMode Mode;
};
+/// Common provider for legacy and new pass managers.
+/// This keeps the state for logging, and sets up and holds the provider.
+/// The legacy pass itself used to keep the logging state and provider,
+/// so this extraction helps the NPM analysis to reuse the logic.
+class RegAllocEvictionAdvisorProvider {
+public:
+ enum class AdvisorMode : int { Default, Release, Development };
+ RegAllocEvictionAdvisorProvider(AdvisorMode Mode) : Mode(Mode) {}
+
+ virtual ~RegAllocEvictionAdvisorProvider() = default;
+
+ virtual bool doInitialization(Module &M) { return false; }
+
+ virtual void logRewardIfNeeded(const MachineFunction &MF,
+ llvm::function_ref<float()> GetReward) {}
+
+ virtual std::unique_ptr<RegAllocEvictionAdvisor>
+ getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
+
+ /// Set the analyses that the advisor needs to use as they might not be
+ /// available before the advisor is created.
+ virtual void setAnalyses(std::initializer_list<llvm::Any> AnalysisP) {}
+
+ AdvisorMode getAdvisorMode() const { return Mode; }
+
+private:
+ const AdvisorMode Mode;
+};
+
+RegAllocEvictionAdvisorProvider *createReleaseModeAdvisorProvider();
+RegAllocEvictionAdvisorProvider *createDevelopmentModeAdvisorProvider();
+
+/// A Module analysis for fetching the Eviction Advisor. This is not a
+/// MachineFunction analysis for two reasons:
+/// - in the ML implementation case, the evaluator is stateless but (especially
+/// in the development mode) expensive to set up. With a Module Analysis, we
+/// `require` it and set it up once.
+/// - in the 'development' mode ML case, we want to capture the training log
+/// during allocation (this is a log of features encountered and decisions
+/// made), and then measure a score, potentially a few steps after allocation
+/// completes. So we need a Module analysis to keep the logger state around
+/// until we can make that measurement.
+class RegAllocEvictionAdvisorAnalysis
+ : public AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis> {
+ static AnalysisKey Key;
+ friend AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis>;
+
+public:
+ using Result = std::unique_ptr<RegAllocEvictionAdvisorProvider>;
+ Result run(Module &MF, ModuleAnalysisManager &MAM);
+};
+
/// Specialization for the API used by the analysis infrastructure to create
/// an instance of the eviction advisor.
-template <> Pass *callDefaultCtor<RegAllocEvictionAdvisorAnalysis>();
+template <> Pass *callDefaultCtor<RegAllocEvictionAdvisorAnalysisLegacy>();
-RegAllocEvictionAdvisorAnalysis *createReleaseModeAdvisor();
+RegAllocEvictionAdvisorAnalysisLegacy *createReleaseModeAdvisorAnalysisLegacy();
-RegAllocEvictionAdvisorAnalysis *createDevelopmentModeAdvisor();
+RegAllocEvictionAdvisorAnalysisLegacy *
+createDevelopmentModeAdvisorAnalysisLegacy();
// TODO: move to RegAllocEvictionAdvisor.cpp when we move implementation
// out of RegAllocGreedy.cpp
diff --git a/llvm/include/llvm/InitializePasses.h b/llvm/include/llvm/InitializePasses.h
index e883aae2758688d..75212cb2a777ab8 100644
--- a/llvm/include/llvm/InitializePasses.h
+++ b/llvm/include/llvm/InitializePasses.h
@@ -252,7 +252,7 @@ void initializePseudoProbeInserterPass(PassRegistry &);
void initializeRAGreedyPass(PassRegistry &);
void initializeReachingDefAnalysisPass(PassRegistry &);
void initializeReassociateLegacyPassPass(PassRegistry &);
-void initializeRegAllocEvictionAdvisorAnalysisPass(PassRegistry &);
+void initializeRegAllocEvictionAdvisorAnalysisLegacyPass(PassRegistry &);
void initializeRegAllocFastPass(PassRegistry &);
void initializeRegAllocPriorityAdvisorAnalysisPass(PassRegistry &);
void initializeRegAllocScoringPass(PassRegistry &);
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index d2e9e8185a2b90c..aa6437c1d2f0172 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -54,6 +54,7 @@
#include "llvm/CodeGen/PHIElimination.h"
#include "llvm/CodeGen/PeepholeOptimizer.h"
#include "llvm/CodeGen/PreISelIntrinsicLowering.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "llvm/CodeGen/RegAllocFast.h"
#include "llvm/CodeGen/RegUsageInfoCollector.h"
#include "llvm/CodeGen/RegUsageInfoPropagate.h"
@@ -1061,9 +1062,10 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
template <typename Derived, typename TargetMachineT>
void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
AddMachinePass &addPass, bool Optimized) const {
- if (Optimized)
+ if (Optimized) {
+ addPass(RequireAnalysis<RegAllocEvictionAdvisorAnalysis>());
addPass(RAGreedyPass());
- else
+ } else
addPass(RegAllocFastPass());
}
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index d099544c2a49187..9c15c327bb202c2 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -11,11 +11,11 @@
//===----------------------------------------------------------------------===//
#include "AllocationOrder.h"
-#include "RegAllocEvictionAdvisor.h"
#include "RegAllocGreedy.h"
#include "llvm/Analysis/InteractiveModelRunner.h"
#include "llvm/Analysis/MLModelRunner.h"
#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#if defined(LLVM_HAVE_TF_AOT_REGALLOCEVICTMODEL) || defined(LLVM_HAVE_TFLITE)
#include "llvm/Analysis/ModelUnderTrainingRunner.h"
#include "llvm/Analysis/NoInferenceModelRunner.h"
@@ -108,7 +108,7 @@ class RegAllocScoring : public MachineFunctionPass {
/// RegAllocReward analysis usage.
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.setPreservesAll();
- AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+ AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
MachineFunctionPass::getAnalysisUsage(AU);
@@ -366,11 +366,12 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
// ===================================
// Release (AOT) - specifics
// ===================================
-class ReleaseModeEvictionAdvisorAnalysis final
- : public RegAllocEvictionAdvisorAnalysis {
+/// Common provider for legacy and new pass managers.
+class ReleaseModeEvictionAdvisorProvider final
+ : public RegAllocEvictionAdvisorProvider {
public:
- ReleaseModeEvictionAdvisorAnalysis()
- : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {
+ ReleaseModeEvictionAdvisorProvider()
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Release) {
if (EnableDevelopmentFeatures) {
InputFeatures = {RA_EVICT_FEATURES_LIST(
_DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -380,17 +381,17 @@ class ReleaseModeEvictionAdvisorAnalysis final
}
}
// support for isa<> and dyn_cast.
- static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+ static bool classof(const RegAllocEvictionAdvisorProvider *R) {
return R->getAdvisorMode() == AdvisorMode::Release;
}
-private:
- std::vector<TensorSpec> InputFeatures;
-
- void getAnalysisUsage(AnalysisUsage &AU) const override {
- AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
- AU.addRequired<MachineLoopInfoWrapperPass>();
- RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
+ void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
+ for (auto Analysis : AnalysisList) {
+ if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
+ this->MBFI = *MBFI;
+ if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
+ this->Loops = *Loops;
+ }
}
std::unique_ptr<RegAllocEvictionAdvisor>
@@ -405,12 +406,47 @@ class ReleaseModeEvictionAdvisorAnalysis final
InteractiveChannelBaseName + ".out",
InteractiveChannelBaseName + ".in");
}
- return std::make_unique<MLEvictAdvisor>(
- MF, RA, Runner.get(),
- getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(),
- getAnalysis<MachineLoopInfoWrapperPass>().getLI());
+ return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI,
+ *Loops);
}
+
+private:
+ std::vector<TensorSpec> InputFeatures;
std::unique_ptr<MLModelRunner> Runner;
+ MachineBlockFrequencyInfo *MBFI;
+ MachineLoopInfo *Loops;
+};
+
+class ReleaseModeEvictionAdvisorAnalysisLegacy final
+ : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+ ReleaseModeEvictionAdvisorAnalysisLegacy()
+ : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {}
+
+ std::unique_ptr<RegAllocEvictionAdvisor>
+ getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+ auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+ auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+ Provider.setAnalyses({MBFI, Loops});
+ return Provider.getAdvisor(MF, RA);
+ }
+
+ bool doInitialization(Module &M) override {
+ return Provider.doInitialization(M);
+ }
+
+ static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+ return R->getAdvisorMode() == AdvisorMode::Release;
+ }
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+ AU.addRequired<MachineLoopInfoWrapperPass>();
+ RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+ }
+
+private:
+ ReleaseModeEvictionAdvisorProvider Provider;
};
// ===================================
@@ -445,11 +481,14 @@ class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
Logger *const Log;
};
-class DevelopmentModeEvictionAdvisorAnalysis final
- : public RegAllocEvictionAdvisorAnalysis {
+class DevelopmentModeEvictionAdvisorProvider final
+ : public RegAllocEvictionAdvisorProvider {
public:
- DevelopmentModeEvictionAdvisorAnalysis()
- : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Development) {
+ DevelopmentModeEvictionAdvisorProvider(
+ MachineBlockFrequencyInfo *MBFI = nullptr,
+ MachineLoopInfo *Loops = nullptr)
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Development), MBFI(MBFI),
+ Loops(Loops) {
if (EnableDevelopmentFeatures) {
InputFeatures = {RA_EVICT_FEATURES_LIST(
_DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -471,7 +510,7 @@ class DevelopmentModeEvictionAdvisorAnalysis final
}
}
// support for isa<> and dyn_cast.
- static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+ static bool classof(const RegAllocEvictionAdvisorProvider *R) {
return R->getAdvisorMode() == AdvisorMode::Development;
}
@@ -491,14 +530,13 @@ class DevelopmentModeEvictionAdvisorAnalysis final
Log->logReward<float>(GetReward());
}
-private:
- std::vector<TensorSpec> InputFeatures;
- std::vector<TensorSpec> TrainingInputFeatures;
-
- void getAnalysisUsage(AnalysisUsage &AU) const override {
- AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
- AU.addRequired<MachineLoopInfoWrapperPass>();
- RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
+ void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
+ for (auto Analysis : AnalysisList) {
+ if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
+ this->MBFI = *MBFI;
+ if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
+ this->Loops = *Loops;
+ }
}
bool doInitialization(Module &M) override {
@@ -544,14 +582,53 @@ class DevelopmentModeEvictionAdvisorAnalysis final
return nullptr;
if (Log)
Log->switchContext(MF.getName());
+ assert((MBFI && Loops) &&
+ "Invalid provider state: must have analysis available");
return std::make_unique<DevelopmentModeEvictAdvisor>(
- MF, RA, Runner.get(),
- getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(),
- getAnalysis<MachineLoopInfoWrapperPass>().getLI(), Log.get());
+ MF, RA, Runner.get(), *MBFI, *Loops, Log.get());
}
+private:
+ std::vector<TensorSpec> InputFeatures;
+ std::vector<TensorSpec> TrainingInputFeatures;
+
std::unique_ptr<MLModelRunner> Runner;
std::unique_ptr<Logger> Log;
+ const MachineBlockFrequencyInfo *MBFI;
+ const MachineLoopInfo *Loops;
+};
+
+class DevelopmentModeEvictionAdvisorAnalysisLegacy final
+ : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+ DevelopmentModeEvictionAdvisorAnalysisLegacy()
+ : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Development) {}
+
+ bool doInitialization(Module &M) override {
+ auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+ auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+ Provider.setAnalyses({MBFI, Loops});
+ return Provider.doInitialization(M);
+ }
+
+ std::unique_ptr<RegAllocEvictionAdvisor>
+ getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+ return Provider.getAdvisor(MF, RA);
+ }
+
+ // support for isa<> and dyn_cast.
+ static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+ return R->getAdvisorMode() == AdvisorMode::Development;
+ }
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
+ AU.addRequired<MachineLoopInfoWrapperPass>();
+ RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU);
+ }
+
+private:
+ DevelopmentModeEvictionAdvisorProvider Provider;
};
#endif //#ifdef LLVM_HAVE_TFLITE
@@ -1079,8 +1156,13 @@ void llvm::extractMBBFrequency(
// Development mode-specific implementations
#ifdef LLVM_HAVE_TFLITE
-RegAllocEvictionAdvisorAnalysis *llvm::createDevelopmentModeAdvisor() {
- return new DevelopmentModeEvictionAdvisorAnalysis();
+RegAllocEvictionAdvisorProvider *llvm::createDevelopmentModeAdvisorProvider() {
+ return new DevelopmentModeEvictionAdvisorProvider();
+}
+
+RegAllocEvictionAdvisorAnalysisLegacy *
+llvm::createDevelopmentModeAdvisorAnalysisLegacy() {
+ return new DevelopmentModeEvictionAdvisorAnalysisLegacy();
}
int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
@@ -1146,18 +1228,23 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
return *CachedReward;
};
- getAnalysis<RegAllocEvictionAdvisorAnalysis>().logRewardIfNeeded(MF,
- GetReward);
+ getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().logRewardIfNeeded(
+ MF, GetReward);
getAnalysis<RegAllocPriorityAdvisorAnalysis>().logRewardIfNeeded(MF,
GetReward);
return false;
}
#endif // #ifdef LLVM_HAVE_TFLITE
-RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
+RegAllocEvictionAdvisorProvider *llvm::createReleaseModeAdvisorProvider() {
+ return new ReleaseModeEvictionAdvisorProvider();
+}
+
+RegAllocEvictionAdvisorAnalysisLegacy *
+llvm::createReleaseModeAdvisorAnalysisLegacy() {
return llvm::isEmbeddedModelEvaluatorValid<CompiledModelType>() ||
!InteractiveChannelBaseName.empty()
- ? new ReleaseModeEvictionAdvisorAnalysis()
+ ? new ReleaseModeEvictionAdvisorAnalysisLegacy()
: nullptr;
}
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index a1f441ebd0d5e41..09fb5df259a97e9 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -10,9 +10,10 @@
//
//===----------------------------------------------------------------------===//
-#include "RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "AllocationOrder.h"
#include "RegAllocGreedy.h"
+#include "RegAllocPriorityAdvisor.h"
#include "llvm/CodeGen/LiveRegMatrix.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
@@ -26,17 +27,18 @@
using namespace llvm;
-static cl::opt<RegAllocEvictionAdvisorAnalysis::AdvisorMode> Mode(
+static cl::opt<RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode> Mode(
"regalloc-enable-advisor", cl::Hidden,
- cl::init(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default),
+ cl::init(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default),
cl::desc("Enable regalloc advisor mode"),
cl::values(
- clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default,
+ clEnumValN(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default,
"default", "Default"),
- clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Release,
+ clEnumValN(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release,
"release", "precompiled"),
- clEnumValN(RegAllocEvictionAdvisorAnalysis::AdvisorMode::Development,
- "development", "for training")));
+ clEnumValN(
+ RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development,
+ "development", "for training")));
static cl::opt<bool> EnableLocalReassignment(
"enable-local-reassign", cl::Hidden,
@@ -59,59 +61,114 @@ cl::opt<unsigned> EvictInterferenceCutoff(
#define LLVM_HAVE_TF_AOT
#endif
-char RegAllocEvictionAdvisorAnalysis::ID = 0;
-INITIALIZE_PASS(RegAllocEvictionAdvisorAnalysis, "regalloc-evict",
+char RegAllocEvictionAdvisorAnalysisLegacy::ID = 0;
+INITIALIZE_PASS(RegAllocEvictionAdvisorAnalysisLegacy, "regalloc-evict",
"Regalloc eviction policy", false, true)
namespace {
-class DefaultEvictionAdvisorAnalysis final
- : public RegAllocEvictionAdvisorAnalysis {
+class DefaultEvictionAdvisorProvider final
+ : public RegAllocEvictionAdvisorProvider {
public:
- DefaultEvictionAdvisorAnalysis(bool NotAsRequested)
- : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Default),
+ DefaultEvictionAdvisorProvider(bool NotAsRequested)
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Default),
NotAsRequested(NotAsRequested) {}
// support for isa<> and dyn_cast.
- static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
+ static bool classof(const RegAllocEvictionAdvisorProvider *R) {
return R->getAdvisorMode() == AdvisorMode::Default;
}
-private:
std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
return std::make_unique<DefaultEvictionAdvisor>(MF, RA);
}
+
bool doInitialization(Module &M) override {
if (NotAsRequested)
M.getContext().emitError("Requested regalloc eviction advisor analysis "
"could not be created. Using default");
- return RegAllocEvictionAdvisorAnalysis::doInitialization(M);
+ return RegAllocEvictionAdvisorProvider::doInitialization(M);
+ }
+
+private:
+ const bool NotAsRequested;
+};
+
+class DefaultEvictionAdvisorAnalysisLegacy final
+ : public RegAllocEvictionAdvisorAnalysisLegacy {
+public:
+ DefaultEvictionAdvisorAnalysisLegacy(bool NotAsRequested)
+ : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default),
+ NotAsRequested(NotAsRequested) {}
+
+ std::unique_ptr<RegAllocEvictionAdvisor>
+ getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+ return Provider->getAdvisor(MF, RA);
+ }
+
+ bool doInitialization(Module &M) override {
+ Provider.reset(new DefaultEvictionAdvisorProvider(NotAsRequested));
+ return Provider->doInitialization(M);
}
+
+ // support for isa<> and dyn_cast.
+ static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
+ return R->getAdvisorMode() == AdvisorMode::Default;
+ }
+
+private:
+ std::unique_ptr<DefaultEvictionAdvisorProvider> Provider;
const bool NotAsRequested;
};
} // namespace
-template <> Pass *llvm::callDefaultCtor<RegAllocEvictionAdvisorAnalysis>() {
+AnalysisKey RegAllocEvictionAdvisorAnalysis::Key;
+
+RegAllocEvictionAdvisorAnalysis::Result
+RegAllocEvictionAdvisorAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
+ std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
+ switch (Mode) {
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
+ Provider.reset(
+ new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false));
+ break;
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
+#if defined(LLVM_HAVE_TFLITE)
+ Provider.reset(createDevelopmentModeAdvisorProvider());
+#endif
+ break;
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
+ Provider.reset(createReleaseModeAdvisorProvider());
+ break;
+ }
+ if (!Provider)
+ Provider.reset(new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true));
+ Provider->doInitialization(M);
+ return Provider;
+}
+
+template <>
+Pass *llvm::callDefaultCtor<RegAllocEvictionAdvisorAnalysisLegacy>() {
Pass *Ret = nullptr;
switch (Mode) {
- case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Default:
- Ret = new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ false);
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
+ Ret = new DefaultEvictionAdvisorAnalysisLegacy(/*NotAsRequested*/ false);
break;
- case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Development:
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
#if defined(LLVM_HAVE_TFLITE)
- Ret = createDevelopmentModeAdvisor();
+ Ret = createDevelopmentModeAdvisorAnalysisLegacy();
#endif
break;
- case RegAllocEvictionAdvisorAnalysis::AdvisorMode::Release:
- Ret = createReleaseModeAdvisor();
+ case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
+ Ret = createReleaseModeAdvisorAnalysisLegacy();
break;
}
if (Ret)
return Ret;
- return new DefaultEvictionAdvisorAnalysis(/*NotAsRequested*/ true);
+ return new DefaultEvictionAdvisorAnalysisLegacy(/*NotAsRequested*/ true);
}
-StringRef RegAllocEvictionAdvisorAnalysis::getPassName() const {
+StringRef RegAllocEvictionAdvisorAnalysisLegacy::getPassName() const {
switch (getAdvisorMode()) {
case AdvisorMode::Default:
return "Default Regalloc Eviction Advisor";
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp
index 3542bfe18af46f9..8f9a7b46691f0d7 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.cpp
+++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp
@@ -15,7 +15,6 @@
#include "AllocationOrder.h"
#include "InterferenceCache.h"
#include "RegAllocBase.h"
-#include "RegAllocEvictionAdvisor.h"
#include "RegAllocPriorityAdvisor.h"
#include "SpillPlacement.h"
#include "SplitKit.h"
@@ -47,6 +46,7 @@
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "llvm/CodeGen/RegAllocRegistry.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
#include "llvm/CodeGen/SlotIndexes.h"
@@ -164,7 +164,7 @@ INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
INITIALIZE_PASS_DEPENDENCY(EdgeBundlesWrapperLegacy)
INITIALIZE_PASS_DEPENDENCY(SpillPlacement)
INITIALIZE_PASS_DEPENDENCY(MachineOptimizationRemarkEmitterPass)
-INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysis)
+INITIALIZE_PASS_DEPENDENCY(RegAllocEvictionAdvisorAnalysisLegacy)
INITIALIZE_PASS_DEPENDENCY(RegAllocPriorityAdvisorAnalysis)
INITIALIZE_PASS_END(RAGreedy, "greedy",
"Greedy Register Allocator", false, false)
@@ -219,7 +219,7 @@ void RAGreedy::getAnalysisUsage(AnalysisUsage &AU) const {
AU.addRequired<EdgeBundlesWrapperLegacy>();
AU.addRequired<SpillPlacement>();
AU.addRequired<MachineOptimizationRemarkEmitterPass>();
- AU.addRequired<RegAllocEvictionAdvisorAnalysis>();
+ AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>();
AU.addRequired<RegAllocPriorityAdvisorAnalysis>();
MachineFunctionPass::getAnalysisUsage(AU);
}
@@ -2748,7 +2748,8 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) {
ExtraInfo.emplace();
EvictAdvisor =
- getAnalysis<RegAllocEvictionAdvisorAnalysis>().getAdvisor(*MF, *this);
+ getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getAdvisor(*MF,
+ *this);
PriorityAdvisor =
getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this);
diff --git a/llvm/lib/CodeGen/RegAllocGreedy.h b/llvm/lib/CodeGen/RegAllocGreedy.h
index 2e7608a53e9cec1..62752a7962db804 100644
--- a/llvm/lib/CodeGen/RegAllocGreedy.h
+++ b/llvm/lib/CodeGen/RegAllocGreedy.h
@@ -14,7 +14,6 @@
#include "InterferenceCache.h"
#include "RegAllocBase.h"
-#include "RegAllocEvictionAdvisor.h"
#include "RegAllocPriorityAdvisor.h"
#include "SpillPlacement.h"
#include "SplitKit.h"
diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
index 1e9fa967214cc87..2d42a43c4c63724 100644
--- a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
+++ b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h
@@ -9,7 +9,7 @@
#ifndef LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
#define LLVM_CODEGEN_REGALLOCPRIORITYADVISOR_H
-#include "RegAllocEvictionAdvisor.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "llvm/CodeGen/SlotIndexes.h"
#include "llvm/Pass.h"
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index cf7ceed63607a67..af2c2e0058d1c7b 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -121,6 +121,7 @@
#include "llvm/CodeGen/PHIElimination.h"
#include "llvm/CodeGen/PeepholeOptimizer.h"
#include "llvm/CodeGen/PreISelIntrinsicLowering.h"
+#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "llvm/CodeGen/RegAllocFast.h"
#include "llvm/CodeGen/RegUsageInfoCollector.h"
#include "llvm/CodeGen/RegUsageInfoPropagate.h"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 7c3798f6462a46d..2a2d2ad674680b7 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -32,6 +32,7 @@ MODULE_ANALYSIS("no-op-module", NoOpModuleAnalysis())
MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
+MODULE_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
MODULE_ANALYSIS("verify", VerifierAnalysis())
>From 6faee3f719576c5dbd448a246aece2eae22db9f6 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 29 Nov 2024 09:18:59 +0000
Subject: [PATCH 2/2] Remove Provider::doInit, change to MF analysis
init is in the constructor. The new MF analysis lazily intiailizes the
provider. The (wrapped) provider is the analysis result.
---
.../llvm/CodeGen/RegAllocEvictionAdvisor.h | 58 ++++++--
llvm/include/llvm/Passes/CodeGenPassBuilder.h | 5 +-
.../llvm/Passes/MachinePassRegistry.def | 1 +
llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 136 ++++++++----------
llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp | 56 ++++----
llvm/lib/Passes/PassRegistry.def | 1 -
6 files changed, 143 insertions(+), 114 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
index 847bf032235c1dd..8261f6d1cf0bbde 100644
--- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
+++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h
@@ -13,6 +13,8 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/Register.h"
#include "llvm/Config/llvm-config.h"
#include "llvm/IR/PassManager.h"
@@ -200,33 +202,43 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass {
class RegAllocEvictionAdvisorProvider {
public:
enum class AdvisorMode : int { Default, Release, Development };
- RegAllocEvictionAdvisorProvider(AdvisorMode Mode) : Mode(Mode) {}
+ RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx)
+ : Ctx(Ctx), Mode(Mode) {}
virtual ~RegAllocEvictionAdvisorProvider() = default;
- virtual bool doInitialization(Module &M) { return false; }
-
virtual void logRewardIfNeeded(const MachineFunction &MF,
llvm::function_ref<float()> GetReward) {}
virtual std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0;
- /// Set the analyses that the advisor needs to use as they might not be
- /// available before the advisor is created.
- virtual void setAnalyses(std::initializer_list<llvm::Any> AnalysisP) {}
+ /// We create this provider in doInitialization which doesn't have these
+ /// analyses. For NPM, we do have them in run(MachineFunction&)
+ virtual void setAnalyses(MachineBlockFrequencyInfo *MBFI,
+ MachineLoopInfo *Loops) {
+ this->MBFI = MBFI;
+ this->Loops = Loops;
+ }
AdvisorMode getAdvisorMode() const { return Mode; }
+protected:
+ LLVMContext &Ctx;
+ MachineBlockFrequencyInfo *MBFI;
+ MachineLoopInfo *Loops;
+
private:
const AdvisorMode Mode;
};
-RegAllocEvictionAdvisorProvider *createReleaseModeAdvisorProvider();
-RegAllocEvictionAdvisorProvider *createDevelopmentModeAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+createReleaseModeAdvisorProvider(LLVMContext &Ctx);
+RegAllocEvictionAdvisorProvider *
+createDevelopmentModeAdvisorProvider(LLVMContext &Ctx);
-/// A Module analysis for fetching the Eviction Advisor. This is not a
-/// MachineFunction analysis for two reasons:
+/// A MachineFunction analysis for fetching the Eviction Advisor.
+/// This sets up the Provider lazily and caches it.
/// - in the ML implementation case, the evaluator is stateless but (especially
/// in the development mode) expensive to set up. With a Module Analysis, we
/// `require` it and set it up once.
@@ -241,8 +253,30 @@ class RegAllocEvictionAdvisorAnalysis
friend AnalysisInfoMixin<RegAllocEvictionAdvisorAnalysis>;
public:
- using Result = std::unique_ptr<RegAllocEvictionAdvisorProvider>;
- Result run(Module &MF, ModuleAnalysisManager &MAM);
+ struct Result {
+ // owned by this analysis
+ RegAllocEvictionAdvisorProvider *Provider;
+
+ bool invalidate(MachineFunction &MF, const PreservedAnalyses &PA,
+ MachineFunctionAnalysisManager::Invalidator &Inv) {
+ auto PAC = PA.getChecker<RegAllocEvictionAdvisorAnalysis>();
+ // If we are in default mode, the provider is always valid.
+ if (Provider->getAdvisorMode() ==
+ RegAllocEvictionAdvisorProvider::AdvisorMode::Default)
+ return !PAC.preservedWhenStateless();
+ // MBFI and Loops are used in release and development modes, so check
+ // those.
+ return !PAC.preservedWhenStateless() ||
+ Inv.invalidate<MachineBlockFrequencyAnalysis>(MF, PA) ||
+ Inv.invalidate<MachineLoopAnalysis>(MF, PA);
+ }
+ };
+
+ Result run(MachineFunction &MF, MachineFunctionAnalysisManager &MAM);
+
+private:
+ void initializeProvider(LLVMContext &Ctx);
+ std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
};
/// Specialization for the API used by the analysis infrastructure to create
diff --git a/llvm/include/llvm/Passes/CodeGenPassBuilder.h b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
index aa6437c1d2f0172..bcd335f2f99427e 100644
--- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h
+++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h
@@ -1062,10 +1062,9 @@ void CodeGenPassBuilder<Derived, TargetMachineT>::addMachineSSAOptimization(
template <typename Derived, typename TargetMachineT>
void CodeGenPassBuilder<Derived, TargetMachineT>::addTargetRegisterAllocator(
AddMachinePass &addPass, bool Optimized) const {
- if (Optimized) {
- addPass(RequireAnalysis<RegAllocEvictionAdvisorAnalysis>());
+ if (Optimized)
addPass(RAGreedyPass());
- } else
+ else
addPass(RegAllocFastPass());
}
diff --git a/llvm/include/llvm/Passes/MachinePassRegistry.def b/llvm/include/llvm/Passes/MachinePassRegistry.def
index cb1c295d824787b..e84f2d19f2bd621 100644
--- a/llvm/include/llvm/Passes/MachinePassRegistry.def
+++ b/llvm/include/llvm/Passes/MachinePassRegistry.def
@@ -112,6 +112,7 @@ MACHINE_FUNCTION_ANALYSIS("machine-post-dom-tree",
MachinePostDominatorTreeAnalysis())
MACHINE_FUNCTION_ANALYSIS("machine-trace-metrics", MachineTraceMetricsAnalysis())
MACHINE_FUNCTION_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
+MACHINE_FUNCTION_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
MACHINE_FUNCTION_ANALYSIS("slot-indexes", SlotIndexesAnalysis())
MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis())
// MACHINE_FUNCTION_ANALYSIS("live-stacks", LiveStacksPass())
diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
index 9c15c327bb202c2..76011b0a6e4bfc5 100644
--- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
+++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
@@ -370,8 +370,8 @@ class MLEvictAdvisor : public RegAllocEvictionAdvisor {
class ReleaseModeEvictionAdvisorProvider final
: public RegAllocEvictionAdvisorProvider {
public:
- ReleaseModeEvictionAdvisorProvider()
- : RegAllocEvictionAdvisorProvider(AdvisorMode::Release) {
+ ReleaseModeEvictionAdvisorProvider(LLVMContext &Ctx)
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Release, Ctx) {
if (EnableDevelopmentFeatures) {
InputFeatures = {RA_EVICT_FEATURES_LIST(
_DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -385,15 +385,6 @@ class ReleaseModeEvictionAdvisorProvider final
return R->getAdvisorMode() == AdvisorMode::Release;
}
- void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
- for (auto Analysis : AnalysisList) {
- if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
- this->MBFI = *MBFI;
- if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
- this->Loops = *Loops;
- }
- }
-
std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
if (!Runner) {
@@ -406,6 +397,8 @@ class ReleaseModeEvictionAdvisorProvider final
InteractiveChannelBaseName + ".out",
InteractiveChannelBaseName + ".in");
}
+ assert((MBFI && Loops) &&
+ "Invalid provider state: must have analysis available");
return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI,
*Loops);
}
@@ -413,8 +406,6 @@ class ReleaseModeEvictionAdvisorProvider final
private:
std::vector<TensorSpec> InputFeatures;
std::unique_ptr<MLModelRunner> Runner;
- MachineBlockFrequencyInfo *MBFI;
- MachineLoopInfo *Loops;
};
class ReleaseModeEvictionAdvisorAnalysisLegacy final
@@ -427,12 +418,19 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
- Provider.setAnalyses({MBFI, Loops});
- return Provider.getAdvisor(MF, RA);
+ Provider->setAnalyses(MBFI, Loops);
+ return Provider->getAdvisor(MF, RA);
+ }
+
+ void logRewardIfNeeded(const MachineFunction &MF,
+ llvm::function_ref<float()> GetReward) override {
+ // No-op in release mode
}
bool doInitialization(Module &M) override {
- return Provider.doInitialization(M);
+ Provider =
+ std::make_unique<ReleaseModeEvictionAdvisorProvider>(M.getContext());
+ return false;
}
static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) {
@@ -446,7 +444,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final
}
private:
- ReleaseModeEvictionAdvisorProvider Provider;
+ std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider;
};
// ===================================
@@ -484,11 +482,8 @@ class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
class DevelopmentModeEvictionAdvisorProvider final
: public RegAllocEvictionAdvisorProvider {
public:
- DevelopmentModeEvictionAdvisorProvider(
- MachineBlockFrequencyInfo *MBFI = nullptr,
- MachineLoopInfo *Loops = nullptr)
- : RegAllocEvictionAdvisorProvider(AdvisorMode::Development), MBFI(MBFI),
- Loops(Loops) {
+ DevelopmentModeEvictionAdvisorProvider(LLVMContext &Ctx)
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Development, Ctx) {
if (EnableDevelopmentFeatures) {
InputFeatures = {RA_EVICT_FEATURES_LIST(
_DECL_FEATURES) RA_EVICT_FIRST_DEVELOPMENT_FEATURE(_DECL_FEATURES)
@@ -508,43 +503,10 @@ class DevelopmentModeEvictionAdvisorProvider final
TensorSpec::createSpec<int32_t>("action_step_type", {1}),
TensorSpec::createSpec<float>("action_reward", {1})};
}
- }
- // support for isa<> and dyn_cast.
- static bool classof(const RegAllocEvictionAdvisorProvider *R) {
- return R->getAdvisorMode() == AdvisorMode::Development;
- }
-
- void logRewardIfNeeded(const MachineFunction &MF,
- llvm::function_ref<float()> GetReward) override {
- if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
- 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());
- }
-
- void setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
- for (auto Analysis : AnalysisList) {
- if (auto **MBFI = llvm::any_cast<MachineBlockFrequencyInfo *>(&Analysis))
- this->MBFI = *MBFI;
- if (auto **Loops = llvm::any_cast<MachineLoopInfo *>(&Analysis))
- this->Loops = *Loops;
- }
- }
-
- 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;
+ return;
}
if (ModelUnderTraining.empty())
Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
@@ -553,15 +515,15 @@ class DevelopmentModeEvictionAdvisorProvider final
Ctx, ModelUnderTraining, DecisionName, TrainingInputFeatures);
if (!Runner) {
Ctx.emitError("Regalloc: could not set up the model runner");
- return false;
+ return;
}
if (TrainingLog.empty())
- return false;
+ return;
std::error_code EC;
auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC);
if (EC) {
- M.getContext().emitError(EC.message() + ":" + TrainingLog);
- return false;
+ Ctx.emitError(EC.message() + ":" + TrainingLog);
+ return;
}
std::vector<TensorSpec> LFS = InputFeatures;
if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(Runner.get()))
@@ -573,7 +535,28 @@ class DevelopmentModeEvictionAdvisorProvider final
Log = std::make_unique<Logger>(std::move(OS), LFS, Reward,
/*IncludeReward*/ true);
- return false;
+ return;
+ }
+
+ // support for isa<> and dyn_cast.
+ static bool classof(const RegAllocEvictionAdvisorProvider *R) {
+ return R->getAdvisorMode() == AdvisorMode::Development;
+ }
+
+ void logRewardIfNeeded(const MachineFunction &MF,
+ llvm::function_ref<float()> GetReward) override {
+ if (!Log || !Log->hasAnyObservationForContext(MF.getName()))
+ 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());
}
std::unique_ptr<RegAllocEvictionAdvisor>
@@ -594,8 +577,6 @@ class DevelopmentModeEvictionAdvisorProvider final
std::unique_ptr<MLModelRunner> Runner;
std::unique_ptr<Logger> Log;
- const MachineBlockFrequencyInfo *MBFI;
- const MachineLoopInfo *Loops;
};
class DevelopmentModeEvictionAdvisorAnalysisLegacy final
@@ -605,15 +586,22 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
: RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Development) {}
bool doInitialization(Module &M) override {
- auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
- auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
- Provider.setAnalyses({MBFI, Loops});
- return Provider.doInitialization(M);
+ Provider = std::make_unique<DevelopmentModeEvictionAdvisorProvider>(
+ M.getContext());
+ return false;
+ }
+
+ void logRewardIfNeeded(const MachineFunction &MF,
+ llvm::function_ref<float()> GetReward) override {
+ Provider->logRewardIfNeeded(MF, GetReward);
}
std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
- return Provider.getAdvisor(MF, RA);
+ auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+ auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI();
+ Provider->setAnalyses(MBFI, Loops);
+ return Provider->getAdvisor(MF, RA);
}
// support for isa<> and dyn_cast.
@@ -628,7 +616,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final
}
private:
- DevelopmentModeEvictionAdvisorProvider Provider;
+ std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider;
};
#endif //#ifdef LLVM_HAVE_TFLITE
@@ -1156,8 +1144,9 @@ void llvm::extractMBBFrequency(
// Development mode-specific implementations
#ifdef LLVM_HAVE_TFLITE
-RegAllocEvictionAdvisorProvider *llvm::createDevelopmentModeAdvisorProvider() {
- return new DevelopmentModeEvictionAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) {
+ return new DevelopmentModeEvictionAdvisorProvider(Ctx);
}
RegAllocEvictionAdvisorAnalysisLegacy *
@@ -1236,8 +1225,9 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
}
#endif // #ifdef LLVM_HAVE_TFLITE
-RegAllocEvictionAdvisorProvider *llvm::createReleaseModeAdvisorProvider() {
- return new ReleaseModeEvictionAdvisorProvider();
+RegAllocEvictionAdvisorProvider *
+llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) {
+ return new ReleaseModeEvictionAdvisorProvider(Ctx);
}
RegAllocEvictionAdvisorAnalysisLegacy *
diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
index 09fb5df259a97e9..4c7528ebffbfd55 100644
--- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
+++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp
@@ -9,13 +9,14 @@
// Implementation of the default eviction advisor and of the Analysis pass.
//
//===----------------------------------------------------------------------===//
-
#include "llvm/CodeGen/RegAllocEvictionAdvisor.h"
#include "AllocationOrder.h"
#include "RegAllocGreedy.h"
#include "RegAllocPriorityAdvisor.h"
#include "llvm/CodeGen/LiveRegMatrix.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
#include "llvm/CodeGen/VirtRegMap.h"
#include "llvm/IR/Module.h"
@@ -69,9 +70,12 @@ namespace {
class DefaultEvictionAdvisorProvider final
: public RegAllocEvictionAdvisorProvider {
public:
- DefaultEvictionAdvisorProvider(bool NotAsRequested)
- : RegAllocEvictionAdvisorProvider(AdvisorMode::Default),
- NotAsRequested(NotAsRequested) {}
+ DefaultEvictionAdvisorProvider(bool NotAsRequested, LLVMContext &Ctx)
+ : RegAllocEvictionAdvisorProvider(AdvisorMode::Default, Ctx) {
+ if (NotAsRequested)
+ Ctx.emitError("Requested regalloc eviction advisor analysis "
+ "could not be created. Using default");
+ }
// support for isa<> and dyn_cast.
static bool classof(const RegAllocEvictionAdvisorProvider *R) {
@@ -82,16 +86,6 @@ class DefaultEvictionAdvisorProvider final
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
return std::make_unique<DefaultEvictionAdvisor>(MF, RA);
}
-
- bool doInitialization(Module &M) override {
- if (NotAsRequested)
- M.getContext().emitError("Requested regalloc eviction advisor analysis "
- "could not be created. Using default");
- return RegAllocEvictionAdvisorProvider::doInitialization(M);
- }
-
-private:
- const bool NotAsRequested;
};
class DefaultEvictionAdvisorAnalysisLegacy final
@@ -103,12 +97,14 @@ class DefaultEvictionAdvisorAnalysisLegacy final
std::unique_ptr<RegAllocEvictionAdvisor>
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
+ // MBFI and Loops not required here.
return Provider->getAdvisor(MF, RA);
}
bool doInitialization(Module &M) override {
- Provider.reset(new DefaultEvictionAdvisorProvider(NotAsRequested));
- return Provider->doInitialization(M);
+ Provider.reset(
+ new DefaultEvictionAdvisorProvider(NotAsRequested, M.getContext()));
+ return false;
}
// support for isa<> and dyn_cast.
@@ -124,27 +120,37 @@ class DefaultEvictionAdvisorAnalysisLegacy final
AnalysisKey RegAllocEvictionAdvisorAnalysis::Key;
-RegAllocEvictionAdvisorAnalysis::Result
-RegAllocEvictionAdvisorAnalysis::run(Module &M, ModuleAnalysisManager &MAM) {
- std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider;
+void RegAllocEvictionAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) {
+ if (Provider)
+ return;
switch (Mode) {
case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default:
Provider.reset(
- new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false));
+ new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx));
break;
case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development:
#if defined(LLVM_HAVE_TFLITE)
- Provider.reset(createDevelopmentModeAdvisorProvider());
+ Provider.reset(createDevelopmentModeAdvisorProvider(Ctx));
#endif
break;
case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release:
- Provider.reset(createReleaseModeAdvisorProvider());
+ Provider.reset(createReleaseModeAdvisorProvider(Ctx));
break;
}
if (!Provider)
- Provider.reset(new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true));
- Provider->doInitialization(M);
- return Provider;
+ Provider.reset(
+ new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx));
+}
+
+RegAllocEvictionAdvisorAnalysis::Result
+RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF,
+ MachineFunctionAnalysisManager &MFAM) {
+ // Lazy initialization of the provider.
+ initializeProvider(MF.getFunction().getContext());
+ auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF);
+ auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF);
+ Provider->setAnalyses(MBFI, Loops);
+ return Result{Provider.get()};
}
template <>
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 2a2d2ad674680b7..7c3798f6462a46d 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -32,7 +32,6 @@ MODULE_ANALYSIS("no-op-module", NoOpModuleAnalysis())
MODULE_ANALYSIS("pass-instrumentation", PassInstrumentationAnalysis(PIC))
MODULE_ANALYSIS("profile-summary", ProfileSummaryAnalysis())
MODULE_ANALYSIS("reg-usage", PhysicalRegisterUsageAnalysis())
-MODULE_ANALYSIS("regalloc-evict", RegAllocEvictionAdvisorAnalysis())
MODULE_ANALYSIS("stack-safety", StackSafetyGlobalAnalysis())
MODULE_ANALYSIS("verify", VerifierAnalysis())
More information about the llvm-commits
mailing list