https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/117309
>From 04586196b610c1348bfe5fbdd67b9faa43431b43 Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Fri, 22 Nov 2024 09:31:50 +0000 Subject: [PATCH 1/8] [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 52dd946a685400..847bf032235c1d 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 1cb9013bc48cc5..8915bf258005a7 100644 --- a/llvm/include/llvm/InitializePasses.h +++ b/llvm/include/llvm/InitializePasses.h @@ -253,7 +253,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 a84164bed46cec..175b5aba215f26 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" @@ -1058,9 +1059,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 9c6487b40d6061..1eca5c92e2fac4 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" @@ -114,7 +114,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); @@ -372,11 +372,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) @@ -386,17 +387,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> @@ -411,12 +412,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; }; // =================================== @@ -451,11 +487,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) @@ -477,7 +516,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; } @@ -497,14 +536,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 { @@ -550,14 +588,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 @@ -1094,8 +1171,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( @@ -1161,18 +1243,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 a1f441ebd0d5e4..09fb5df259a97e 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 e61dad5cf64d08..ffe93d4a3fb76e 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 "SplitKit.h" #include "llvm/ADT/ArrayRef.h" @@ -46,6 +45,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(SpillPlacementWrapperLegacy) 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<SpillPlacementWrapperLegacy>(); AU.addRequired<MachineOptimizationRemarkEmitterPass>(); - AU.addRequired<RegAllocEvictionAdvisorAnalysis>(); + AU.addRequired<RegAllocEvictionAdvisorAnalysisLegacy>(); AU.addRequired<RegAllocPriorityAdvisorAnalysis>(); MachineFunctionPass::getAnalysisUsage(AU); } @@ -2765,7 +2765,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 594c481826cf09..001a1e86a96914 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 "SplitKit.h" #include "llvm/ADT/ArrayRef.h" diff --git a/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h b/llvm/lib/CodeGen/RegAllocPriorityAdvisor.h index 32e4598b715392..0758743c2b1403 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 30b8d7c9499488..a34fd6b8bfbf69 100644 --- a/llvm/lib/Passes/PassBuilder.cpp +++ b/llvm/lib/Passes/PassBuilder.cpp @@ -123,6 +123,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 9f0b09278edcca..ca00ac55a9687e 100644 --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def @@ -33,6 +33,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 513a767be21fb4fe75ea9dc7f6ce9d6b56bc5cde Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Fri, 29 Nov 2024 09:18:59 +0000 Subject: [PATCH 2/8] 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 847bf032235c1d..8261f6d1cf0bbd 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 175b5aba215f26..50f548b3a92e04 100644 --- a/llvm/include/llvm/Passes/CodeGenPassBuilder.h +++ b/llvm/include/llvm/Passes/CodeGenPassBuilder.h @@ -1059,10 +1059,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 29763995e8b516..827875551e7c04 100644 --- a/llvm/include/llvm/Passes/MachinePassRegistry.def +++ b/llvm/include/llvm/Passes/MachinePassRegistry.def @@ -114,6 +114,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("spill-code-placement", SpillPlacementAnalysis()) MACHINE_FUNCTION_ANALYSIS("virtregmap", VirtRegMapAnalysis()) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 1eca5c92e2fac4..9523183745d119 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -376,8 +376,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) @@ -391,15 +391,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) { @@ -412,6 +403,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); } @@ -419,8 +412,6 @@ class ReleaseModeEvictionAdvisorProvider final private: std::vector<TensorSpec> InputFeatures; std::unique_ptr<MLModelRunner> Runner; - MachineBlockFrequencyInfo *MBFI; - MachineLoopInfo *Loops; }; class ReleaseModeEvictionAdvisorAnalysisLegacy final @@ -433,12 +424,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) { @@ -452,7 +450,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final } private: - ReleaseModeEvictionAdvisorProvider Provider; + std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider; }; // =================================== @@ -490,11 +488,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) @@ -514,43 +509,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); @@ -559,15 +521,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())) @@ -579,7 +541,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> @@ -600,8 +583,6 @@ class DevelopmentModeEvictionAdvisorProvider final std::unique_ptr<MLModelRunner> Runner; std::unique_ptr<Logger> Log; - const MachineBlockFrequencyInfo *MBFI; - const MachineLoopInfo *Loops; }; class DevelopmentModeEvictionAdvisorAnalysisLegacy final @@ -611,15 +592,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. @@ -634,7 +622,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final } private: - DevelopmentModeEvictionAdvisorProvider Provider; + std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider; }; #endif // #ifdef LLVM_HAVE_TFLITE @@ -1171,8 +1159,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 * @@ -1251,8 +1240,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 09fb5df259a97e..4c7528ebffbfd5 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 ca00ac55a9687e..9f0b09278edcca 100644 --- a/llvm/lib/Passes/PassRegistry.def +++ b/llvm/lib/Passes/PassRegistry.def @@ -33,7 +33,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()) >From 2db92cf37002ca6e466e555a599dfb2800b8013b Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Tue, 3 Dec 2024 06:17:46 +0000 Subject: [PATCH 3/8] Remove parenthesis --- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 9523183745d119..2eab8fccceced9 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -571,7 +571,7 @@ class DevelopmentModeEvictionAdvisorProvider final return nullptr; if (Log) Log->switchContext(MF.getName()); - assert((MBFI && Loops) && + assert(MBFI && Loops && "Invalid provider state: must have analysis available"); return std::make_unique<DevelopmentModeEvictAdvisor>( MF, RA, Runner.get(), *MBFI, *Loops, Log.get()); >From dc6a359483e6045003ef902799654c50d433790b Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Wed, 4 Dec 2024 06:44:24 +0000 Subject: [PATCH 4/8] Add TODO to remove provider, remove paren --- llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h | 1 + llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h index 8261f6d1cf0bbd..9ebe5c22f47399 100644 --- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h @@ -199,6 +199,7 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass { /// 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. +/// TODO: Coalesce this with the NPM analysis when legacy PM is removed. class RegAllocEvictionAdvisorProvider { public: enum class AdvisorMode : int { Default, Release, Development }; diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 2eab8fccceced9..534237713f0cb5 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -403,7 +403,7 @@ class ReleaseModeEvictionAdvisorProvider final InteractiveChannelBaseName + ".out", InteractiveChannelBaseName + ".in"); } - assert((MBFI && Loops) && + assert(MBFI && Loops && "Invalid provider state: must have analysis available"); return std::make_unique<MLEvictAdvisor>(MF, RA, Runner.get(), *MBFI, *Loops); >From 98bdd44adac21f430bc8c376e9aceb249067f050 Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Mon, 9 Dec 2024 08:28:04 +0000 Subject: [PATCH 5/8] Remove createProvider*(), add new method in ML.cpp --- .../llvm/CodeGen/RegAllocEvictionAdvisor.h | 13 +++++----- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 24 ++++++++++++------- llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp | 20 +++++----------- 3 files changed, 29 insertions(+), 28 deletions(-) diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h index 9ebe5c22f47399..3e6127351911bd 100644 --- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h @@ -233,11 +233,6 @@ class RegAllocEvictionAdvisorProvider { const AdvisorMode Mode; }; -RegAllocEvictionAdvisorProvider * -createReleaseModeAdvisorProvider(LLVMContext &Ctx); -RegAllocEvictionAdvisorProvider * -createDevelopmentModeAdvisorProvider(LLVMContext &Ctx); - /// 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 @@ -276,7 +271,13 @@ class RegAllocEvictionAdvisorAnalysis Result run(MachineFunction &MF, MachineFunctionAnalysisManager &MAM); private: - void initializeProvider(LLVMContext &Ctx); + void + initializeProvider(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, + LLVMContext &Ctx); + void + initializeMLProvider(RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, + LLVMContext &Ctx); + std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider; }; diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 534237713f0cb5..26bd241fc6ae37 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -1159,11 +1159,6 @@ void llvm::extractMBBFrequency( // Development mode-specific implementations #ifdef LLVM_HAVE_TFLITE -RegAllocEvictionAdvisorProvider * -llvm::createDevelopmentModeAdvisorProvider(LLVMContext &Ctx) { - return new DevelopmentModeEvictionAdvisorProvider(Ctx); -} - RegAllocEvictionAdvisorAnalysisLegacy * llvm::createDevelopmentModeAdvisorAnalysisLegacy() { return new DevelopmentModeEvictionAdvisorAnalysisLegacy(); @@ -1240,9 +1235,22 @@ bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) { } #endif // #ifdef LLVM_HAVE_TFLITE -RegAllocEvictionAdvisorProvider * -llvm::createReleaseModeAdvisorProvider(LLVMContext &Ctx) { - return new ReleaseModeEvictionAdvisorProvider(Ctx); +void RegAllocEvictionAdvisorAnalysis::initializeMLProvider( + RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, LLVMContext &Ctx) { + if (Provider) + return; + switch (Mode) { + case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development: +#if defined(LLVM_HAVE_TFLITE) + Provider.reset(new DevelopmentModeEvictionAdvisorProvider(Ctx)); +#endif + break; + case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release: + Provider.reset(new ReleaseModeEvictionAdvisorProvider(Ctx)); + break; + default: + break; + } } RegAllocEvictionAdvisorAnalysisLegacy * diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp index 4c7528ebffbfd5..1307ed2651a91a 100644 --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp @@ -120,23 +120,15 @@ class DefaultEvictionAdvisorAnalysisLegacy final AnalysisKey RegAllocEvictionAdvisorAnalysis::Key; -void RegAllocEvictionAdvisorAnalysis::initializeProvider(LLVMContext &Ctx) { +void RegAllocEvictionAdvisorAnalysis::initializeProvider( + RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode Mode, LLVMContext &Ctx) { if (Provider) return; - switch (Mode) { - case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default: + if (Mode == RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default) Provider.reset( new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx)); - break; - case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Development: -#if defined(LLVM_HAVE_TFLITE) - Provider.reset(createDevelopmentModeAdvisorProvider(Ctx)); -#endif - break; - case RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Release: - Provider.reset(createReleaseModeAdvisorProvider(Ctx)); - break; - } + else + initializeMLProvider(Mode, Ctx); if (!Provider) Provider.reset( new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx)); @@ -146,7 +138,7 @@ RegAllocEvictionAdvisorAnalysis::Result RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) { // Lazy initialization of the provider. - initializeProvider(MF.getFunction().getContext()); + initializeProvider(::Mode, MF.getFunction().getContext()); auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF); auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF); Provider->setAnalyses(MBFI, Loops); >From d4878a52e482312a95f81bad09de60fd952858df Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Wed, 11 Dec 2024 09:08:06 +0000 Subject: [PATCH 6/8] Use the provider as analysis result --- .../llvm/CodeGen/RegAllocEvictionAdvisor.h | 81 ++++++++++--------- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 16 ++-- llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp | 7 +- llvm/lib/CodeGen/RegAllocGreedy.cpp | 3 +- 4 files changed, 53 insertions(+), 54 deletions(-) diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h index 3e6127351911bd..25df101ee512e9 100644 --- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h @@ -153,6 +153,44 @@ class RegAllocEvictionAdvisor { const bool EnableLocalReassign; }; +/// 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. +/// TODO: Coalesce this with the NPM analysis when legacy PM is removed. +class RegAllocEvictionAdvisorProvider { +public: + enum class AdvisorMode : int { Default, Release, Development }; + RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx) + : Ctx(Ctx), Mode(Mode) {} + + virtual ~RegAllocEvictionAdvisorProvider() = default; + + virtual void logRewardIfNeeded(const MachineFunction &MF, + llvm::function_ref<float()> GetReward) {} + + virtual std::unique_ptr<RegAllocEvictionAdvisor> + getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; + + /// 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; +}; + /// ImmutableAnalysis abstraction for fetching the Eviction Advisor. We model it /// as an analysis to decouple the user from the implementation insofar as /// dependencies on other analyses goes. The motivation for it being an @@ -177,8 +215,8 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass { static char ID; /// Get an advisor for the given context (i.e. machine function, etc) - virtual std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; + virtual std::unique_ptr<RegAllocEvictionAdvisorProvider>& + getProvider() = 0; AdvisorMode getAdvisorMode() const { return Mode; } virtual void logRewardIfNeeded(const MachineFunction &MF, llvm::function_ref<float()> GetReward) {}; @@ -189,50 +227,13 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass { void getAnalysisUsage(AnalysisUsage &AU) const override { AU.setPreservesAll(); } + std::unique_ptr<RegAllocEvictionAdvisorProvider> Provider; private: StringRef getPassName() const override; 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. -/// TODO: Coalesce this with the NPM analysis when legacy PM is removed. -class RegAllocEvictionAdvisorProvider { -public: - enum class AdvisorMode : int { Default, Release, Development }; - RegAllocEvictionAdvisorProvider(AdvisorMode Mode, LLVMContext &Ctx) - : Ctx(Ctx), Mode(Mode) {} - - virtual ~RegAllocEvictionAdvisorProvider() = default; - - virtual void logRewardIfNeeded(const MachineFunction &MF, - llvm::function_ref<float()> GetReward) {} - - virtual std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; - - /// 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; -}; - /// 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 diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 26bd241fc6ae37..368bf592ebac90 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -420,12 +420,12 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final ReleaseModeEvictionAdvisorAnalysisLegacy() : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {} - std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + std::unique_ptr<RegAllocEvictionAdvisorProvider>& + getProvider() override { auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(); auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI(); Provider->setAnalyses(MBFI, Loops); - return Provider->getAdvisor(MF, RA); + return Provider; } void logRewardIfNeeded(const MachineFunction &MF, @@ -450,7 +450,7 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final } private: - std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider; + // std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider; }; // =================================== @@ -602,12 +602,12 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final Provider->logRewardIfNeeded(MF, GetReward); } - std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + std::unique_ptr<RegAllocEvictionAdvisorProvider>& + getProvider() override { auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(); auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI(); Provider->setAnalyses(MBFI, Loops); - return Provider->getAdvisor(MF, RA); + return Provider; } // support for isa<> and dyn_cast. @@ -622,7 +622,7 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final } private: - std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider; + // std::unique_ptr<DevelopmentModeEvictionAdvisorProvider> Provider; }; #endif // #ifdef LLVM_HAVE_TFLITE diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp index 1307ed2651a91a..e00b81d8dcf61d 100644 --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp @@ -95,10 +95,10 @@ class DefaultEvictionAdvisorAnalysisLegacy final : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default), NotAsRequested(NotAsRequested) {} - std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + std::unique_ptr<RegAllocEvictionAdvisorProvider>& + getProvider() override { // MBFI and Loops not required here. - return Provider->getAdvisor(MF, RA); + return Provider; } bool doInitialization(Module &M) override { @@ -113,7 +113,6 @@ class DefaultEvictionAdvisorAnalysisLegacy final } private: - std::unique_ptr<DefaultEvictionAdvisorProvider> Provider; const bool NotAsRequested; }; } // namespace diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp index ffe93d4a3fb76e..8d8b12278a30c3 100644 --- a/llvm/lib/CodeGen/RegAllocGreedy.cpp +++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp @@ -2765,8 +2765,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) { ExtraInfo.emplace(); EvictAdvisor = - getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getAdvisor(*MF, - *this); + getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider()->getAdvisor(*MF, *this); PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this); >From 220fd47097157e0b17c3d31a99e9bf81d6f37ddf Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Thu, 26 Dec 2024 08:25:11 +0000 Subject: [PATCH 7/8] Remove setAnalyses(); pass analysis in getAdvisor --- .../llvm/CodeGen/RegAllocEvictionAdvisor.h | 19 +++++--------- llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp | 25 +++---------------- llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp | 16 +++--------- llvm/lib/CodeGen/RegAllocGreedy.cpp | 7 ++++-- 4 files changed, 19 insertions(+), 48 deletions(-) diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h index 25df101ee512e9..781afbc571abbf 100644 --- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h @@ -170,22 +170,13 @@ class RegAllocEvictionAdvisorProvider { llvm::function_ref<float()> GetReward) {} virtual std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) = 0; - - /// 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; - } + getAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) = 0; AdvisorMode getAdvisorMode() const { return Mode; } protected: LLVMContext &Ctx; - MachineBlockFrequencyInfo *MBFI; - MachineLoopInfo *Loops; private: const AdvisorMode Mode; @@ -215,8 +206,10 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass { static char ID; /// Get an advisor for the given context (i.e. machine function, etc) - virtual std::unique_ptr<RegAllocEvictionAdvisorProvider>& - getProvider() = 0; + std::unique_ptr<RegAllocEvictionAdvisorProvider> &getProvider() { + return Provider; + } + AdvisorMode getAdvisorMode() const { return Mode; } virtual void logRewardIfNeeded(const MachineFunction &MF, llvm::function_ref<float()> GetReward) {}; diff --git a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp index 368bf592ebac90..fa315146090b8a 100644 --- a/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp +++ b/llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp @@ -392,7 +392,8 @@ class ReleaseModeEvictionAdvisorProvider final } std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + getAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override { if (!Runner) { if (InteractiveChannelBaseName.empty()) Runner = std::make_unique<ReleaseModeModelRunner<CompiledModelType>>( @@ -420,14 +421,6 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final ReleaseModeEvictionAdvisorAnalysisLegacy() : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Release) {} - std::unique_ptr<RegAllocEvictionAdvisorProvider>& - getProvider() override { - auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(); - auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI(); - Provider->setAnalyses(MBFI, Loops); - return Provider; - } - void logRewardIfNeeded(const MachineFunction &MF, llvm::function_ref<float()> GetReward) override { // No-op in release mode @@ -448,9 +441,6 @@ class ReleaseModeEvictionAdvisorAnalysisLegacy final AU.addRequired<MachineLoopInfoWrapperPass>(); RegAllocEvictionAdvisorAnalysisLegacy::getAnalysisUsage(AU); } - -private: - // std::unique_ptr<ReleaseModeEvictionAdvisorProvider> Provider; }; // =================================== @@ -566,7 +556,8 @@ class DevelopmentModeEvictionAdvisorProvider final } std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + getAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MachineBlockFrequencyInfo *MBFI, MachineLoopInfo *Loops) override { if (!Runner) return nullptr; if (Log) @@ -602,14 +593,6 @@ class DevelopmentModeEvictionAdvisorAnalysisLegacy final Provider->logRewardIfNeeded(MF, GetReward); } - std::unique_ptr<RegAllocEvictionAdvisorProvider>& - getProvider() override { - auto *MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI(); - auto *Loops = &getAnalysis<MachineLoopInfoWrapperPass>().getLI(); - Provider->setAnalyses(MBFI, Loops); - return Provider; - } - // support for isa<> and dyn_cast. static bool classof(const RegAllocEvictionAdvisorAnalysisLegacy *R) { return R->getAdvisorMode() == AdvisorMode::Development; diff --git a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp index e00b81d8dcf61d..73de74c5aad925 100644 --- a/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp +++ b/llvm/lib/CodeGen/RegAllocEvictionAdvisor.cpp @@ -83,7 +83,8 @@ class DefaultEvictionAdvisorProvider final } std::unique_ptr<RegAllocEvictionAdvisor> - getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override { + getAdvisor(const MachineFunction &MF, const RAGreedy &RA, + MachineBlockFrequencyInfo *, MachineLoopInfo *) override { return std::make_unique<DefaultEvictionAdvisor>(MF, RA); } }; @@ -95,12 +96,6 @@ class DefaultEvictionAdvisorAnalysisLegacy final : RegAllocEvictionAdvisorAnalysisLegacy(AdvisorMode::Default), NotAsRequested(NotAsRequested) {} - std::unique_ptr<RegAllocEvictionAdvisorProvider>& - getProvider() override { - // MBFI and Loops not required here. - return Provider; - } - bool doInitialization(Module &M) override { Provider.reset( new DefaultEvictionAdvisorProvider(NotAsRequested, M.getContext())); @@ -125,12 +120,12 @@ void RegAllocEvictionAdvisorAnalysis::initializeProvider( return; if (Mode == RegAllocEvictionAdvisorAnalysisLegacy::AdvisorMode::Default) Provider.reset( - new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ false, Ctx)); + new DefaultEvictionAdvisorProvider(/*NotAsRequested=*/false, Ctx)); else initializeMLProvider(Mode, Ctx); if (!Provider) Provider.reset( - new DefaultEvictionAdvisorProvider(/*NotAsRequested*/ true, Ctx)); + new DefaultEvictionAdvisorProvider(/*NotAsRequested=*/true, Ctx)); } RegAllocEvictionAdvisorAnalysis::Result @@ -138,9 +133,6 @@ RegAllocEvictionAdvisorAnalysis::run(MachineFunction &MF, MachineFunctionAnalysisManager &MFAM) { // Lazy initialization of the provider. initializeProvider(::Mode, MF.getFunction().getContext()); - auto *MBFI = &MFAM.getResult<MachineBlockFrequencyAnalysis>(MF); - auto *Loops = &MFAM.getResult<MachineLoopAnalysis>(MF); - Provider->setAnalyses(MBFI, Loops); return Result{Provider.get()}; } diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp index 8d8b12278a30c3..0cb5377d9dc76a 100644 --- a/llvm/lib/CodeGen/RegAllocGreedy.cpp +++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp @@ -2764,8 +2764,11 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) { : TRI->reverseLocalAssignment(); ExtraInfo.emplace(); - EvictAdvisor = - getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider()->getAdvisor(*MF, *this); + + auto &EvictAdvisorProvider = + getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider(); + EvictAdvisor = EvictAdvisorProvider->getAdvisor(*MF, *this, MBFI, Loops); + PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this); >From 3ea391fe2e94d75e30ff8e430295d13bbbd0d1a8 Mon Sep 17 00:00:00 2001 From: Akshat Oke <akshat....@amd.com> Date: Tue, 7 Jan 2025 09:38:25 +0000 Subject: [PATCH 8/8] Apply suggestions --- llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h | 6 ++---- llvm/lib/CodeGen/RegAllocGreedy.cpp | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h index 781afbc571abbf..beb99ba2bc457c 100644 --- a/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h +++ b/llvm/include/llvm/CodeGen/RegAllocEvictionAdvisor.h @@ -206,13 +206,11 @@ class RegAllocEvictionAdvisorAnalysisLegacy : public ImmutablePass { static char ID; /// Get an advisor for the given context (i.e. machine function, etc) - std::unique_ptr<RegAllocEvictionAdvisorProvider> &getProvider() { - return Provider; - } + RegAllocEvictionAdvisorProvider &getProvider() { return *Provider; } AdvisorMode getAdvisorMode() const { return Mode; } virtual void logRewardIfNeeded(const MachineFunction &MF, - llvm::function_ref<float()> GetReward) {}; + function_ref<float()> GetReward) {}; protected: // This analysis preserves everything, and subclasses may have additional diff --git a/llvm/lib/CodeGen/RegAllocGreedy.cpp b/llvm/lib/CodeGen/RegAllocGreedy.cpp index 0cb5377d9dc76a..9b8101ca62b7bb 100644 --- a/llvm/lib/CodeGen/RegAllocGreedy.cpp +++ b/llvm/lib/CodeGen/RegAllocGreedy.cpp @@ -2767,7 +2767,7 @@ bool RAGreedy::runOnMachineFunction(MachineFunction &mf) { auto &EvictAdvisorProvider = getAnalysis<RegAllocEvictionAdvisorAnalysisLegacy>().getProvider(); - EvictAdvisor = EvictAdvisorProvider->getAdvisor(*MF, *this, MBFI, Loops); + EvictAdvisor = EvictAdvisorProvider.getAdvisor(*MF, *this, MBFI, Loops); PriorityAdvisor = getAnalysis<RegAllocPriorityAdvisorAnalysis>().getAdvisor(*MF, *this); _______________________________________________ llvm-branch-commits mailing list llvm-branch-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits