[llvm] [AMDGPU] Graph-based Module Splitting Rewrite (PR #104763)
Matt Arsenault via llvm-commits
llvm-commits at lists.llvm.org
Mon Aug 19 07:05:14 PDT 2024
================
@@ -651,138 +1317,176 @@ static void splitAMDGPUModule(
if (!NoExternalizeGlobals) {
for (auto &GV : M.globals()) {
if (GV.hasLocalLinkage())
- SML << "[externalize] GV " << GV.getName() << '\n';
+ LLVM_DEBUG(dbgs() << "[externalize] GV " << GV.getName() << '\n');
externalize(GV);
}
}
// Start by calculating the cost of every function in the module, as well as
// the module's overall cost.
- DenseMap<const Function *, CostType> FnCosts;
- const CostType ModuleCost = calculateFunctionCosts(SML, GetTTI, M, FnCosts);
-
- // First, gather ever kernel into the worklist.
- SmallVector<FunctionWithDependencies> WorkList;
- for (auto &Fn : M) {
- if (isEntryPoint(&Fn) && !Fn.isDeclaration())
- WorkList.emplace_back(SML, CG, FnCosts, &Fn);
+ FunctionsCostMap FnCosts;
+ const CostType ModuleCost = calculateFunctionCosts(GetTTI, M, FnCosts);
+
+ // Build the SplitGraph, which represents the module's functions and models
+ // their dependencies accurately.
+ SplitGraph SG(M, FnCosts, ModuleCost);
+ SG.buildGraph(CG);
+
+ if (SG.empty()) {
+ LLVM_DEBUG(
+ dbgs()
+ << "[!] no nodes in graph, input is empty - no splitting possible\n");
+ ModuleCallback(cloneAll(M));
+ return;
}
- // Then, find missing functions that need to be considered as additional
- // roots. These can't be called in theory, but in practice we still have to
- // handle them to avoid linker errors.
- {
- DenseSet<const Function *> SeenFunctions;
- for (const auto &FWD : WorkList) {
- SeenFunctions.insert(FWD.Fn);
- SeenFunctions.insert(FWD.Dependencies.begin(), FWD.Dependencies.end());
+ LLVM_DEBUG({
+ dbgs() << "[graph] nodes:\n";
+ for (const SplitGraph::Node *N : SG.nodes()) {
+ dbgs() << " - [" << N->getID() << "]: " << N->getName() << " "
+ << (N->isGraphEntryPoint() ? "(entry)" : "") << "\n";
}
+ });
- for (auto &Fn : M) {
- // If this function is not part of any kernel's dependencies and isn't
- // directly called, consider it as a root.
- if (!Fn.isDeclaration() && !isEntryPoint(&Fn) &&
- !SeenFunctions.count(&Fn) && !hasDirectCaller(Fn)) {
- WorkList.emplace_back(SML, CG, FnCosts, &Fn);
- }
- }
- }
+ writeDOTGraph(SG);
- // Sort the worklist so the most expensive roots are seen first.
- sort(WorkList, [&](auto &A, auto &B) {
- // Sort by total cost, and if the total cost is identical, sort
- // alphabetically.
- if (A.TotalCost == B.TotalCost)
- return A.Fn->getName() < B.Fn->getName();
- return A.TotalCost > B.TotalCost;
- });
+ LLVM_DEBUG(dbgs() << "[search] testing splitting strategies\n");
- if (SML) {
- SML << "Worklist\n";
- for (const auto &FWD : WorkList) {
- SML << "[root] " << getName(*FWD.Fn) << " (totalCost:" << FWD.TotalCost
- << " indirect:" << FWD.HasIndirectCall
- << " hasNonDuplicatableDep:" << FWD.HasNonDuplicatableDependecy
- << ")\n";
- // Sort function names before printing to ensure determinism.
- SmallVector<std::string> SortedDepNames;
- SortedDepNames.reserve(FWD.Dependencies.size());
- for (const auto *Dep : FWD.Dependencies)
- SortedDepNames.push_back(getName(*Dep));
- sort(SortedDepNames);
-
- for (const auto &Name : SortedDepNames)
- SML << " [dependency] " << Name << '\n';
- }
+ std::optional<SplitProposal> Proposal;
+ const auto EvaluateProposal = [&](SplitProposal SP) {
+ if (!Proposal)
+ Proposal = std::move(SP);
+ else
+ evaluateProposal(*Proposal, std::move(SP));
+ };
+
+ // TODO: It would be very easy to create new strategies by just adding a base
+ // class to RecursiveSearchSplitting and abstracting it away.
+ RecursiveSearchSplitting(SG, NumParts, EvaluateProposal).run();
+ LLVM_DEBUG(if (Proposal) dbgs() << "[search done] selected proposal: "
+ << Proposal->getName() << "\n";);
+
+ if (!Proposal) {
+ LLVM_DEBUG(dbgs() << "[!] no proposal made, no splitting possible!\n");
+ ModuleCallback(cloneAll(M));
+ return;
}
- // This performs all of the partitioning work.
- auto Partitions = doPartitioning(SML, M, N, ModuleCost, FnCosts, WorkList);
- assert(Partitions.size() == N);
+ LLVM_DEBUG(Proposal->print(dbgs()););
- // If we didn't externalize GVs, then local GVs need to be conservatively
- // imported into every module (including their initializers), and then cleaned
- // up afterwards.
- const auto NeedsConservativeImport = [&](const GlobalValue *GV) {
- // We conservatively import private/internal GVs into every module and clean
- // them up afterwards.
- const auto *Var = dyn_cast<GlobalVariable>(GV);
- return Var && Var->hasLocalLinkage();
- };
+ std::optional<raw_fd_ostream> SummariesOS;
+ if (!PartitionSummariesOutput.empty()) {
+ std::error_code EC;
+ SummariesOS.emplace(PartitionSummariesOutput, EC);
+ if (EC)
+ errs() << "[" DEBUG_TYPE "]: cannot open '" << PartitionSummariesOutput
+ << "' - Partition summaries will not be printed\n";
+ }
+
+ for (unsigned PID = 0; PID < NumParts; ++PID) {
+ SplitModuleTimer SMT2("modules_creation",
+ "creating modules for each partition");
+ LLVM_DEBUG(dbgs() << "[split] creating new modules\n");
- SML << "Creating " << N << " modules...\n";
- unsigned TotalFnImpls = 0;
- for (unsigned I = 0; I < N; ++I) {
- const auto &FnsInPart = Partitions[I];
+ DenseSet<const Function *> FnsInPart;
+ for (unsigned NodeID : (*Proposal)[PID].set_bits())
+ FnsInPart.insert(&SG.getNode(NodeID).getFunction());
ValueToValueMapTy VMap;
+ CostType PartCost = 0;
std::unique_ptr<Module> MPart(
CloneModule(M, VMap, [&](const GlobalValue *GV) {
// Functions go in their assigned partition.
- if (const auto *Fn = dyn_cast<Function>(GV))
- return FnsInPart.contains(Fn);
-
- if (NeedsConservativeImport(GV))
- return true;
+ if (const auto *Fn = dyn_cast<Function>(GV)) {
+ if (FnsInPart.contains(Fn)) {
+ PartCost += SG.getCost(*Fn);
+ return true;
+ }
+ return false;
+ }
// Everything else goes in the first partition.
- return I == 0;
+ return needsConservativeImport(GV) || PID == 0;
}));
+ // FIXME: Aliases aren't seen often, and their handling isn't perfect so
+ // bugs are possible.
+
// Clean-up conservatively imported GVs without any users.
- for (auto &GV : make_early_inc_range(MPart->globals())) {
- if (NeedsConservativeImport(&GV) && GV.use_empty())
+ for (auto &GV : make_early_inc_range(MPart->global_values())) {
+ if (needsConservativeImport(&GV) && GV.use_empty())
GV.eraseFromParent();
}
- unsigned NumAllFns = 0, NumKernels = 0;
- for (auto &Cur : *MPart) {
- if (!Cur.isDeclaration()) {
- ++NumAllFns;
- if (isEntryPoint(&Cur))
- ++NumKernels;
- }
- }
- TotalFnImpls += NumAllFns;
- SML << " - Module " << I << " with " << NumAllFns << " functions ("
- << NumKernels << " kernels)\n";
+ if (SummariesOS)
+ printPartitionSummary(*SummariesOS, PID, *MPart, PartCost, ModuleCost);
+
+ LLVM_DEBUG(
+ printPartitionSummary(dbgs(), PID, *MPart, PartCost, ModuleCost));
+
ModuleCallback(std::move(MPart));
}
-
- SML << TotalFnImpls << " function definitions across all modules ("
- << format("%0.2f", (float(TotalFnImpls) / FnCosts.size()) * 100)
- << "% of original module)\n";
}
} // namespace
PreservedAnalyses AMDGPUSplitModulePass::run(Module &M,
ModuleAnalysisManager &MAM) {
+ SplitModuleTimer SMT(
+ "total", "total pass runtime (incl. potentially waiting for lockfile)");
+
FunctionAnalysisManager &FAM =
MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
const auto TTIGetter = [&FAM](Function &F) -> const TargetTransformInfo & {
return FAM.getResult<TargetIRAnalysis>(F);
};
- splitAMDGPUModule(TTIGetter, M, N, ModuleCallback);
- // We don't change the original module.
- return PreservedAnalyses::all();
+
+ bool Done = false;
+#ifndef NDEBUG
+ if (UseLockFile) {
+ SmallString<128> LockFilePath;
+ sys::path::system_temp_directory(/*ErasedOnReboot=*/true, LockFilePath);
+ sys::path::append(LockFilePath, "amdgpu-split-module-debug");
+ LLVM_DEBUG(dbgs() << DEBUG_TYPE " using lockfile '" << LockFilePath
+ << "'\n");
+
+ while (true) {
+ llvm::LockFileManager Locked(LockFilePath.str());
+ switch (Locked) {
+ case LockFileManager::LFS_Error:
+ errs() << "[amdgpu-split-module] unable to acquire lockfile, debug "
+ "output may be mangled by other processes\n";
+ Locked.unsafeRemoveLockFile();
+ break;
+ case LockFileManager::LFS_Owned:
+ break;
+ case LockFileManager::LFS_Shared: {
+ switch (Locked.waitForUnlock()) {
+ case LockFileManager::Res_Success:
+ break;
+ case LockFileManager::Res_OwnerDied:
+ continue; // try again to get the lock.
+ case LockFileManager::Res_Timeout:
+ errs() << "[amdgpu-split-module] unable to acquire lockfile, debug "
+ "output may be mangled by other processes\n";
----------------
arsenm wrote:
Shouldn't spam errors to console, keep it under LLVM_DEBUG?
https://github.com/llvm/llvm-project/pull/104763
More information about the llvm-commits
mailing list