1//===- AMDGPUSplitModule.cpp ----------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9/// \file Implements a module splitting algorithm designed to support the
10/// FullLTO --lto-partitions option for parallel codegen.
11///
12/// The role of this module splitting pass is the same as
13/// lib/Transforms/Utils/SplitModule.cpp: load-balance the module's functions
14/// across a set of N partitions to allow for parallel codegen.
15///
16/// The similarities mostly end here, as this pass achieves load-balancing in a
17/// more elaborate fashion which is targeted towards AMDGPU modules. It can take
18/// advantage of the structure of AMDGPU modules (which are mostly
19/// self-contained) to allow for more efficient splitting without affecting
20/// codegen negatively, or causing innaccurate resource usage analysis.
21///
22/// High-level pass overview:
23/// - SplitGraph & associated classes
24/// - Graph representation of the module and of the dependencies that
25/// matter for splitting.
26/// - RecursiveSearchSplitting
27/// - Core splitting algorithm.
28/// - SplitProposal
29/// - Represents a suggested solution for splitting the input module. These
30/// solutions can be scored to determine the best one when multiple
31/// solutions are available.
32/// - Driver/pass "run" function glues everything together.
33
34#include "AMDGPUSplitModule.h"
35#include "Utils/AMDGPUBaseInfo.h"
36#include "llvm/ADT/EquivalenceClasses.h"
37#include "llvm/ADT/GraphTraits.h"
38#include "llvm/ADT/SmallVector.h"
39#include "llvm/ADT/StringExtras.h"
40#include "llvm/ADT/StringRef.h"
41#include "llvm/Analysis/CallGraph.h"
42#include "llvm/Analysis/TargetTransformInfo.h"
43#include "llvm/IR/Function.h"
44#include "llvm/IR/GlobalAlias.h"
45#include "llvm/IR/InstIterator.h"
46#include "llvm/IR/Instruction.h"
47#include "llvm/IR/Module.h"
48#include "llvm/IR/PassTimingInfo.h"
49#include "llvm/IR/Value.h"
50#include "llvm/Support/Allocator.h"
51#include "llvm/Support/Casting.h"
52#include "llvm/Support/DOTGraphTraits.h"
53#include "llvm/Support/Debug.h"
54#include "llvm/Support/GraphWriter.h"
55#include "llvm/Support/Path.h"
56#include "llvm/Support/Timer.h"
57#include "llvm/Support/raw_ostream.h"
58#include "llvm/Transforms/Utils/Cloning.h"
59#include <cassert>
60#include <cmath>
61#include <utility>
62
63#ifndef NDEBUG
64#include "llvm/Support/LockFileManager.h"
65#endif
66
67#define DEBUG_TYPE "amdgpu-split-module"
68
69namespace llvm {
70namespace {
71
72static cl::opt<unsigned> MaxDepth(
73 "amdgpu-module-splitting-max-depth",
74 cl::desc(
75 "maximum search depth. 0 forces a greedy approach. "
76 "warning: the algorithm is up to O(2^N), where N is the max depth."),
77 cl::init(Val: 8));
78
79static cl::opt<float> LargeFnFactor(
80 "amdgpu-module-splitting-large-threshold", cl::init(Val: 2.0f), cl::Hidden,
81 cl::desc(
82 "when max depth is reached and we can no longer branch out, this "
83 "value determines if a function is worth merging into an already "
84 "existing partition to reduce code duplication. This is a factor "
85 "of the ideal partition size, e.g. 2.0 means we consider the "
86 "function for merging if its cost (including its callees) is 2x the "
87 "size of an ideal partition."));
88
89static cl::opt<float> LargeFnOverlapForMerge(
90 "amdgpu-module-splitting-merge-threshold", cl::init(Val: 0.7f), cl::Hidden,
91 cl::desc("when a function is considered for merging into a partition that "
92 "already contains some of its callees, do the merge if at least "
93 "n% of the code it can reach is already present inside the "
94 "partition; e.g. 0.7 means only merge >70%"));
95
96static cl::opt<bool> NoExternalizeGlobals(
97 "amdgpu-module-splitting-no-externalize-globals", cl::Hidden,
98 cl::desc("disables externalization of global variable with local linkage; "
99 "may cause globals to be duplicated which increases binary size"));
100
101static cl::opt<bool> NoExternalizeOnAddrTaken(
102 "amdgpu-module-splitting-no-externalize-address-taken", cl::Hidden,
103 cl::desc(
104 "disables externalization of functions whose addresses are taken"));
105
106static cl::opt<std::string>
107 ModuleDotCfgOutput("amdgpu-module-splitting-print-module-dotcfg",
108 cl::Hidden,
109 cl::desc("output file to write out the dotgraph "
110 "representation of the input module"));
111
112static cl::opt<std::string> PartitionSummariesOutput(
113 "amdgpu-module-splitting-print-partition-summaries", cl::Hidden,
114 cl::desc("output file to write out a summary of "
115 "the partitions created for each module"));
116
117#ifndef NDEBUG
118static cl::opt<bool>
119 UseLockFile("amdgpu-module-splitting-serial-execution", cl::Hidden,
120 cl::desc("use a lock file so only one process in the system "
121 "can run this pass at once. useful to avoid mangled "
122 "debug output in multithreaded environments."));
123
124static cl::opt<bool>
125 DebugProposalSearch("amdgpu-module-splitting-debug-proposal-search",
126 cl::Hidden,
127 cl::desc("print all proposals received and whether "
128 "they were rejected or accepted"));
129#endif
130
131struct SplitModuleTimer : NamedRegionTimer {
132 SplitModuleTimer(StringRef Name, StringRef Desc)
133 : NamedRegionTimer(Name, Desc, DEBUG_TYPE, "AMDGPU Module Splitting",
134 TimePassesIsEnabled) {}
135};
136
137//===----------------------------------------------------------------------===//
138// Utils
139//===----------------------------------------------------------------------===//
140
141using CostType = InstructionCost::CostType;
142using FunctionsCostMap = DenseMap<const Function *, CostType>;
143using GetTTIFn = function_ref<const TargetTransformInfo &(Function &)>;
144static constexpr unsigned InvalidPID = -1;
145
146/// \param Num numerator
147/// \param Dem denominator
148/// \returns a printable object to print (Num/Dem) using "%0.2f".
149static auto formatRatioOf(CostType Num, CostType Dem) {
150 CostType DemOr1 = Dem ? Dem : 1;
151 return format(Fmt: "%0.2f", Vals: (static_cast<double>(Num) / DemOr1) * 100);
152}
153
154/// Checks whether a given function is non-copyable.
155///
156/// Non-copyable functions cannot be cloned into multiple partitions, and only
157/// one copy of the function can be present across all partitions.
158///
159/// Kernel functions and external functions fall into this category. If we were
160/// to clone them, we would end up with multiple symbol definitions and a very
161/// unhappy linker.
162static bool isNonCopyable(const Function &F) {
163 return F.hasExternalLinkage() || !F.isDefinitionExact() ||
164 AMDGPU::isEntryFunctionCC(CC: F.getCallingConv());
165}
166
167/// Cost analysis function. Calculates the cost of each function in \p M
168///
169/// \param GetTTI Abstract getter for TargetTransformInfo.
170/// \param M Module to analyze.
171/// \param CostMap[out] Resulting Function -> Cost map.
172/// \return The module's total cost.
173static CostType calculateFunctionCosts(GetTTIFn GetTTI, Module &M,
174 FunctionsCostMap &CostMap) {
175 SplitModuleTimer SMT("calculateFunctionCosts", "cost analysis");
176
177 LLVM_DEBUG(dbgs() << "[cost analysis] calculating function costs\n");
178 CostType ModuleCost = 0;
179 [[maybe_unused]] CostType KernelCost = 0;
180
181 for (auto &Fn : M) {
182 if (Fn.isDeclaration())
183 continue;
184
185 CostType FnCost = 0;
186 const auto &TTI = GetTTI(Fn);
187 for (const auto &BB : Fn) {
188 for (const auto &I : BB) {
189 auto Cost =
190 TTI.getInstructionCost(U: &I, CostKind: TargetTransformInfo::TCK_CodeSize);
191 assert(Cost != InstructionCost::getMax());
192 // Assume expensive if we can't tell the cost of an instruction.
193 CostType CostVal = Cost.isValid()
194 ? Cost.getValue()
195 : (CostType)TargetTransformInfo::TCC_Expensive;
196 assert((FnCost + CostVal) >= FnCost && "Overflow!");
197 FnCost += CostVal;
198 }
199 }
200
201 assert(FnCost != 0);
202
203 CostMap[&Fn] = FnCost;
204 assert((ModuleCost + FnCost) >= ModuleCost && "Overflow!");
205 ModuleCost += FnCost;
206
207 if (AMDGPU::isEntryFunctionCC(CC: Fn.getCallingConv()))
208 KernelCost += FnCost;
209 }
210
211 if (CostMap.empty())
212 return 0;
213
214 assert(ModuleCost);
215 LLVM_DEBUG({
216 const CostType FnCost = ModuleCost - KernelCost;
217 dbgs() << " - total module cost is " << ModuleCost << ". kernels cost "
218 << "" << KernelCost << " ("
219 << format("%0.2f", (float(KernelCost) / ModuleCost) * 100)
220 << "% of the module), functions cost " << FnCost << " ("
221 << format("%0.2f", (float(FnCost) / ModuleCost) * 100)
222 << "% of the module)\n";
223 });
224
225 return ModuleCost;
226}
227
228/// \return true if \p F can be indirectly called
229static bool canBeIndirectlyCalled(const Function &F) {
230 if (F.isDeclaration() || AMDGPU::isEntryFunctionCC(CC: F.getCallingConv()))
231 return false;
232 return !F.hasLocalLinkage() ||
233 F.hasAddressTaken(/*PutOffender=*/nullptr,
234 /*IgnoreCallbackUses=*/false,
235 /*IgnoreAssumeLikeCalls=*/true,
236 /*IgnoreLLVMUsed=*/IngoreLLVMUsed: true,
237 /*IgnoreARCAttachedCall=*/false,
238 /*IgnoreCastedDirectCall=*/true);
239}
240
241//===----------------------------------------------------------------------===//
242// Graph-based Module Representation
243//===----------------------------------------------------------------------===//
244
245/// AMDGPUSplitModule's view of the source Module, as a graph of all components
246/// that can be split into different modules.
247///
248/// The most trivial instance of this graph is just the CallGraph of the module,
249/// but it is not guaranteed that the graph is strictly equal to the CG. It
250/// currently always is but it's designed in a way that would eventually allow
251/// us to create abstract nodes, or nodes for different entities such as global
252/// variables or any other meaningful constraint we must consider.
253///
254/// The graph is only mutable by this class, and is generally not modified
255/// after \ref SplitGraph::buildGraph runs. No consumers of the graph can
256/// mutate it.
257class SplitGraph {
258public:
259 class Node;
260
261 enum class EdgeKind : uint8_t {
262 /// The nodes are related through a direct call. This is a "strong" edge as
263 /// it means the Src will directly reference the Dst.
264 DirectCall,
265 /// The nodes are related through an indirect call.
266 /// This is a "weaker" edge and is only considered when traversing the graph
267 /// starting from a kernel. We need this edge for resource usage analysis.
268 ///
269 /// The reason why we have this edge in the first place is due to how
270 /// AMDGPUResourceUsageAnalysis works. In the presence of an indirect call,
271 /// the resource usage of the kernel containing the indirect call is the
272 /// max resource usage of all functions that can be indirectly called.
273 IndirectCall,
274 };
275
276 /// An edge between two nodes. Edges are directional, and tagged with a
277 /// "kind".
278 struct Edge {
279 Edge(Node *Src, Node *Dst, EdgeKind Kind)
280 : Src(Src), Dst(Dst), Kind(Kind) {}
281
282 Node *Src; ///< Source
283 Node *Dst; ///< Destination
284 EdgeKind Kind;
285 };
286
287 using EdgesVec = SmallVector<const Edge *, 0>;
288 using edges_iterator = EdgesVec::const_iterator;
289 using nodes_iterator = const Node *const *;
290
291 SplitGraph(const Module &M, const FunctionsCostMap &CostMap,
292 CostType ModuleCost)
293 : M(M), CostMap(CostMap), ModuleCost(ModuleCost) {}
294
295 void buildGraph(CallGraph &CG);
296
297#ifndef NDEBUG
298 bool verifyGraph() const;
299#endif
300
301 bool empty() const { return Nodes.empty(); }
302 iterator_range<nodes_iterator> nodes() const { return Nodes; }
303 const Node &getNode(unsigned ID) const { return *Nodes[ID]; }
304
305 unsigned getNumNodes() const { return Nodes.size(); }
306 BitVector createNodesBitVector() const { return BitVector(Nodes.size()); }
307
308 const Module &getModule() const { return M; }
309
310 CostType getModuleCost() const { return ModuleCost; }
311 CostType getCost(const Function &F) const { return CostMap.at(Val: &F); }
312
313 /// \returns the aggregated cost of all nodes in \p BV (bits set to 1 = node
314 /// IDs).
315 CostType calculateCost(const BitVector &BV) const;
316
317private:
318 /// Retrieves the node for \p GV in \p Cache, or creates a new node for it and
319 /// updates \p Cache.
320 Node &getNode(DenseMap<const GlobalValue *, Node *> &Cache,
321 const GlobalValue &GV);
322
323 // Create a new edge between two nodes and add it to both nodes.
324 const Edge &createEdge(Node &Src, Node &Dst, EdgeKind EK);
325
326 const Module &M;
327 const FunctionsCostMap &CostMap;
328 CostType ModuleCost;
329
330 // Final list of nodes with stable ordering.
331 SmallVector<Node *> Nodes;
332
333 SpecificBumpPtrAllocator<Node> NodesPool;
334
335 // Edges are trivially destructible objects, so as a small optimization we
336 // use a BumpPtrAllocator which avoids destructor calls but also makes
337 // allocation faster.
338 static_assert(
339 std::is_trivially_destructible_v<Edge>,
340 "Edge must be trivially destructible to use the BumpPtrAllocator");
341 BumpPtrAllocator EdgesPool;
342};
343
344/// Nodes in the SplitGraph contain both incoming, and outgoing edges.
345/// Incoming edges have this node as their Dst, and Outgoing ones have this node
346/// as their Src.
347///
348/// Edge objects are shared by both nodes in Src/Dst. They provide immediate
349/// feedback on how two nodes are related, and in which direction they are
350/// related, which is valuable information to make splitting decisions.
351///
352/// Nodes are fundamentally abstract, and any consumers of the graph should
353/// treat them as such. While a node will be a function most of the time, we
354/// could also create nodes for any other reason. In the future, we could have
355/// single nodes for multiple functions, or nodes for GVs, etc.
356class SplitGraph::Node {
357 friend class SplitGraph;
358
359public:
360 Node(unsigned ID, const GlobalValue &GV, CostType IndividualCost,
361 bool IsNonCopyable)
362 : ID(ID), GV(GV), IndividualCost(IndividualCost),
363 IsNonCopyable(IsNonCopyable), IsEntryFnCC(false), IsGraphEntry(false) {
364 if (auto *Fn = dyn_cast<Function>(Val: &GV))
365 IsEntryFnCC = AMDGPU::isEntryFunctionCC(CC: Fn->getCallingConv());
366 }
367
368 /// An 0-indexed ID for the node. The maximum ID (exclusive) is the number of
369 /// nodes in the graph. This ID can be used as an index in a BitVector.
370 unsigned getID() const { return ID; }
371
372 const Function &getFunction() const { return cast<Function>(Val: GV); }
373
374 /// \returns the cost to import this component into a given module, not
375 /// accounting for any dependencies that may need to be imported as well.
376 CostType getIndividualCost() const { return IndividualCost; }
377
378 bool isNonCopyable() const { return IsNonCopyable; }
379 bool isEntryFunctionCC() const { return IsEntryFnCC; }
380
381 /// \returns whether this is an entry point in the graph. Entry points are
382 /// defined as follows: if you take all entry points in the graph, and iterate
383 /// their dependencies, you are guaranteed to visit all nodes in the graph at
384 /// least once.
385 bool isGraphEntryPoint() const { return IsGraphEntry; }
386
387 StringRef getName() const { return GV.getName(); }
388
389 bool hasAnyIncomingEdges() const { return IncomingEdges.size(); }
390 bool hasAnyIncomingEdgesOfKind(EdgeKind EK) const {
391 return any_of(Range: IncomingEdges, P: [&](const auto *E) { return E->Kind == EK; });
392 }
393
394 bool hasAnyOutgoingEdges() const { return OutgoingEdges.size(); }
395 bool hasAnyOutgoingEdgesOfKind(EdgeKind EK) const {
396 return any_of(Range: OutgoingEdges, P: [&](const auto *E) { return E->Kind == EK; });
397 }
398
399 iterator_range<edges_iterator> incoming_edges() const {
400 return IncomingEdges;
401 }
402
403 iterator_range<edges_iterator> outgoing_edges() const {
404 return OutgoingEdges;
405 }
406
407 bool shouldFollowIndirectCalls() const { return isEntryFunctionCC(); }
408
409 /// Visit all children of this node in a recursive fashion. Also visits Self.
410 /// If \ref shouldFollowIndirectCalls returns false, then this only follows
411 /// DirectCall edges.
412 ///
413 /// \param Visitor Visitor Function.
414 void visitAllDependencies(std::function<void(const Node &)> Visitor) const;
415
416 /// Adds the depedencies of this node in \p BV by setting the bit
417 /// corresponding to each node.
418 ///
419 /// Implemented using \ref visitAllDependencies, hence it follows the same
420 /// rules regarding dependencies traversal.
421 ///
422 /// \param[out] BV The bitvector where the bits should be set.
423 void getDependencies(BitVector &BV) const {
424 visitAllDependencies(Visitor: [&](const Node &N) { BV.set(N.getID()); });
425 }
426
427private:
428 void markAsGraphEntry() { IsGraphEntry = true; }
429
430 unsigned ID;
431 const GlobalValue &GV;
432 CostType IndividualCost;
433 bool IsNonCopyable : 1;
434 bool IsEntryFnCC : 1;
435 bool IsGraphEntry : 1;
436
437 // TODO: Use a single sorted vector (with all incoming/outgoing edges grouped
438 // together)
439 EdgesVec IncomingEdges;
440 EdgesVec OutgoingEdges;
441};
442
443void SplitGraph::Node::visitAllDependencies(
444 std::function<void(const Node &)> Visitor) const {
445 const bool FollowIndirect = shouldFollowIndirectCalls();
446 // FIXME: If this can access SplitGraph in the future, use a BitVector
447 // instead.
448 DenseSet<const Node *> Seen;
449 SmallVector<const Node *, 8> WorkList({this});
450 while (!WorkList.empty()) {
451 const Node *CurN = WorkList.pop_back_val();
452 if (auto [It, Inserted] = Seen.insert(V: CurN); !Inserted)
453 continue;
454
455 Visitor(*CurN);
456
457 for (const Edge *E : CurN->outgoing_edges()) {
458 if (!FollowIndirect && E->Kind == EdgeKind::IndirectCall)
459 continue;
460 WorkList.push_back(Elt: E->Dst);
461 }
462 }
463}
464
465/// Checks if \p I has MD_callees and if it does, parse it and put the function
466/// in \p Callees.
467///
468/// \returns true if there was metadata and it was parsed correctly. false if
469/// there was no MD or if it contained unknown entries and parsing failed.
470/// If this returns false, \p Callees will contain incomplete information
471/// and must not be used.
472static bool handleCalleesMD(const Instruction &I,
473 SetVector<Function *> &Callees) {
474 auto *MD = I.getMetadata(KindID: LLVMContext::MD_callees);
475 if (!MD)
476 return false;
477
478 for (const auto &Op : MD->operands()) {
479 Function *Callee = mdconst::extract_or_null<Function>(MD: Op);
480 if (!Callee)
481 return false;
482 Callees.insert(X: Callee);
483 }
484
485 return true;
486}
487
488void SplitGraph::buildGraph(CallGraph &CG) {
489 SplitModuleTimer SMT("buildGraph", "graph construction");
490 LLVM_DEBUG(
491 dbgs()
492 << "[build graph] constructing graph representation of the input\n");
493
494 // FIXME(?): Is the callgraph really worth using if we have to iterate the
495 // function again whenever it fails to give us enough information?
496
497 // We build the graph by just iterating all functions in the module and
498 // working on their direct callees. At the end, all nodes should be linked
499 // together as expected.
500 DenseMap<const GlobalValue *, Node *> Cache;
501 SmallVector<const Function *> FnsWithIndirectCalls, IndirectlyCallableFns;
502 for (const Function &Fn : M) {
503 if (Fn.isDeclaration())
504 continue;
505
506 // Look at direct callees and create the necessary edges in the graph.
507 SetVector<const Function *> DirectCallees;
508 bool CallsExternal = false;
509 for (auto &CGEntry : *CG[&Fn]) {
510 auto *CGNode = CGEntry.second;
511 if (auto *Callee = CGNode->getFunction()) {
512 if (!Callee->isDeclaration())
513 DirectCallees.insert(X: Callee);
514 } else if (CGNode == CG.getCallsExternalNode())
515 CallsExternal = true;
516 }
517
518 // Keep track of this function if it contains an indirect call and/or if it
519 // can be indirectly called.
520 if (CallsExternal) {
521 LLVM_DEBUG(dbgs() << " [!] callgraph is incomplete for ";
522 Fn.printAsOperand(dbgs());
523 dbgs() << " - analyzing function\n");
524
525 SetVector<Function *> KnownCallees;
526 bool HasUnknownIndirectCall = false;
527 for (const auto &Inst : instructions(F: Fn)) {
528 // look at all calls without a direct callee.
529 const auto *CB = dyn_cast<CallBase>(Val: &Inst);
530 if (!CB || CB->getCalledFunction())
531 continue;
532
533 // inline assembly can be ignored, unless InlineAsmIsIndirectCall is
534 // true.
535 if (CB->isInlineAsm()) {
536 LLVM_DEBUG(dbgs() << " found inline assembly\n");
537 continue;
538 }
539
540 if (handleCalleesMD(I: Inst, Callees&: KnownCallees))
541 continue;
542 // If we failed to parse any !callees MD, or some was missing,
543 // the entire KnownCallees list is now unreliable.
544 KnownCallees.clear();
545
546 // Everything else is handled conservatively. If we fall into the
547 // conservative case don't bother analyzing further.
548 HasUnknownIndirectCall = true;
549 break;
550 }
551
552 if (HasUnknownIndirectCall) {
553 LLVM_DEBUG(dbgs() << " indirect call found\n");
554 FnsWithIndirectCalls.push_back(Elt: &Fn);
555 } else if (!KnownCallees.empty())
556 DirectCallees.insert_range(R&: KnownCallees);
557 }
558
559 Node &N = getNode(Cache, GV: Fn);
560 for (const auto *Callee : DirectCallees)
561 createEdge(Src&: N, Dst&: getNode(Cache, GV: *Callee), EK: EdgeKind::DirectCall);
562
563 if (canBeIndirectlyCalled(F: Fn))
564 IndirectlyCallableFns.push_back(Elt: &Fn);
565 }
566
567 // Post-process functions with indirect calls.
568 for (const Function *Fn : FnsWithIndirectCalls) {
569 for (const Function *Candidate : IndirectlyCallableFns) {
570 Node &Src = getNode(Cache, GV: *Fn);
571 Node &Dst = getNode(Cache, GV: *Candidate);
572 createEdge(Src, Dst, EK: EdgeKind::IndirectCall);
573 }
574 }
575
576 // Now, find all entry points.
577 SmallVector<Node *, 16> CandidateEntryPoints;
578 BitVector NodesReachableByKernels = createNodesBitVector();
579 for (Node *N : Nodes) {
580 // Functions with an Entry CC are always graph entry points too.
581 if (N->isEntryFunctionCC()) {
582 N->markAsGraphEntry();
583 N->getDependencies(BV&: NodesReachableByKernels);
584 } else if (!N->hasAnyIncomingEdgesOfKind(EK: EdgeKind::DirectCall))
585 CandidateEntryPoints.push_back(Elt: N);
586 }
587
588 for (Node *N : CandidateEntryPoints) {
589 // This can be another entry point if it's not reachable by a kernel
590 // TODO: We could sort all of the possible new entries in a stable order
591 // (e.g. by cost), then consume them one by one until
592 // NodesReachableByKernels is all 1s. It'd allow us to avoid
593 // considering some nodes as non-entries in some specific cases.
594 if (!NodesReachableByKernels.test(Idx: N->getID()))
595 N->markAsGraphEntry();
596 }
597
598#ifndef NDEBUG
599 assert(verifyGraph());
600#endif
601}
602
603#ifndef NDEBUG
604bool SplitGraph::verifyGraph() const {
605 unsigned ExpectedID = 0;
606 // Exceptionally using a set here in case IDs are messed up.
607 DenseSet<const Node *> SeenNodes;
608 DenseSet<const Function *> SeenFunctionNodes;
609 for (const Node *N : Nodes) {
610 if (N->getID() != (ExpectedID++)) {
611 errs() << "Node IDs are incorrect!\n";
612 return false;
613 }
614
615 if (!SeenNodes.insert(N).second) {
616 errs() << "Node seen more than once!\n";
617 return false;
618 }
619
620 if (&getNode(N->getID()) != N) {
621 errs() << "getNode doesn't return the right node\n";
622 return false;
623 }
624
625 for (const Edge *E : N->IncomingEdges) {
626 if (!E->Src || !E->Dst || (E->Dst != N) ||
627 (find(E->Src->OutgoingEdges, E) == E->Src->OutgoingEdges.end())) {
628 errs() << "ill-formed incoming edges\n";
629 return false;
630 }
631 }
632
633 for (const Edge *E : N->OutgoingEdges) {
634 if (!E->Src || !E->Dst || (E->Src != N) ||
635 (find(E->Dst->IncomingEdges, E) == E->Dst->IncomingEdges.end())) {
636 errs() << "ill-formed outgoing edges\n";
637 return false;
638 }
639 }
640
641 const Function &Fn = N->getFunction();
642 if (AMDGPU::isEntryFunctionCC(Fn.getCallingConv())) {
643 if (N->hasAnyIncomingEdges()) {
644 errs() << "Kernels cannot have incoming edges\n";
645 return false;
646 }
647 }
648
649 if (Fn.isDeclaration()) {
650 errs() << "declarations shouldn't have nodes!\n";
651 return false;
652 }
653
654 auto [It, Inserted] = SeenFunctionNodes.insert(&Fn);
655 if (!Inserted) {
656 errs() << "one function has multiple nodes!\n";
657 return false;
658 }
659 }
660
661 if (ExpectedID != Nodes.size()) {
662 errs() << "Node IDs out of sync!\n";
663 return false;
664 }
665
666 if (createNodesBitVector().size() != getNumNodes()) {
667 errs() << "nodes bit vector doesn't have the right size!\n";
668 return false;
669 }
670
671 // Check we respect the promise of Node::isKernel
672 BitVector BV = createNodesBitVector();
673 for (const Node *N : nodes()) {
674 if (N->isGraphEntryPoint())
675 N->getDependencies(BV);
676 }
677
678 // Ensure each function in the module has an associated node.
679 for (const auto &Fn : M) {
680 if (!Fn.isDeclaration()) {
681 if (!SeenFunctionNodes.contains(&Fn)) {
682 errs() << "Fn has no associated node in the graph!\n";
683 return false;
684 }
685 }
686 }
687
688 if (!BV.all()) {
689 errs() << "not all nodes are reachable through the graph's entry points!\n";
690 return false;
691 }
692
693 return true;
694}
695#endif
696
697CostType SplitGraph::calculateCost(const BitVector &BV) const {
698 CostType Cost = 0;
699 for (unsigned NodeID : BV.set_bits())
700 Cost += getNode(ID: NodeID).getIndividualCost();
701 return Cost;
702}
703
704SplitGraph::Node &
705SplitGraph::getNode(DenseMap<const GlobalValue *, Node *> &Cache,
706 const GlobalValue &GV) {
707 auto &N = Cache[&GV];
708 if (N)
709 return *N;
710
711 CostType Cost = 0;
712 bool NonCopyable = false;
713 if (const Function *Fn = dyn_cast<Function>(Val: &GV)) {
714 NonCopyable = isNonCopyable(F: *Fn);
715 Cost = CostMap.at(Val: Fn);
716 }
717 N = new (NodesPool.Allocate()) Node(Nodes.size(), GV, Cost, NonCopyable);
718 Nodes.push_back(Elt: N);
719 assert(&getNode(N->getID()) == N);
720 return *N;
721}
722
723const SplitGraph::Edge &SplitGraph::createEdge(Node &Src, Node &Dst,
724 EdgeKind EK) {
725 const Edge *E = new (EdgesPool.Allocate<Edge>(Num: 1)) Edge(&Src, &Dst, EK);
726 Src.OutgoingEdges.push_back(Elt: E);
727 Dst.IncomingEdges.push_back(Elt: E);
728 return *E;
729}
730
731//===----------------------------------------------------------------------===//
732// Split Proposals
733//===----------------------------------------------------------------------===//
734
735/// Represents a module splitting proposal.
736///
737/// Proposals are made of N BitVectors, one for each partition, where each bit
738/// set indicates that the node is present and should be copied inside that
739/// partition.
740///
741/// Proposals have several metrics attached so they can be compared/sorted,
742/// which the driver to try multiple strategies resultings in multiple proposals
743/// and choose the best one out of them.
744class SplitProposal {
745public:
746 SplitProposal(const SplitGraph &SG, unsigned MaxPartitions) : SG(&SG) {
747 Partitions.resize(new_size: MaxPartitions, x: {0, SG.createNodesBitVector()});
748 }
749
750 void setName(StringRef NewName) { Name = NewName; }
751 StringRef getName() const { return Name; }
752
753 const BitVector &operator[](unsigned PID) const {
754 return Partitions[PID].second;
755 }
756
757 void add(unsigned PID, const BitVector &BV) {
758 Partitions[PID].second |= BV;
759 updateScore(PID);
760 }
761
762 void print(raw_ostream &OS) const;
763 LLVM_DUMP_METHOD void dump() const { print(OS&: dbgs()); }
764
765 // Find the cheapest partition (lowest cost). In case of ties, always returns
766 // the highest partition number.
767 unsigned findCheapestPartition() const;
768
769 /// Calculate the CodeSize and Bottleneck scores.
770 void calculateScores();
771
772#ifndef NDEBUG
773 void verifyCompleteness() const;
774#endif
775
776 /// Only available after \ref calculateScores is called.
777 ///
778 /// A positive number indicating the % of code duplication that this proposal
779 /// creates. e.g. 0.2 means this proposal adds roughly 20% code size by
780 /// duplicating some functions across partitions.
781 ///
782 /// Value is always rounded up to 3 decimal places.
783 ///
784 /// A perfect score would be 0.0, and anything approaching 1.0 is very bad.
785 double getCodeSizeScore() const { return CodeSizeScore; }
786
787 /// Only available after \ref calculateScores is called.
788 ///
789 /// A number between [0, 1] which indicates how big of a bottleneck is
790 /// expected from the largest partition.
791 ///
792 /// A score of 1.0 means the biggest partition is as big as the source module,
793 /// so build time will be equal to or greater than the build time of the
794 /// initial input.
795 ///
796 /// Value is always rounded up to 3 decimal places.
797 ///
798 /// This is one of the metrics used to estimate this proposal's build time.
799 double getBottleneckScore() const { return BottleneckScore; }
800
801private:
802 void updateScore(unsigned PID) {
803 assert(SG);
804 for (auto &[PCost, Nodes] : Partitions) {
805 TotalCost -= PCost;
806 PCost = SG->calculateCost(BV: Nodes);
807 TotalCost += PCost;
808 }
809 }
810
811 /// \see getCodeSizeScore
812 double CodeSizeScore = 0.0;
813 /// \see getBottleneckScore
814 double BottleneckScore = 0.0;
815 /// Aggregated cost of all partitions
816 CostType TotalCost = 0;
817
818 const SplitGraph *SG = nullptr;
819 std::string Name;
820
821 std::vector<std::pair<CostType, BitVector>> Partitions;
822};
823
824void SplitProposal::print(raw_ostream &OS) const {
825 assert(SG);
826
827 OS << "[proposal] " << Name << ", total cost:" << TotalCost
828 << ", code size score:" << format(Fmt: "%0.3f", Vals: CodeSizeScore)
829 << ", bottleneck score:" << format(Fmt: "%0.3f", Vals: BottleneckScore) << '\n';
830 for (const auto &[PID, Part] : enumerate(First: Partitions)) {
831 const auto &[Cost, NodeIDs] = Part;
832 OS << " - P" << PID << " nodes:" << NodeIDs.count() << " cost: " << Cost
833 << '|' << formatRatioOf(Num: Cost, Dem: SG->getModuleCost()) << "%\n";
834 }
835}
836
837unsigned SplitProposal::findCheapestPartition() const {
838 assert(!Partitions.empty());
839 CostType CurCost = std::numeric_limits<CostType>::max();
840 unsigned CurPID = InvalidPID;
841 for (const auto &[Idx, Part] : enumerate(First: Partitions)) {
842 if (Part.first <= CurCost) {
843 CurPID = Idx;
844 CurCost = Part.first;
845 }
846 }
847 assert(CurPID != InvalidPID);
848 return CurPID;
849}
850
851void SplitProposal::calculateScores() {
852 if (Partitions.empty())
853 return;
854
855 assert(SG);
856 CostType LargestPCost = 0;
857 for (auto &[PCost, Nodes] : Partitions) {
858 if (PCost > LargestPCost)
859 LargestPCost = PCost;
860 }
861
862 CostType ModuleCost = SG->getModuleCost();
863 CodeSizeScore = double(TotalCost) / ModuleCost;
864 assert(CodeSizeScore >= 0.0);
865
866 BottleneckScore = double(LargestPCost) / ModuleCost;
867
868 CodeSizeScore = std::ceil(x: CodeSizeScore * 100.0) / 100.0;
869 BottleneckScore = std::ceil(x: BottleneckScore * 100.0) / 100.0;
870}
871
872#ifndef NDEBUG
873void SplitProposal::verifyCompleteness() const {
874 if (Partitions.empty())
875 return;
876
877 BitVector Result = Partitions[0].second;
878 for (const auto &P : drop_begin(Partitions))
879 Result |= P.second;
880 assert(Result.all() && "some nodes are missing from this proposal!");
881}
882#endif
883
884//===-- RecursiveSearchStrategy -------------------------------------------===//
885
886/// Partitioning algorithm.
887///
888/// This is a recursive search algorithm that can explore multiple possiblities.
889///
890/// When a cluster of nodes can go into more than one partition, and we haven't
891/// reached maximum search depth, we recurse and explore both options and their
892/// consequences. Both branches will yield a proposal, and the driver will grade
893/// both and choose the best one.
894///
895/// If max depth is reached, we will use some heuristics to make a choice. Most
896/// of the time we will just use the least-pressured (cheapest) partition, but
897/// if a cluster is particularly big and there is a good amount of overlap with
898/// an existing partition, we will choose that partition instead.
899class RecursiveSearchSplitting {
900public:
901 using SubmitProposalFn = function_ref<void(SplitProposal)>;
902
903 RecursiveSearchSplitting(const SplitGraph &SG, unsigned NumParts,
904 SubmitProposalFn SubmitProposal);
905
906 void run();
907
908private:
909 struct WorkListEntry {
910 WorkListEntry(const BitVector &BV) : Cluster(BV) {}
911
912 unsigned NumNonEntryNodes = 0;
913 CostType TotalCost = 0;
914 CostType CostExcludingGraphEntryPoints = 0;
915 BitVector Cluster;
916 };
917
918 /// Collects all graph entry points's clusters and sort them so the most
919 /// expensive clusters are viewed first. This will merge clusters together if
920 /// they share a non-copyable dependency.
921 void setupWorkList();
922
923 /// Recursive function that assigns the worklist item at \p Idx into a
924 /// partition of \p SP.
925 ///
926 /// \p Depth is the current search depth. When this value is equal to
927 /// \ref MaxDepth, we can no longer recurse.
928 ///
929 /// This function only recurses if there is more than one possible assignment,
930 /// otherwise it is iterative to avoid creating a call stack that is as big as
931 /// \ref WorkList.
932 void pickPartition(unsigned Depth, unsigned Idx, SplitProposal SP);
933
934 /// \return A pair: first element is the PID of the partition that has the
935 /// most similarities with \p Entry, or \ref InvalidPID if no partition was
936 /// found with at least one element in common. The second element is the
937 /// aggregated cost of all dependencies in common between \p Entry and that
938 /// partition.
939 std::pair<unsigned, CostType>
940 findMostSimilarPartition(const WorkListEntry &Entry, const SplitProposal &SP);
941
942 const SplitGraph &SG;
943 unsigned NumParts;
944 SubmitProposalFn SubmitProposal;
945
946 // A Cluster is considered large when its cost, excluding entry points,
947 // exceeds this value.
948 CostType LargeClusterThreshold = 0;
949 unsigned NumProposalsSubmitted = 0;
950 SmallVector<WorkListEntry> WorkList;
951};
952
953RecursiveSearchSplitting::RecursiveSearchSplitting(
954 const SplitGraph &SG, unsigned NumParts, SubmitProposalFn SubmitProposal)
955 : SG(SG), NumParts(NumParts), SubmitProposal(SubmitProposal) {
956 // arbitrary max value as a safeguard. Anything above 10 will already be
957 // slow, this is just a max value to prevent extreme resource exhaustion or
958 // unbounded run time.
959 if (MaxDepth > 16)
960 report_fatal_error(reason: "[amdgpu-split-module] search depth of " +
961 Twine(MaxDepth) + " is too high!");
962 LargeClusterThreshold =
963 (LargeFnFactor != 0.0)
964 ? CostType(((SG.getModuleCost() / NumParts) * LargeFnFactor))
965 : std::numeric_limits<CostType>::max();
966 LLVM_DEBUG(dbgs() << "[recursive search] large cluster threshold set at "
967 << LargeClusterThreshold << "\n");
968}
969
970void RecursiveSearchSplitting::run() {
971 {
972 SplitModuleTimer SMT("recursive_search_prepare", "preparing worklist");
973 setupWorkList();
974 }
975
976 {
977 SplitModuleTimer SMT("recursive_search_pick", "partitioning");
978 SplitProposal SP(SG, NumParts);
979 pickPartition(/*BranchDepth=*/Depth: 0, /*Idx=*/0, SP: std::move(SP));
980 }
981}
982
983void RecursiveSearchSplitting::setupWorkList() {
984 // e.g. if A and B are two worklist item, and they both call a non copyable
985 // dependency C, this does:
986 // A=C
987 // B=C
988 // => NodeEC will create a single group (A, B, C) and we create a new
989 // WorkList entry for that group.
990
991 EquivalenceClasses<unsigned> NodeEC;
992 for (const SplitGraph::Node *N : SG.nodes()) {
993 if (!N->isGraphEntryPoint())
994 continue;
995
996 NodeEC.insert(Data: N->getID());
997 N->visitAllDependencies(Visitor: [&](const SplitGraph::Node &Dep) {
998 if (&Dep != N && Dep.isNonCopyable())
999 NodeEC.unionSets(V1: N->getID(), V2: Dep.getID());
1000 });
1001 }
1002
1003 for (const auto &Node : NodeEC) {
1004 if (!Node->isLeader())
1005 continue;
1006
1007 BitVector Cluster = SG.createNodesBitVector();
1008 for (unsigned M : NodeEC.members(ECV: *Node)) {
1009 const SplitGraph::Node &N = SG.getNode(ID: M);
1010 if (N.isGraphEntryPoint())
1011 N.getDependencies(BV&: Cluster);
1012 }
1013 WorkList.emplace_back(Args: std::move(Cluster));
1014 }
1015
1016 // Calculate costs and other useful information.
1017 for (WorkListEntry &Entry : WorkList) {
1018 for (unsigned NodeID : Entry.Cluster.set_bits()) {
1019 const SplitGraph::Node &N = SG.getNode(ID: NodeID);
1020 const CostType Cost = N.getIndividualCost();
1021
1022 Entry.TotalCost += Cost;
1023 if (!N.isGraphEntryPoint()) {
1024 Entry.CostExcludingGraphEntryPoints += Cost;
1025 ++Entry.NumNonEntryNodes;
1026 }
1027 }
1028 }
1029
1030 stable_sort(Range&: WorkList, C: [](const WorkListEntry &A, const WorkListEntry &B) {
1031 if (A.TotalCost != B.TotalCost)
1032 return A.TotalCost > B.TotalCost;
1033
1034 if (A.CostExcludingGraphEntryPoints != B.CostExcludingGraphEntryPoints)
1035 return A.CostExcludingGraphEntryPoints > B.CostExcludingGraphEntryPoints;
1036
1037 if (A.NumNonEntryNodes != B.NumNonEntryNodes)
1038 return A.NumNonEntryNodes > B.NumNonEntryNodes;
1039
1040 return A.Cluster.count() > B.Cluster.count();
1041 });
1042
1043 LLVM_DEBUG({
1044 dbgs() << "[recursive search] worklist:\n";
1045 for (const auto &[Idx, Entry] : enumerate(WorkList)) {
1046 dbgs() << " - [" << Idx << "]: ";
1047 for (unsigned NodeID : Entry.Cluster.set_bits())
1048 dbgs() << NodeID << " ";
1049 dbgs() << "(total_cost:" << Entry.TotalCost
1050 << ", cost_excl_entries:" << Entry.CostExcludingGraphEntryPoints
1051 << ")\n";
1052 }
1053 });
1054}
1055
1056void RecursiveSearchSplitting::pickPartition(unsigned Depth, unsigned Idx,
1057 SplitProposal SP) {
1058 while (Idx < WorkList.size()) {
1059 // Step 1: Determine candidate PIDs.
1060 //
1061 const WorkListEntry &Entry = WorkList[Idx];
1062 const BitVector &Cluster = Entry.Cluster;
1063
1064 // Default option is to do load-balancing, AKA assign to least pressured
1065 // partition.
1066 const unsigned CheapestPID = SP.findCheapestPartition();
1067 assert(CheapestPID != InvalidPID);
1068
1069 // Explore assigning to the kernel that contains the most dependencies in
1070 // common.
1071 const auto [MostSimilarPID, SimilarDepsCost] =
1072 findMostSimilarPartition(Entry, SP);
1073
1074 // We can chose to explore only one path if we only have one valid path, or
1075 // if we reached maximum search depth and can no longer branch out.
1076 unsigned SinglePIDToTry = InvalidPID;
1077 if (MostSimilarPID == InvalidPID) // no similar PID found
1078 SinglePIDToTry = CheapestPID;
1079 else if (MostSimilarPID == CheapestPID) // both landed on the same PID
1080 SinglePIDToTry = CheapestPID;
1081 else if (Depth >= MaxDepth) {
1082 // We have to choose one path. Use a heuristic to guess which one will be
1083 // more appropriate.
1084 if (Entry.CostExcludingGraphEntryPoints > LargeClusterThreshold) {
1085 // Check if the amount of code in common makes it worth it.
1086 assert(SimilarDepsCost && Entry.CostExcludingGraphEntryPoints);
1087 const double Ratio = static_cast<double>(SimilarDepsCost) /
1088 Entry.CostExcludingGraphEntryPoints;
1089 assert(Ratio >= 0.0 && Ratio <= 1.0);
1090 if (Ratio > LargeFnOverlapForMerge) {
1091 // For debug, just print "L", so we'll see "L3=P3" for instance, which
1092 // will mean we reached max depth and chose P3 based on this
1093 // heuristic.
1094 LLVM_DEBUG(dbgs() << 'L');
1095 SinglePIDToTry = MostSimilarPID;
1096 }
1097 } else
1098 SinglePIDToTry = CheapestPID;
1099 }
1100
1101 // Step 2: Explore candidates.
1102
1103 // When we only explore one possible path, and thus branch depth doesn't
1104 // increase, do not recurse, iterate instead.
1105 if (SinglePIDToTry != InvalidPID) {
1106 LLVM_DEBUG(dbgs() << Idx << "=P" << SinglePIDToTry << ' ');
1107 // Only one path to explore, don't clone SP, don't increase depth.
1108 SP.add(PID: SinglePIDToTry, BV: Cluster);
1109 ++Idx;
1110 continue;
1111 }
1112
1113 assert(MostSimilarPID != InvalidPID);
1114
1115 // We explore multiple paths: recurse at increased depth, then stop this
1116 // function.
1117
1118 LLVM_DEBUG(dbgs() << '\n');
1119
1120 // lb = load balancing = put in cheapest partition
1121 {
1122 SplitProposal BranchSP = SP;
1123 LLVM_DEBUG(dbgs().indent(Depth)
1124 << " [lb] " << Idx << "=P" << CheapestPID << "? ");
1125 BranchSP.add(PID: CheapestPID, BV: Cluster);
1126 pickPartition(Depth: Depth + 1, Idx: Idx + 1, SP: std::move(BranchSP));
1127 }
1128
1129 // ms = most similar = put in partition with the most in common
1130 {
1131 SplitProposal BranchSP = SP;
1132 LLVM_DEBUG(dbgs().indent(Depth)
1133 << " [ms] " << Idx << "=P" << MostSimilarPID << "? ");
1134 BranchSP.add(PID: MostSimilarPID, BV: Cluster);
1135 pickPartition(Depth: Depth + 1, Idx: Idx + 1, SP: std::move(BranchSP));
1136 }
1137
1138 return;
1139 }
1140
1141 // Step 3: If we assigned all WorkList items, submit the proposal.
1142
1143 assert(Idx == WorkList.size());
1144 assert(NumProposalsSubmitted <= (2u << MaxDepth) &&
1145 "Search got out of bounds?");
1146 SP.setName("recursive_search (depth=" + std::to_string(val: Depth) + ") #" +
1147 std::to_string(val: NumProposalsSubmitted++));
1148 LLVM_DEBUG(dbgs() << '\n');
1149 SubmitProposal(std::move(SP));
1150}
1151
1152std::pair<unsigned, CostType>
1153RecursiveSearchSplitting::findMostSimilarPartition(const WorkListEntry &Entry,
1154 const SplitProposal &SP) {
1155 if (!Entry.NumNonEntryNodes)
1156 return {InvalidPID, 0};
1157
1158 // We take the partition that is the most similar using Cost as a metric.
1159 // So we take the set of nodes in common, compute their aggregated cost, and
1160 // pick the partition with the highest cost in common.
1161 unsigned ChosenPID = InvalidPID;
1162 CostType ChosenCost = 0;
1163 for (unsigned PID = 0; PID < NumParts; ++PID) {
1164 BitVector BV = SP[PID];
1165 BV &= Entry.Cluster; // FIXME: & doesn't work between BVs?!
1166
1167 if (BV.none())
1168 continue;
1169
1170 const CostType Cost = SG.calculateCost(BV);
1171
1172 if (ChosenPID == InvalidPID || ChosenCost < Cost ||
1173 (ChosenCost == Cost && PID > ChosenPID)) {
1174 ChosenPID = PID;
1175 ChosenCost = Cost;
1176 }
1177 }
1178
1179 return {ChosenPID, ChosenCost};
1180}
1181
1182//===----------------------------------------------------------------------===//
1183// DOTGraph Printing Support
1184//===----------------------------------------------------------------------===//
1185
1186const SplitGraph::Node *mapEdgeToDst(const SplitGraph::Edge *E) {
1187 return E->Dst;
1188}
1189
1190using SplitGraphEdgeDstIterator =
1191 mapped_iterator<SplitGraph::edges_iterator, decltype(&mapEdgeToDst)>;
1192
1193} // namespace
1194
1195template <> struct GraphTraits<SplitGraph> {
1196 using NodeRef = const SplitGraph::Node *;
1197 using nodes_iterator = SplitGraph::nodes_iterator;
1198 using ChildIteratorType = SplitGraphEdgeDstIterator;
1199
1200 using EdgeRef = const SplitGraph::Edge *;
1201 using ChildEdgeIteratorType = SplitGraph::edges_iterator;
1202
1203 static NodeRef getEntryNode(NodeRef N) { return N; }
1204
1205 static ChildIteratorType child_begin(NodeRef Ref) {
1206 return {Ref->outgoing_edges().begin(), mapEdgeToDst};
1207 }
1208 static ChildIteratorType child_end(NodeRef Ref) {
1209 return {Ref->outgoing_edges().end(), mapEdgeToDst};
1210 }
1211
1212 static nodes_iterator nodes_begin(const SplitGraph &G) {
1213 return G.nodes().begin();
1214 }
1215 static nodes_iterator nodes_end(const SplitGraph &G) {
1216 return G.nodes().end();
1217 }
1218};
1219
1220template <> struct DOTGraphTraits<SplitGraph> : public DefaultDOTGraphTraits {
1221 DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
1222
1223 static std::string getGraphName(const SplitGraph &SG) {
1224 return SG.getModule().getName().str();
1225 }
1226
1227 std::string getNodeLabel(const SplitGraph::Node *N, const SplitGraph &SG) {
1228 return N->getName().str();
1229 }
1230
1231 static std::string getNodeDescription(const SplitGraph::Node *N,
1232 const SplitGraph &SG) {
1233 std::string Result;
1234 if (N->isEntryFunctionCC())
1235 Result += "entry-fn-cc ";
1236 if (N->isNonCopyable())
1237 Result += "non-copyable ";
1238 Result += "cost:" + std::to_string(val: N->getIndividualCost());
1239 return Result;
1240 }
1241
1242 static std::string getNodeAttributes(const SplitGraph::Node *N,
1243 const SplitGraph &SG) {
1244 return N->hasAnyIncomingEdges() ? "" : "color=\"red\"";
1245 }
1246
1247 static std::string getEdgeAttributes(const SplitGraph::Node *N,
1248 SplitGraphEdgeDstIterator EI,
1249 const SplitGraph &SG) {
1250
1251 switch ((*EI.getCurrent())->Kind) {
1252 case SplitGraph::EdgeKind::DirectCall:
1253 return "";
1254 case SplitGraph::EdgeKind::IndirectCall:
1255 return "style=\"dashed\"";
1256 }
1257 llvm_unreachable("Unknown SplitGraph::EdgeKind enum");
1258 }
1259};
1260
1261//===----------------------------------------------------------------------===//
1262// Driver
1263//===----------------------------------------------------------------------===//
1264
1265namespace {
1266
1267// If we didn't externalize GVs, then local GVs need to be conservatively
1268// imported into every module (including their initializers), and then cleaned
1269// up afterwards.
1270static bool needsConservativeImport(const GlobalValue *GV) {
1271 if (const auto *Var = dyn_cast<GlobalVariable>(Val: GV))
1272 return Var->hasLocalLinkage();
1273 if (const auto *GA = dyn_cast<GlobalAlias>(Val: GV))
1274 return GA->hasLocalLinkage();
1275 return false;
1276}
1277
1278/// Prints a summary of the partition \p N, represented by module \p M, to \p
1279/// OS.
1280static void printPartitionSummary(raw_ostream &OS, unsigned N, const Module &M,
1281 unsigned PartCost, unsigned ModuleCost) {
1282 OS << "*** Partition P" << N << " ***\n";
1283
1284 for (const auto &Fn : M) {
1285 if (!Fn.isDeclaration())
1286 OS << " - [function] " << Fn.getName() << "\n";
1287 }
1288
1289 for (const auto &GV : M.globals()) {
1290 if (GV.hasInitializer())
1291 OS << " - [global] " << GV.getName() << "\n";
1292 }
1293
1294 OS << "Partition contains " << formatRatioOf(Num: PartCost, Dem: ModuleCost)
1295 << "% of the source\n";
1296}
1297
1298static void evaluateProposal(SplitProposal &Best, SplitProposal New) {
1299 SplitModuleTimer SMT("proposal_evaluation", "proposal ranking algorithm");
1300
1301 LLVM_DEBUG({
1302 New.verifyCompleteness();
1303 if (DebugProposalSearch)
1304 New.print(dbgs());
1305 });
1306
1307 const double CurBScore = Best.getBottleneckScore();
1308 const double CurCSScore = Best.getCodeSizeScore();
1309 const double NewBScore = New.getBottleneckScore();
1310 const double NewCSScore = New.getCodeSizeScore();
1311
1312 // TODO: Improve this
1313 // We can probably lower the precision of the comparison at first
1314 // e.g. if we have
1315 // - (Current): BScore: 0.489 CSCore 1.105
1316 // - (New): BScore: 0.475 CSCore 1.305
1317 // Currently we'd choose the new one because the bottleneck score is
1318 // lower, but the new one duplicates more code. It may be worth it to
1319 // discard the new proposal as the impact on build time is negligible.
1320
1321 // Compare them
1322 bool IsBest = false;
1323 if (NewBScore < CurBScore)
1324 IsBest = true;
1325 else if (NewBScore == CurBScore)
1326 IsBest = (NewCSScore < CurCSScore); // Use code size as tie breaker.
1327
1328 if (IsBest)
1329 Best = std::move(New);
1330
1331 LLVM_DEBUG(if (DebugProposalSearch) {
1332 if (IsBest)
1333 dbgs() << "[search] new best proposal!\n";
1334 else
1335 dbgs() << "[search] discarding - not profitable\n";
1336 });
1337}
1338
1339/// Trivial helper to create an identical copy of \p M.
1340static std::unique_ptr<Module> cloneAll(const Module &M) {
1341 ValueToValueMapTy VMap;
1342 return CloneModule(M, VMap, ShouldCloneDefinition: [&](const GlobalValue *GV) { return true; });
1343}
1344
1345/// Writes \p SG as a DOTGraph to \ref ModuleDotCfgDir if requested.
1346static void writeDOTGraph(const SplitGraph &SG) {
1347 if (ModuleDotCfgOutput.empty())
1348 return;
1349
1350 std::error_code EC;
1351 raw_fd_ostream OS(ModuleDotCfgOutput, EC);
1352 if (EC) {
1353 errs() << "[" DEBUG_TYPE "]: cannot open '" << ModuleDotCfgOutput
1354 << "' - DOTGraph will not be printed\n";
1355 }
1356 WriteGraph(O&: OS, G: SG, /*ShortName=*/ShortNames: false,
1357 /*Title=*/SG.getModule().getName());
1358}
1359
1360static void splitAMDGPUModule(
1361 GetTTIFn GetTTI, Module &M, unsigned NumParts,
1362 function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback) {
1363 CallGraph CG(M);
1364
1365 // Externalize functions whose address are taken.
1366 //
1367 // This is needed because partitioning is purely based on calls, but sometimes
1368 // a kernel/function may just look at the address of another local function
1369 // and not do anything (no calls). After partitioning, that local function may
1370 // end up in a different module (so it's just a declaration in the module
1371 // where its address is taken), which emits a "undefined hidden symbol" linker
1372 // error.
1373 //
1374 // Additionally, it guides partitioning to not duplicate this function if it's
1375 // called directly at some point.
1376 //
1377 // TODO: Could we be smarter about this ? This makes all functions whose
1378 // addresses are taken non-copyable. We should probably model this type of
1379 // constraint in the graph and use it to guide splitting, instead of
1380 // externalizing like this. Maybe non-copyable should really mean "keep one
1381 // visible copy, then internalize all other copies" for some functions?
1382 if (!NoExternalizeOnAddrTaken) {
1383 for (auto &Fn : M) {
1384 if (Fn.hasLocalLinkage() && Fn.hasAddressTaken()) {
1385 LLVM_DEBUG(dbgs() << "[externalize] "; Fn.printAsOperand(dbgs());
1386 dbgs() << " because its address is taken\n");
1387 Fn.externalize();
1388 }
1389 }
1390 }
1391
1392 // Externalize local GVs, which avoids duplicating their initializers, which
1393 // in turns helps keep code size in check.
1394 if (!NoExternalizeGlobals) {
1395 for (auto &GV : M.globals()) {
1396 if (GV.hasLocalLinkage())
1397 LLVM_DEBUG(dbgs() << "[externalize] GV " << GV.getName() << '\n');
1398 GV.externalize();
1399 }
1400 }
1401
1402 for (auto &GA : M.aliases()) {
1403 if (GA.hasLocalLinkage()) {
1404 LLVM_DEBUG(dbgs() << "[externalize] alias " << GA.getName() << '\n');
1405 GA.externalize();
1406 }
1407 }
1408
1409 // Start by calculating the cost of every function in the module, as well as
1410 // the module's overall cost.
1411 FunctionsCostMap FnCosts;
1412 const CostType ModuleCost = calculateFunctionCosts(GetTTI, M, CostMap&: FnCosts);
1413
1414 // Build the SplitGraph, which represents the module's functions and models
1415 // their dependencies accurately.
1416 SplitGraph SG(M, FnCosts, ModuleCost);
1417 SG.buildGraph(CG);
1418
1419 if (SG.empty()) {
1420 LLVM_DEBUG(
1421 dbgs()
1422 << "[!] no nodes in graph, input is empty - no splitting possible\n");
1423 ModuleCallback(cloneAll(M));
1424 return;
1425 }
1426
1427 LLVM_DEBUG({
1428 dbgs() << "[graph] nodes:\n";
1429 for (const SplitGraph::Node *N : SG.nodes()) {
1430 dbgs() << " - [" << N->getID() << "]: " << N->getName() << " "
1431 << (N->isGraphEntryPoint() ? "(entry)" : "") << " "
1432 << (N->isNonCopyable() ? "(noncopyable)" : "") << "\n";
1433 }
1434 });
1435
1436 writeDOTGraph(SG);
1437
1438 LLVM_DEBUG(dbgs() << "[search] testing splitting strategies\n");
1439
1440 std::optional<SplitProposal> Proposal;
1441 const auto EvaluateProposal = [&](SplitProposal SP) {
1442 SP.calculateScores();
1443 if (!Proposal)
1444 Proposal = std::move(SP);
1445 else
1446 evaluateProposal(Best&: *Proposal, New: std::move(SP));
1447 };
1448
1449 // TODO: It would be very easy to create new strategies by just adding a base
1450 // class to RecursiveSearchSplitting and abstracting it away.
1451 RecursiveSearchSplitting(SG, NumParts, EvaluateProposal).run();
1452 LLVM_DEBUG(if (Proposal) dbgs() << "[search done] selected proposal: "
1453 << Proposal->getName() << "\n";);
1454
1455 if (!Proposal) {
1456 LLVM_DEBUG(dbgs() << "[!] no proposal made, no splitting possible!\n");
1457 ModuleCallback(cloneAll(M));
1458 return;
1459 }
1460
1461 LLVM_DEBUG(Proposal->print(dbgs()););
1462
1463 std::optional<raw_fd_ostream> SummariesOS;
1464 if (!PartitionSummariesOutput.empty()) {
1465 std::error_code EC;
1466 SummariesOS.emplace(args&: PartitionSummariesOutput, args&: EC);
1467 if (EC)
1468 errs() << "[" DEBUG_TYPE "]: cannot open '" << PartitionSummariesOutput
1469 << "' - Partition summaries will not be printed\n";
1470 }
1471
1472 // One module will import all GlobalValues that are not Functions
1473 // and are not subject to conservative import.
1474 bool ImportAllGVs = true;
1475
1476 for (unsigned PID = 0; PID < NumParts; ++PID) {
1477 SplitModuleTimer SMT2("modules_creation",
1478 "creating modules for each partition");
1479 LLVM_DEBUG(dbgs() << "[split] creating new modules\n");
1480
1481 DenseSet<const Function *> FnsInPart;
1482 for (unsigned NodeID : (*Proposal)[PID].set_bits())
1483 FnsInPart.insert(V: &SG.getNode(ID: NodeID).getFunction());
1484
1485 // Don't create empty modules.
1486 if (FnsInPart.empty()) {
1487 LLVM_DEBUG(dbgs() << "[split] P" << PID
1488 << " is empty, not creating module\n");
1489 continue;
1490 }
1491
1492 ValueToValueMapTy VMap;
1493 CostType PartCost = 0;
1494 std::unique_ptr<Module> MPart(
1495 CloneModule(M, VMap, ShouldCloneDefinition: [&](const GlobalValue *GV) {
1496 // Functions go in their assigned partition.
1497 if (const auto *Fn = dyn_cast<Function>(Val: GV)) {
1498 if (FnsInPart.contains(V: Fn)) {
1499 PartCost += SG.getCost(F: *Fn);
1500 return true;
1501 }
1502 return false;
1503 }
1504
1505 // Aliases should not be separated from their underlying object.
1506 if (const auto *GA = dyn_cast<GlobalAlias>(Val: GV)) {
1507 if (const auto *Fn = dyn_cast<Function>(Val: GA->getAliaseeObject()))
1508 return FnsInPart.contains(V: Fn);
1509 }
1510
1511 // Everything else goes in the first non-empty module we create.
1512 return ImportAllGVs || needsConservativeImport(GV);
1513 }));
1514
1515 ImportAllGVs = false;
1516
1517 // Clean-up conservatively imported GVs without any users.
1518 for (auto &GV : make_early_inc_range(Range: MPart->global_values())) {
1519 if (needsConservativeImport(GV: &GV) && GV.use_empty())
1520 GV.eraseFromParent();
1521 }
1522
1523 if (SummariesOS)
1524 printPartitionSummary(OS&: *SummariesOS, N: PID, M: *MPart, PartCost, ModuleCost);
1525
1526 LLVM_DEBUG(
1527 printPartitionSummary(dbgs(), PID, *MPart, PartCost, ModuleCost));
1528
1529 ModuleCallback(std::move(MPart));
1530 }
1531}
1532} // namespace
1533
1534PreservedAnalyses AMDGPUSplitModulePass::run(Module &M,
1535 ModuleAnalysisManager &MAM) {
1536 SplitModuleTimer SMT(
1537 "total", "total pass runtime (incl. potentially waiting for lockfile)");
1538
1539 FunctionAnalysisManager &FAM =
1540 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
1541 const auto TTIGetter = [&FAM](Function &F) -> const TargetTransformInfo & {
1542 return FAM.getResult<TargetIRAnalysis>(IR&: F);
1543 };
1544
1545 bool Done = false;
1546#ifndef NDEBUG
1547 if (UseLockFile) {
1548 SmallString<128> LockFilePath;
1549 sys::path::system_temp_directory(/*ErasedOnReboot=*/true, LockFilePath);
1550 sys::path::append(LockFilePath, "amdgpu-split-module-debug");
1551 LLVM_DEBUG(dbgs() << DEBUG_TYPE " using lockfile '" << LockFilePath
1552 << "'\n");
1553
1554 while (true) {
1555 llvm::LockFileManager Lock(LockFilePath.str());
1556 bool Owned;
1557 if (Error Err = Lock.tryLock().moveInto(Owned)) {
1558 consumeError(std::move(Err));
1559 LLVM_DEBUG(
1560 dbgs() << "[amdgpu-split-module] unable to acquire lockfile, debug "
1561 "output may be mangled by other processes\n");
1562 } else if (!Owned) {
1563 switch (Lock.waitForUnlockFor(std::chrono::seconds(90))) {
1564 case WaitForUnlockResult::Success:
1565 break;
1566 case WaitForUnlockResult::OwnerDied:
1567 continue; // try again to get the lock.
1568 case WaitForUnlockResult::Timeout:
1569 LLVM_DEBUG(
1570 dbgs()
1571 << "[amdgpu-split-module] unable to acquire lockfile, debug "
1572 "output may be mangled by other processes\n");
1573 Lock.unsafeUnlock();
1574 break; // give up
1575 }
1576 }
1577
1578 splitAMDGPUModule(TTIGetter, M, N, ModuleCallback);
1579 Done = true;
1580 break;
1581 }
1582 }
1583#endif
1584
1585 if (!Done)
1586 splitAMDGPUModule(GetTTI: TTIGetter, M, NumParts: N, ModuleCallback);
1587
1588 // We can change linkage/visibilities in the input, consider that nothing is
1589 // preserved just to be safe. This pass runs last anyway.
1590 return PreservedAnalyses::none();
1591}
1592} // namespace llvm
1593