Skip to content

[CodeGen] Support start/stop in CodeGenPassBuilder #70912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 18, 2024
Merged

[CodeGen] Support start/stop in CodeGenPassBuilder #70912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 113 additions & 20 deletions llvm/include/llvm/CodeGen/CodeGenPassBuilder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
#include "llvm/CodeGen/ShadowStackGCLowering.h"
#include "llvm/CodeGen/SjLjEHPrepare.h"
#include "llvm/CodeGen/StackProtector.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/CodeGen/UnreachableBlockElim.h"
#include "llvm/CodeGen/WasmEHPrepare.h"
#include "llvm/CodeGen/WinEHPrepare.h"
Expand Down Expand Up @@ -175,73 +176,80 @@ template <typename DerivedT> class CodeGenPassBuilder {
// Function object to maintain state while adding codegen IR passes.
class AddIRPass {
public:
AddIRPass(ModulePassManager &MPM) : MPM(MPM) {}
AddIRPass(ModulePassManager &MPM, const DerivedT &PB) : MPM(MPM), PB(PB) {}
~AddIRPass() {
if (!FPM.isEmpty())
MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
}

template <typename PassT> void operator()(PassT &&Pass) {
template <typename PassT>
void operator()(PassT &&Pass, StringRef Name = PassT::name()) {
static_assert((is_detected<is_function_pass_t, PassT>::value ||
is_detected<is_module_pass_t, PassT>::value) &&
"Only module pass and function pass are supported.");

if (!PB.runBeforeAdding(Name))
return;

// Add Function Pass
if constexpr (is_detected<is_function_pass_t, PassT>::value) {
FPM.addPass(std::forward<PassT>(Pass));

for (auto &C : PB.AfterCallbacks)
C(Name);
} else {
// Add Module Pass
if (!FPM.isEmpty()) {
MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
FPM = FunctionPassManager();
}

MPM.addPass(std::forward<PassT>(Pass));

for (auto &C : PB.AfterCallbacks)
C(Name);
}
}

private:
ModulePassManager &MPM;
FunctionPassManager FPM;
const DerivedT &PB;
};

// Function object to maintain state while adding codegen machine passes.
class AddMachinePass {
public:
AddMachinePass(MachineFunctionPassManager &PM) : PM(PM) {}
AddMachinePass(MachineFunctionPassManager &PM, const DerivedT &PB)
: PM(PM), PB(PB) {}

template <typename PassT> void operator()(PassT &&Pass) {
static_assert(
is_detected<has_key_t, PassT>::value,
"Machine function pass must define a static member variable `Key`.");
for (auto &C : BeforeCallbacks)
if (!C(&PassT::Key))
return;

if (!PB.runBeforeAdding(PassT::name()))
return;

PM.addPass(std::forward<PassT>(Pass));
for (auto &C : AfterCallbacks)
C(&PassT::Key);

for (auto &C : PB.AfterCallbacks)
C(PassT::name());
}

template <typename PassT> void insertPass(MachinePassKey *ID, PassT Pass) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unrelated, and I think I mentioned this before, but I'd like to see if we can come up with a nicer way of doing this, perhaps with something like PassBuilder extension points (e.g. PassBuilder::registerPeepholeEPCallback()). it seems in line with various per-target TargetMachines adding their own passes anyway. do you know if we could potentially replace insertPass() with the various TargetMachines simply adding passes in the right place, or are the passes they add too ad hoc?

Copy link
Contributor Author

@paperchalice paperchalice Jan 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently AMDGPU, Hexagon, PowerPC and RISCV uses insertPass to adjust the regalloc pipeline, we need add extension points after PHIElimination, TwoAddressInstruction, MachineScheduler, RenameIndependentSubregs, LiveVariables and DetectDeadLanes. They inject passes between these passes directly and only happens in regalloc pipeline.

AfterCallbacks.emplace_back(
PB.AfterCallbacks.emplace_back(
[this, ID, Pass = std::move(Pass)](MachinePassKey *PassID) {
if (PassID == ID)
this->PM.addPass(std::move(Pass));
});
}

void disablePass(MachinePassKey *ID) {
BeforeCallbacks.emplace_back(
[ID](MachinePassKey *PassID) { return PassID != ID; });
}

MachineFunctionPassManager releasePM() { return std::move(PM); }

private:
MachineFunctionPassManager &PM;
SmallVector<llvm::unique_function<bool(MachinePassKey *)>, 4>
BeforeCallbacks;
SmallVector<llvm::unique_function<void(MachinePassKey *)>, 4>
AfterCallbacks;
const DerivedT &PB;
};

LLVMTargetMachine &TM;
Expand Down Expand Up @@ -469,20 +477,43 @@ template <typename DerivedT> class CodeGenPassBuilder {
const DerivedT &derived() const {
return static_cast<const DerivedT &>(*this);
}

bool runBeforeAdding(StringRef Name) const {
bool ShouldAdd = true;
for (auto &C : BeforeCallbacks)
ShouldAdd &= C(Name);
return ShouldAdd;
}

void setStartStopPasses(const TargetPassConfig::StartStopInfo &Info) const;

Error verifyStartStop(const TargetPassConfig::StartStopInfo &Info) const;

mutable SmallVector<llvm::unique_function<bool(StringRef)>, 4>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it seems nicer to keep this in AddPass. is the reason that these are moved out of AddPass because it references Started/Stopped in CodeGenPassBuilder? and there's no way to easily validate them if these are somewhere else?

Copy link
Contributor Author

@paperchalice paperchalice Jan 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently IR passes and machine function passes are handled by different classes, and -start/stop options should support both IR and machine function passes, but it's OK to merge them into one class, because only buildPipeline creates instance of AddIRPass and AddMachinePass, but IMHO keeping these in CodeGenPassBuilder can let target disable passes in constructor.

BeforeCallbacks;
mutable SmallVector<llvm::unique_function<void(StringRef)>, 4> AfterCallbacks;

/// Helper variable for `-start-before/-start-after/-stop-before/-stop-after`
mutable bool Started = true;
mutable bool Stopped = true;
};

template <typename Derived>
Error CodeGenPassBuilder<Derived>::buildPipeline(
ModulePassManager &MPM, MachineFunctionPassManager &MFPM,
raw_pwrite_stream &Out, raw_pwrite_stream *DwoOut,
CodeGenFileType FileType) const {
AddIRPass addIRPass(MPM);
auto StartStopInfo = TargetPassConfig::getStartStopInfo(*PIC);
if (!StartStopInfo)
return StartStopInfo.takeError();
setStartStopPasses(*StartStopInfo);
AddIRPass addIRPass(MPM, derived());
// `ProfileSummaryInfo` is always valid.
addIRPass(RequireAnalysisPass<ProfileSummaryAnalysis, Module>());
addIRPass(RequireAnalysisPass<CollectorMetadataAnalysis, Module>());
addISelPasses(addIRPass);

AddMachinePass addPass(MFPM);
AddMachinePass addPass(MFPM, derived());
if (auto Err = addCoreISelPasses(addPass))
return std::move(Err);

Expand All @@ -495,6 +526,68 @@ Error CodeGenPassBuilder<Derived>::buildPipeline(
});

addPass(FreeMachineFunctionPass());
return verifyStartStop(*StartStopInfo);
}

template <typename Derived>
void CodeGenPassBuilder<Derived>::setStartStopPasses(
const TargetPassConfig::StartStopInfo &Info) const {
if (!Info.StartPass.empty()) {
Started = false;
BeforeCallbacks.emplace_back([this, &Info, AfterFlag = Info.StartAfter,
Count = 0u](StringRef ClassName) mutable {
if (Count == Info.StartInstanceNum) {
if (AfterFlag) {
AfterFlag = false;
Started = true;
}
return Started;
}

auto PassName = PIC->getPassNameForClassName(ClassName);
if (Info.StartPass == PassName && ++Count == Info.StartInstanceNum)
Started = !Info.StartAfter;

return Started;
});
}

if (!Info.StopPass.empty()) {
Stopped = false;
BeforeCallbacks.emplace_back([this, &Info, AfterFlag = Info.StopAfter,
Count = 0u](StringRef ClassName) mutable {
if (Count == Info.StopInstanceNum) {
if (AfterFlag) {
AfterFlag = false;
Stopped = true;
}
return !Stopped;
}

auto PassName = PIC->getPassNameForClassName(ClassName);
if (Info.StopPass == PassName && ++Count == Info.StopInstanceNum)
Stopped = !Info.StopAfter;
return !Stopped;
});
}
}

template <typename Derived>
Error CodeGenPassBuilder<Derived>::verifyStartStop(
const TargetPassConfig::StartStopInfo &Info) const {
if (Started && Stopped)
return Error::success();

if (!Started)
return make_error<StringError>(
"Can't find start pass \"" +
PIC->getPassNameForClassName(Info.StartPass) + "\".",
std::make_error_code(std::errc::invalid_argument));
if (!Stopped)
return make_error<StringError>(
"Can't find stop pass \"" +
PIC->getPassNameForClassName(Info.StopPass) + "\".",
std::make_error_code(std::errc::invalid_argument));
return Error::success();
}

Expand Down
15 changes: 15 additions & 0 deletions llvm/include/llvm/CodeGen/TargetPassConfig.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

#include "llvm/Pass.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/Support/Error.h"
#include <cassert>
#include <string>

Expand Down Expand Up @@ -176,6 +177,20 @@ class TargetPassConfig : public ImmutablePass {
static std::string
getLimitedCodeGenPipelineReason(const char *Separator = "/");

struct StartStopInfo {
bool StartAfter;
bool StopAfter;
unsigned StartInstanceNum;
unsigned StopInstanceNum;
StringRef StartPass;
StringRef StopPass;
};

/// Returns pass name in `-stop-before` or `-stop-after`
/// NOTE: New pass manager migration only
static Expected<StartStopInfo>
getStartStopInfo(PassInstrumentationCallbacks &PIC);

void setDisableVerify(bool Disable) { setOpt(DisableVerify, Disable); }

bool getEnableTailMerge() const { return EnableTailMerge; }
Expand Down
34 changes: 34 additions & 0 deletions llvm/lib/CodeGen/TargetPassConfig.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -609,6 +609,40 @@ void llvm::registerCodeGenCallback(PassInstrumentationCallbacks &PIC,
registerPartialPipelineCallback(PIC, LLVMTM);
}

Expected<TargetPassConfig::StartStopInfo>
TargetPassConfig::getStartStopInfo(PassInstrumentationCallbacks &PIC) {
auto [StartBefore, StartBeforeInstanceNum] =
getPassNameAndInstanceNum(StartBeforeOpt);
auto [StartAfter, StartAfterInstanceNum] =
getPassNameAndInstanceNum(StartAfterOpt);
auto [StopBefore, StopBeforeInstanceNum] =
getPassNameAndInstanceNum(StopBeforeOpt);
auto [StopAfter, StopAfterInstanceNum] =
getPassNameAndInstanceNum(StopAfterOpt);

if (!StartBefore.empty() && !StartAfter.empty())
return make_error<StringError>(
Twine(StartBeforeOptName) + " and " + StartAfterOptName + " specified!",
std::make_error_code(std::errc::invalid_argument));
if (!StopBefore.empty() && !StopAfter.empty())
return make_error<StringError>(
Twine(StopBeforeOptName) + " and " + StopAfterOptName + " specified!",
std::make_error_code(std::errc::invalid_argument));

StartStopInfo Result;
Result.StartPass = StartBefore.empty() ? StartAfter : StartBefore;
Result.StopPass = StopBefore.empty() ? StopAfter : StopBefore;
Result.StartInstanceNum =
StartBefore.empty() ? StartAfterInstanceNum : StartBeforeInstanceNum;
Result.StopInstanceNum =
StopBefore.empty() ? StopAfterInstanceNum : StopBeforeInstanceNum;
Result.StartAfter = !StartAfter.empty();
Result.StopAfter = !StopAfter.empty();
Result.StartInstanceNum += Result.StartInstanceNum == 0;
Result.StopInstanceNum += Result.StopInstanceNum == 0;
return Result;
}

// Out of line constructor provides default values for pass options and
// registers all common codegen passes.
TargetPassConfig::TargetPassConfig(LLVMTargetMachine &TM, PassManagerBase &pm)
Expand Down
4 changes: 3 additions & 1 deletion llvm/lib/Passes/PassBuilder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@
#include "llvm/CodeGen/ShadowStackGCLowering.h"
#include "llvm/CodeGen/SjLjEHPrepare.h"
#include "llvm/CodeGen/StackProtector.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/CodeGen/TypePromotion.h"
#include "llvm/CodeGen/WasmEHPrepare.h"
#include "llvm/CodeGen/WinEHPrepare.h"
Expand Down Expand Up @@ -315,7 +316,8 @@ namespace {
/// We currently only use this for --print-before/after.
bool shouldPopulateClassToPassNames() {
return PrintPipelinePasses || !printBeforePasses().empty() ||
!printAfterPasses().empty() || !isFilterPassesEmpty();
!printAfterPasses().empty() || !isFilterPassesEmpty() ||
TargetPassConfig::hasLimitedCodeGenPipeline();
}

// A pass for testing -print-on-crash.
Expand Down
41 changes: 41 additions & 0 deletions llvm/unittests/CodeGen/CodeGenPassBuilderTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,4 +138,45 @@ TEST_F(CodeGenPassBuilderTest, basic) {
EXPECT_EQ(MIRPipeline, ExpectedMIRPipeline);
}

// TODO: Move this to lit test when llc support new pm.
TEST_F(CodeGenPassBuilderTest, start_stop) {
static const char *argv[] = {
"test",
"-start-after=no-op-module",
"-stop-before=no-op-function,2",
};
int argc = std::size(argv);
cl::ParseCommandLineOptions(argc, argv);

LoopAnalysisManager LAM;
FunctionAnalysisManager FAM;
CGSCCAnalysisManager CGAM;
ModuleAnalysisManager MAM;

PassInstrumentationCallbacks PIC;
DummyCodeGenPassBuilder CGPB(*TM, getCGPassBuilderOption(), &PIC);
PipelineTuningOptions PTO;
PassBuilder PB(TM.get(), PTO, std::nullopt, &PIC);

PB.registerModuleAnalyses(MAM);
PB.registerCGSCCAnalyses(CGAM);
PB.registerFunctionAnalyses(FAM);
PB.registerLoopAnalyses(LAM);
PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);

ModulePassManager MPM;
MachineFunctionPassManager MFPM;

Error Err =
CGPB.buildPipeline(MPM, MFPM, outs(), nullptr, CodeGenFileType::Null);
EXPECT_FALSE(Err);
std::string IRPipeline;
raw_string_ostream IROS(IRPipeline);
MPM.printPipeline(IROS, [&PIC](StringRef Name) {
auto PassName = PIC.getPassNameForClassName(Name);
return PassName.empty() ? Name : PassName;
});
EXPECT_EQ(IRPipeline, "function(no-op-function)");
}

} // namespace