Skip to content

Commit

Permalink
[CodeGen][NewPM] Port RegAllocEvictionAdvisor analysis to NPM
Browse files Browse the repository at this point in the history
  • Loading branch information
@optimisan
optimisan committed Nov 22, 2024
1 parent ef20644 commit 50b7ba3
Show file tree
Hide file tree
Showing 9 changed files with 279 additions and 78 deletions.
67 changes: 60 additions & 7 deletions llvm/lib/CodeGen/RegAllocEvictionAdvisor.h → ...de/llvm/CodeGen/RegAllocEvictionAdvisor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -164,20 +166,20 @@ 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)
virtual std::unique_ptr<RegAllocEvictionAdvisor>
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
Expand All @@ -191,13 +193,64 @@ 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;

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
Expand Down
2 changes: 1 addition & 1 deletion llvm/include/llvm/InitializePasses.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 &);
Expand Down
167 changes: 128 additions & 39 deletions llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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)
Expand All @@ -380,17 +381,19 @@ 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 setAnalyses(std::initializer_list<llvm::Any> AnalysisList) override {
// llvm::errs() << "here";
for (auto Analysis : AnalysisList) {

void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addRequired<MachineBlockFrequencyInfoWrapperPass>();
AU.addRequired<MachineLoopInfoWrapperPass>();
RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
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>
Expand All @@ -405,12 +408,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;
};

// ===================================
Expand Down Expand Up @@ -445,11 +483,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)
Expand All @@ -471,7 +512,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;
}

Expand All @@ -491,14 +532,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 {
Expand Down Expand Up @@ -544,14 +584,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
Expand Down Expand Up @@ -1079,8 +1158,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(
Expand Down Expand Up @@ -1146,18 +1230,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;
}

Expand Down
Loading

0 comments on commit 50b7ba3

Please sign in to comment.