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 iterator_range<edges_iterator> outgoing_edges() const {
395 return OutgoingEdges;
396 }
397
398 bool shouldFollowIndirectCalls() const { return isEntryFunctionCC(); }
399
400 /// Visit all children of this node in a recursive fashion. Also visits Self.
401 /// If \ref shouldFollowIndirectCalls returns false, then this only follows
402 /// DirectCall edges.
403 ///
404 /// \param Visitor Visitor Function.
405 void visitAllDependencies(std::function<void(const Node &)> Visitor) const;
406
407 /// Adds the depedencies of this node in \p BV by setting the bit
408 /// corresponding to each node.
409 ///
410 /// Implemented using \ref visitAllDependencies, hence it follows the same
411 /// rules regarding dependencies traversal.
412 ///
413 /// \param[out] BV The bitvector where the bits should be set.
414 void getDependencies(BitVector &BV) const {
415 visitAllDependencies(Visitor: [&](const Node &N) { BV.set(N.getID()); });
416 }
417
418private:
419 void markAsGraphEntry() { IsGraphEntry = true; }
420
421 unsigned ID;
422 const GlobalValue &GV;
423 CostType IndividualCost;
424 bool IsNonCopyable : 1;
425 bool IsEntryFnCC : 1;
426 bool IsGraphEntry : 1;
427
428 // TODO: Use a single sorted vector (with all incoming/outgoing edges grouped
429 // together)
430 EdgesVec IncomingEdges;
431 EdgesVec OutgoingEdges;
432};
433
434void SplitGraph::Node::visitAllDependencies(
435 std::function<void(const Node &)> Visitor) const {
436 const bool FollowIndirect = shouldFollowIndirectCalls();
437 // FIXME: If this can access SplitGraph in the future, use a BitVector
438 // instead.
439 DenseSet<const Node *> Seen;
440 SmallVector<const Node *, 8> WorkList({this});
441 while (!WorkList.empty()) {
442 const Node *CurN = WorkList.pop_back_val();
443 if (auto [It, Inserted] = Seen.insert(V: CurN); !Inserted)
444 continue;
445
446 Visitor(*CurN);
447
448 for (const Edge *E : CurN->outgoing_edges()) {
449 if (!FollowIndirect && E->Kind == EdgeKind::IndirectCall)
450 continue;
451 WorkList.push_back(Elt: E->Dst);
452 }
453 }
454}
455
456/// Checks if \p I has MD_callees and if it does, parse it and put the function
457/// in \p Callees.
458///
459/// \returns true if there was metadata and it was parsed correctly. false if
460/// there was no MD or if it contained unknown entries and parsing failed.
461/// If this returns false, \p Callees will contain incomplete information
462/// and must not be used.
463static bool handleCalleesMD(const Instruction &I,
464 SetVector<Function *> &Callees) {
465 auto *MD = I.getMetadata(KindID: LLVMContext::MD_callees);
466 if (!MD)
467 return false;
468
469 for (const auto &Op : MD->operands()) {
470 Function *Callee = mdconst::extract_or_null<Function>(MD: Op);
471 if (!Callee)
472 return false;
473 Callees.insert(X: Callee);
474 }
475
476 return true;
477}
478
479void SplitGraph::buildGraph(CallGraph &CG) {
480 SplitModuleTimer SMT("buildGraph", "graph construction");
481 LLVM_DEBUG(
482 dbgs()
483 << "[build graph] constructing graph representation of the input\n");
484
485 // FIXME(?): Is the callgraph really worth using if we have to iterate the
486 // function again whenever it fails to give us enough information?
487
488 // We build the graph by just iterating all functions in the module and
489 // working on their direct callees. At the end, all nodes should be linked
490 // together as expected.
491 DenseMap<const GlobalValue *, Node *> Cache;
492 SmallVector<const Function *> FnsWithIndirectCalls, IndirectlyCallableFns;
493 for (const Function &Fn : M) {
494 if (Fn.isDeclaration())
495 continue;
496
497 // Look at direct callees and create the necessary edges in the graph.
498 SetVector<const Function *> DirectCallees;
499 bool CallsExternal = false;
500 for (auto &CGEntry : *CG[&Fn]) {
501 auto *CGNode = CGEntry.second;
502 if (auto *Callee = CGNode->getFunction()) {
503 if (!Callee->isDeclaration())
504 DirectCallees.insert(X: Callee);
505 } else if (CGNode == CG.getCallsExternalNode())
506 CallsExternal = true;
507 }
508
509 // Keep track of this function if it contains an indirect call and/or if it
510 // can be indirectly called.
511 if (CallsExternal) {
512 LLVM_DEBUG(dbgs() << " [!] callgraph is incomplete for ";
513 Fn.printAsOperand(dbgs());
514 dbgs() << " - analyzing function\n");
515
516 SetVector<Function *> KnownCallees;
517 bool HasUnknownIndirectCall = false;
518 for (const auto &Inst : instructions(F: Fn)) {
519 // look at all calls without a direct callee.
520 const auto *CB = dyn_cast<CallBase>(Val: &Inst);
521 if (!CB || CB->getCalledFunction())
522 continue;
523
524 // inline assembly can be ignored, unless InlineAsmIsIndirectCall is
525 // true.
526 if (CB->isInlineAsm()) {
527 LLVM_DEBUG(dbgs() << " found inline assembly\n");
528 continue;
529 }
530
531 if (handleCalleesMD(I: Inst, Callees&: KnownCallees))
532 continue;
533 // If we failed to parse any !callees MD, or some was missing,
534 // the entire KnownCallees list is now unreliable.
535 KnownCallees.clear();
536
537 // Everything else is handled conservatively. If we fall into the
538 // conservative case don't bother analyzing further.
539 HasUnknownIndirectCall = true;
540 break;
541 }
542
543 if (HasUnknownIndirectCall) {
544 LLVM_DEBUG(dbgs() << " indirect call found\n");
545 FnsWithIndirectCalls.push_back(Elt: &Fn);
546 } else if (!KnownCallees.empty())
547 DirectCallees.insert_range(R&: KnownCallees);
548 }
549
550 Node &N = getNode(Cache, GV: Fn);
551 for (const auto *Callee : DirectCallees)
552 createEdge(Src&: N, Dst&: getNode(Cache, GV: *Callee), EK: EdgeKind::DirectCall);
553
554 if (canBeIndirectlyCalled(F: Fn))
555 IndirectlyCallableFns.push_back(Elt: &Fn);
556 }
557
558 // Post-process functions with indirect calls.
559 for (const Function *Fn : FnsWithIndirectCalls) {
560 for (const Function *Candidate : IndirectlyCallableFns) {
561 Node &Src = getNode(Cache, GV: *Fn);
562 Node &Dst = getNode(Cache, GV: *Candidate);
563 createEdge(Src, Dst, EK: EdgeKind::IndirectCall);
564 }
565 }
566
567 // Now, find all entry points.
568 SmallVector<Node *, 16> CandidateEntryPoints;
569 BitVector NodesReachableByKernels = createNodesBitVector();
570 for (Node *N : Nodes) {
571 // Functions with an Entry CC are always graph entry points too.
572 if (N->isEntryFunctionCC()) {
573 N->markAsGraphEntry();
574 N->getDependencies(BV&: NodesReachableByKernels);
575 } else if (!N->hasAnyIncomingEdgesOfKind(EK: EdgeKind::DirectCall))
576 CandidateEntryPoints.push_back(Elt: N);
577 }
578
579 for (Node *N : CandidateEntryPoints) {
580 // This can be another entry point if it's not reachable by a kernel
581 // TODO: We could sort all of the possible new entries in a stable order
582 // (e.g. by cost), then consume them one by one until
583 // NodesReachableByKernels is all 1s. It'd allow us to avoid
584 // considering some nodes as non-entries in some specific cases.
585 if (!NodesReachableByKernels.test(Idx: N->getID()))
586 N->markAsGraphEntry();
587 }
588
589#ifndef NDEBUG
590 assert(verifyGraph());
591#endif
592}
593
594#ifndef NDEBUG
595bool SplitGraph::verifyGraph() const {
596 unsigned ExpectedID = 0;
597 // Exceptionally using a set here in case IDs are messed up.
598 DenseSet<const Node *> SeenNodes;
599 DenseSet<const Function *> SeenFunctionNodes;
600 for (const Node *N : Nodes) {
601 if (N->getID() != (ExpectedID++)) {
602 errs() << "Node IDs are incorrect!\n";
603 return false;
604 }
605
606 if (!SeenNodes.insert(N).second) {
607 errs() << "Node seen more than once!\n";
608 return false;
609 }
610
611 if (&getNode(N->getID()) != N) {
612 errs() << "getNode doesn't return the right node\n";
613 return false;
614 }
615
616 for (const Edge *E : N->IncomingEdges) {
617 if (!E->Src || !E->Dst || (E->Dst != N) ||
618 (find(E->Src->OutgoingEdges, E) == E->Src->OutgoingEdges.end())) {
619 errs() << "ill-formed incoming edges\n";
620 return false;
621 }
622 }
623
624 for (const Edge *E : N->OutgoingEdges) {
625 if (!E->Src || !E->Dst || (E->Src != N) ||
626 (find(E->Dst->IncomingEdges, E) == E->Dst->IncomingEdges.end())) {
627 errs() << "ill-formed outgoing edges\n";
628 return false;
629 }
630 }
631
632 const Function &Fn = N->getFunction();
633 if (AMDGPU::isEntryFunctionCC(Fn.getCallingConv())) {
634 if (N->hasAnyIncomingEdges()) {
635 errs() << "Kernels cannot have incoming edges\n";
636 return false;
637 }
638 }
639
640 if (Fn.isDeclaration()) {
641 errs() << "declarations shouldn't have nodes!\n";
642 return false;
643 }
644
645 auto [It, Inserted] = SeenFunctionNodes.insert(&Fn);
646 if (!Inserted) {
647 errs() << "one function has multiple nodes!\n";
648 return false;
649 }
650 }
651
652 if (ExpectedID != Nodes.size()) {
653 errs() << "Node IDs out of sync!\n";
654 return false;
655 }
656
657 if (createNodesBitVector().size() != getNumNodes()) {
658 errs() << "nodes bit vector doesn't have the right size!\n";
659 return false;
660 }
661
662 // Check we respect the promise of Node::isKernel
663 BitVector BV = createNodesBitVector();
664 for (const Node *N : nodes()) {
665 if (N->isGraphEntryPoint())
666 N->getDependencies(BV);
667 }
668
669 // Ensure each function in the module has an associated node.
670 for (const auto &Fn : M) {
671 if (!Fn.isDeclaration()) {
672 if (!SeenFunctionNodes.contains(&Fn)) {
673 errs() << "Fn has no associated node in the graph!\n";
674 return false;
675 }
676 }
677 }
678
679 if (!BV.all()) {
680 errs() << "not all nodes are reachable through the graph's entry points!\n";
681 return false;
682 }
683
684 return true;
685}
686#endif
687
688CostType SplitGraph::calculateCost(const BitVector &BV) const {
689 CostType Cost = 0;
690 for (unsigned NodeID : BV.set_bits())
691 Cost += getNode(ID: NodeID).getIndividualCost();
692 return Cost;
693}
694
695SplitGraph::Node &
696SplitGraph::getNode(DenseMap<const GlobalValue *, Node *> &Cache,
697 const GlobalValue &GV) {
698 auto &N = Cache[&GV];
699 if (N)
700 return *N;
701
702 CostType Cost = 0;
703 bool NonCopyable = false;
704 if (const Function *Fn = dyn_cast<Function>(Val: &GV)) {
705 NonCopyable = isNonCopyable(F: *Fn);
706 Cost = CostMap.at(Val: Fn);
707 }
708 N = new (NodesPool.Allocate()) Node(Nodes.size(), GV, Cost, NonCopyable);
709 Nodes.push_back(Elt: N);
710 assert(&getNode(N->getID()) == N);
711 return *N;
712}
713
714const SplitGraph::Edge &SplitGraph::createEdge(Node &Src, Node &Dst,
715 EdgeKind EK) {
716 const Edge *E = new (EdgesPool.Allocate<Edge>(Num: 1)) Edge(&Src, &Dst, EK);
717 Src.OutgoingEdges.push_back(Elt: E);
718 Dst.IncomingEdges.push_back(Elt: E);
719 return *E;
720}
721
722//===----------------------------------------------------------------------===//
723// Split Proposals
724//===----------------------------------------------------------------------===//
725
726/// Represents a module splitting proposal.
727///
728/// Proposals are made of N BitVectors, one for each partition, where each bit
729/// set indicates that the node is present and should be copied inside that
730/// partition.
731///
732/// Proposals have several metrics attached so they can be compared/sorted,
733/// which the driver to try multiple strategies resultings in multiple proposals
734/// and choose the best one out of them.
735class SplitProposal {
736public:
737 SplitProposal(const SplitGraph &SG, unsigned MaxPartitions) : SG(&SG) {
738 Partitions.resize(new_size: MaxPartitions, x: {0, SG.createNodesBitVector()});
739 }
740
741 void setName(StringRef NewName) { Name = NewName; }
742 StringRef getName() const { return Name; }
743
744 const BitVector &operator[](unsigned PID) const {
745 return Partitions[PID].second;
746 }
747
748 void add(unsigned PID, const BitVector &BV) {
749 Partitions[PID].second |= BV;
750 updateScore(PID);
751 }
752
753 void print(raw_ostream &OS) const;
754 LLVM_DUMP_METHOD void dump() const { print(OS&: dbgs()); }
755
756 // Find the cheapest partition (lowest cost). In case of ties, always returns
757 // the highest partition number.
758 unsigned findCheapestPartition() const;
759
760 /// Calculate the CodeSize and Bottleneck scores.
761 void calculateScores();
762
763#ifndef NDEBUG
764 void verifyCompleteness() const;
765#endif
766
767 /// Only available after \ref calculateScores is called.
768 ///
769 /// A positive number indicating the % of code duplication that this proposal
770 /// creates. e.g. 0.2 means this proposal adds roughly 20% code size by
771 /// duplicating some functions across partitions.
772 ///
773 /// Value is always rounded up to 3 decimal places.
774 ///
775 /// A perfect score would be 0.0, and anything approaching 1.0 is very bad.
776 double getCodeSizeScore() const { return CodeSizeScore; }
777
778 /// Only available after \ref calculateScores is called.
779 ///
780 /// A number between [0, 1] which indicates how big of a bottleneck is
781 /// expected from the largest partition.
782 ///
783 /// A score of 1.0 means the biggest partition is as big as the source module,
784 /// so build time will be equal to or greater than the build time of the
785 /// initial input.
786 ///
787 /// Value is always rounded up to 3 decimal places.
788 ///
789 /// This is one of the metrics used to estimate this proposal's build time.
790 double getBottleneckScore() const { return BottleneckScore; }
791
792private:
793 void updateScore(unsigned PID) {
794 assert(SG);
795 for (auto &[PCost, Nodes] : Partitions) {
796 TotalCost -= PCost;
797 PCost = SG->calculateCost(BV: Nodes);
798 TotalCost += PCost;
799 }
800 }
801
802 /// \see getCodeSizeScore
803 double CodeSizeScore = 0.0;
804 /// \see getBottleneckScore
805 double BottleneckScore = 0.0;
806 /// Aggregated cost of all partitions
807 CostType TotalCost = 0;
808
809 const SplitGraph *SG = nullptr;
810 std::string Name;
811
812 std::vector<std::pair<CostType, BitVector>> Partitions;
813};
814
815void SplitProposal::print(raw_ostream &OS) const {
816 assert(SG);
817
818 OS << "[proposal] " << Name << ", total cost:" << TotalCost
819 << ", code size score:" << format(Fmt: "%0.3f", Vals: CodeSizeScore)
820 << ", bottleneck score:" << format(Fmt: "%0.3f", Vals: BottleneckScore) << '\n';
821 for (const auto &[PID, Part] : enumerate(First: Partitions)) {
822 const auto &[Cost, NodeIDs] = Part;
823 OS << " - P" << PID << " nodes:" << NodeIDs.count() << " cost: " << Cost
824 << '|' << formatRatioOf(Num: Cost, Dem: SG->getModuleCost()) << "%\n";
825 }
826}
827
828unsigned SplitProposal::findCheapestPartition() const {
829 assert(!Partitions.empty());
830 CostType CurCost = std::numeric_limits<CostType>::max();
831 unsigned CurPID = InvalidPID;
832 for (const auto &[Idx, Part] : enumerate(First: Partitions)) {
833 if (Part.first <= CurCost) {
834 CurPID = Idx;
835 CurCost = Part.first;
836 }
837 }
838 assert(CurPID != InvalidPID);
839 return CurPID;
840}
841
842void SplitProposal::calculateScores() {
843 if (Partitions.empty())
844 return;
845
846 assert(SG);
847 CostType LargestPCost = 0;
848 for (auto &[PCost, Nodes] : Partitions) {
849 if (PCost > LargestPCost)
850 LargestPCost = PCost;
851 }
852
853 CostType ModuleCost = SG->getModuleCost();
854 CodeSizeScore = double(TotalCost) / ModuleCost;
855 assert(CodeSizeScore >= 0.0);
856
857 BottleneckScore = double(LargestPCost) / ModuleCost;
858
859 CodeSizeScore = std::ceil(x: CodeSizeScore * 100.0) / 100.0;
860 BottleneckScore = std::ceil(x: BottleneckScore * 100.0) / 100.0;
861}
862
863#ifndef NDEBUG
864void SplitProposal::verifyCompleteness() const {
865 if (Partitions.empty())
866 return;
867
868 BitVector Result = Partitions[0].second;
869 for (const auto &P : drop_begin(Partitions))
870 Result |= P.second;
871 assert(Result.all() && "some nodes are missing from this proposal!");
872}
873#endif
874
875//===-- RecursiveSearchStrategy -------------------------------------------===//
876
877/// Partitioning algorithm.
878///
879/// This is a recursive search algorithm that can explore multiple possiblities.
880///
881/// When a cluster of nodes can go into more than one partition, and we haven't
882/// reached maximum search depth, we recurse and explore both options and their
883/// consequences. Both branches will yield a proposal, and the driver will grade
884/// both and choose the best one.
885///
886/// If max depth is reached, we will use some heuristics to make a choice. Most
887/// of the time we will just use the least-pressured (cheapest) partition, but
888/// if a cluster is particularly big and there is a good amount of overlap with
889/// an existing partition, we will choose that partition instead.
890class RecursiveSearchSplitting {
891public:
892 using SubmitProposalFn = function_ref<void(SplitProposal)>;
893
894 RecursiveSearchSplitting(const SplitGraph &SG, unsigned NumParts,
895 SubmitProposalFn SubmitProposal);
896
897 void run();
898
899private:
900 struct WorkListEntry {
901 WorkListEntry(const BitVector &BV) : Cluster(BV) {}
902
903 unsigned NumNonEntryNodes = 0;
904 CostType TotalCost = 0;
905 CostType CostExcludingGraphEntryPoints = 0;
906 BitVector Cluster;
907 };
908
909 /// Collects all graph entry points's clusters and sort them so the most
910 /// expensive clusters are viewed first. This will merge clusters together if
911 /// they share a non-copyable dependency.
912 void setupWorkList();
913
914 /// Recursive function that assigns the worklist item at \p Idx into a
915 /// partition of \p SP.
916 ///
917 /// \p Depth is the current search depth. When this value is equal to
918 /// \ref MaxDepth, we can no longer recurse.
919 ///
920 /// This function only recurses if there is more than one possible assignment,
921 /// otherwise it is iterative to avoid creating a call stack that is as big as
922 /// \ref WorkList.
923 void pickPartition(unsigned Depth, unsigned Idx, SplitProposal SP);
924
925 /// \return A pair: first element is the PID of the partition that has the
926 /// most similarities with \p Entry, or \ref InvalidPID if no partition was
927 /// found with at least one element in common. The second element is the
928 /// aggregated cost of all dependencies in common between \p Entry and that
929 /// partition.
930 std::pair<unsigned, CostType>
931 findMostSimilarPartition(const WorkListEntry &Entry, const SplitProposal &SP);
932
933 const SplitGraph &SG;
934 unsigned NumParts;
935 SubmitProposalFn SubmitProposal;
936
937 // A Cluster is considered large when its cost, excluding entry points,
938 // exceeds this value.
939 CostType LargeClusterThreshold = 0;
940 unsigned NumProposalsSubmitted = 0;
941 SmallVector<WorkListEntry> WorkList;
942};
943
944RecursiveSearchSplitting::RecursiveSearchSplitting(
945 const SplitGraph &SG, unsigned NumParts, SubmitProposalFn SubmitProposal)
946 : SG(SG), NumParts(NumParts), SubmitProposal(SubmitProposal) {
947 // arbitrary max value as a safeguard. Anything above 10 will already be
948 // slow, this is just a max value to prevent extreme resource exhaustion or
949 // unbounded run time.
950 if (MaxDepth > 16)
951 report_fatal_error(reason: "[amdgpu-split-module] search depth of " +
952 Twine(MaxDepth) + " is too high!");
953 LargeClusterThreshold =
954 (LargeFnFactor != 0.0)
955 ? CostType(((SG.getModuleCost() / NumParts) * LargeFnFactor))
956 : std::numeric_limits<CostType>::max();
957 LLVM_DEBUG(dbgs() << "[recursive search] large cluster threshold set at "
958 << LargeClusterThreshold << "\n");
959}
960
961void RecursiveSearchSplitting::run() {
962 {
963 SplitModuleTimer SMT("recursive_search_prepare", "preparing worklist");
964 setupWorkList();
965 }
966
967 {
968 SplitModuleTimer SMT("recursive_search_pick", "partitioning");
969 SplitProposal SP(SG, NumParts);
970 pickPartition(/*BranchDepth=*/Depth: 0, /*Idx=*/0, SP: std::move(SP));
971 }
972}
973
974void RecursiveSearchSplitting::setupWorkList() {
975 // e.g. if A and B are two worklist item, and they both call a non copyable
976 // dependency C, this does:
977 // A=C
978 // B=C
979 // => NodeEC will create a single group (A, B, C) and we create a new
980 // WorkList entry for that group.
981
982 EquivalenceClasses<unsigned> NodeEC;
983 for (const SplitGraph::Node *N : SG.nodes()) {
984 if (!N->isGraphEntryPoint())
985 continue;
986
987 NodeEC.insert(Data: N->getID());
988 N->visitAllDependencies(Visitor: [&](const SplitGraph::Node &Dep) {
989 if (&Dep != N && Dep.isNonCopyable())
990 NodeEC.unionSets(V1: N->getID(), V2: Dep.getID());
991 });
992 }
993
994 for (const auto &Node : NodeEC) {
995 if (!Node->isLeader())
996 continue;
997
998 BitVector Cluster = SG.createNodesBitVector();
999 for (unsigned M : NodeEC.members(ECV: *Node)) {
1000 const SplitGraph::Node &N = SG.getNode(ID: M);
1001 if (N.isGraphEntryPoint())
1002 N.getDependencies(BV&: Cluster);
1003 }
1004 WorkList.emplace_back(Args: std::move(Cluster));
1005 }
1006
1007 // Calculate costs and other useful information.
1008 for (WorkListEntry &Entry : WorkList) {
1009 for (unsigned NodeID : Entry.Cluster.set_bits()) {
1010 const SplitGraph::Node &N = SG.getNode(ID: NodeID);
1011 const CostType Cost = N.getIndividualCost();
1012
1013 Entry.TotalCost += Cost;
1014 if (!N.isGraphEntryPoint()) {
1015 Entry.CostExcludingGraphEntryPoints += Cost;
1016 ++Entry.NumNonEntryNodes;
1017 }
1018 }
1019 }
1020
1021 stable_sort(Range&: WorkList, C: [](const WorkListEntry &A, const WorkListEntry &B) {
1022 if (A.TotalCost != B.TotalCost)
1023 return A.TotalCost > B.TotalCost;
1024
1025 if (A.CostExcludingGraphEntryPoints != B.CostExcludingGraphEntryPoints)
1026 return A.CostExcludingGraphEntryPoints > B.CostExcludingGraphEntryPoints;
1027
1028 if (A.NumNonEntryNodes != B.NumNonEntryNodes)
1029 return A.NumNonEntryNodes > B.NumNonEntryNodes;
1030
1031 return A.Cluster.count() > B.Cluster.count();
1032 });
1033
1034 LLVM_DEBUG({
1035 dbgs() << "[recursive search] worklist:\n";
1036 for (const auto &[Idx, Entry] : enumerate(WorkList)) {
1037 dbgs() << " - [" << Idx << "]: ";
1038 for (unsigned NodeID : Entry.Cluster.set_bits())
1039 dbgs() << NodeID << " ";
1040 dbgs() << "(total_cost:" << Entry.TotalCost
1041 << ", cost_excl_entries:" << Entry.CostExcludingGraphEntryPoints
1042 << ")\n";
1043 }
1044 });
1045}
1046
1047void RecursiveSearchSplitting::pickPartition(unsigned Depth, unsigned Idx,
1048 SplitProposal SP) {
1049 while (Idx < WorkList.size()) {
1050 // Step 1: Determine candidate PIDs.
1051 //
1052 const WorkListEntry &Entry = WorkList[Idx];
1053 const BitVector &Cluster = Entry.Cluster;
1054
1055 // Default option is to do load-balancing, AKA assign to least pressured
1056 // partition.
1057 const unsigned CheapestPID = SP.findCheapestPartition();
1058 assert(CheapestPID != InvalidPID);
1059
1060 // Explore assigning to the kernel that contains the most dependencies in
1061 // common.
1062 const auto [MostSimilarPID, SimilarDepsCost] =
1063 findMostSimilarPartition(Entry, SP);
1064
1065 // We can chose to explore only one path if we only have one valid path, or
1066 // if we reached maximum search depth and can no longer branch out.
1067 unsigned SinglePIDToTry = InvalidPID;
1068 if (MostSimilarPID == InvalidPID) // no similar PID found
1069 SinglePIDToTry = CheapestPID;
1070 else if (MostSimilarPID == CheapestPID) // both landed on the same PID
1071 SinglePIDToTry = CheapestPID;
1072 else if (Depth >= MaxDepth) {
1073 // We have to choose one path. Use a heuristic to guess which one will be
1074 // more appropriate.
1075 if (Entry.CostExcludingGraphEntryPoints > LargeClusterThreshold) {
1076 // Check if the amount of code in common makes it worth it.
1077 assert(SimilarDepsCost && Entry.CostExcludingGraphEntryPoints);
1078 const double Ratio = static_cast<double>(SimilarDepsCost) /
1079 Entry.CostExcludingGraphEntryPoints;
1080 assert(Ratio >= 0.0 && Ratio <= 1.0);
1081 if (Ratio > LargeFnOverlapForMerge) {
1082 // For debug, just print "L", so we'll see "L3=P3" for instance, which
1083 // will mean we reached max depth and chose P3 based on this
1084 // heuristic.
1085 LLVM_DEBUG(dbgs() << 'L');
1086 SinglePIDToTry = MostSimilarPID;
1087 }
1088 } else
1089 SinglePIDToTry = CheapestPID;
1090 }
1091
1092 // Step 2: Explore candidates.
1093
1094 // When we only explore one possible path, and thus branch depth doesn't
1095 // increase, do not recurse, iterate instead.
1096 if (SinglePIDToTry != InvalidPID) {
1097 LLVM_DEBUG(dbgs() << Idx << "=P" << SinglePIDToTry << ' ');
1098 // Only one path to explore, don't clone SP, don't increase depth.
1099 SP.add(PID: SinglePIDToTry, BV: Cluster);
1100 ++Idx;
1101 continue;
1102 }
1103
1104 assert(MostSimilarPID != InvalidPID);
1105
1106 // We explore multiple paths: recurse at increased depth, then stop this
1107 // function.
1108
1109 LLVM_DEBUG(dbgs() << '\n');
1110
1111 // lb = load balancing = put in cheapest partition
1112 {
1113 SplitProposal BranchSP = SP;
1114 LLVM_DEBUG(dbgs().indent(Depth)
1115 << " [lb] " << Idx << "=P" << CheapestPID << "? ");
1116 BranchSP.add(PID: CheapestPID, BV: Cluster);
1117 pickPartition(Depth: Depth + 1, Idx: Idx + 1, SP: std::move(BranchSP));
1118 }
1119
1120 // ms = most similar = put in partition with the most in common
1121 {
1122 SplitProposal BranchSP = SP;
1123 LLVM_DEBUG(dbgs().indent(Depth)
1124 << " [ms] " << Idx << "=P" << MostSimilarPID << "? ");
1125 BranchSP.add(PID: MostSimilarPID, BV: Cluster);
1126 pickPartition(Depth: Depth + 1, Idx: Idx + 1, SP: std::move(BranchSP));
1127 }
1128
1129 return;
1130 }
1131
1132 // Step 3: If we assigned all WorkList items, submit the proposal.
1133
1134 assert(Idx == WorkList.size());
1135 assert(NumProposalsSubmitted <= (2u << MaxDepth) &&
1136 "Search got out of bounds?");
1137 SP.setName("recursive_search (depth=" + std::to_string(val: Depth) + ") #" +
1138 std::to_string(val: NumProposalsSubmitted++));
1139 LLVM_DEBUG(dbgs() << '\n');
1140 SubmitProposal(std::move(SP));
1141}
1142
1143std::pair<unsigned, CostType>
1144RecursiveSearchSplitting::findMostSimilarPartition(const WorkListEntry &Entry,
1145 const SplitProposal &SP) {
1146 if (!Entry.NumNonEntryNodes)
1147 return {InvalidPID, 0};
1148
1149 // We take the partition that is the most similar using Cost as a metric.
1150 // So we take the set of nodes in common, compute their aggregated cost, and
1151 // pick the partition with the highest cost in common.
1152 unsigned ChosenPID = InvalidPID;
1153 CostType ChosenCost = 0;
1154 for (unsigned PID = 0; PID < NumParts; ++PID) {
1155 BitVector BV = SP[PID];
1156 BV &= Entry.Cluster; // FIXME: & doesn't work between BVs?!
1157
1158 if (BV.none())
1159 continue;
1160
1161 const CostType Cost = SG.calculateCost(BV);
1162
1163 if (ChosenPID == InvalidPID || ChosenCost < Cost ||
1164 (ChosenCost == Cost && PID > ChosenPID)) {
1165 ChosenPID = PID;
1166 ChosenCost = Cost;
1167 }
1168 }
1169
1170 return {ChosenPID, ChosenCost};
1171}
1172
1173//===----------------------------------------------------------------------===//
1174// DOTGraph Printing Support
1175//===----------------------------------------------------------------------===//
1176
1177const SplitGraph::Node *mapEdgeToDst(const SplitGraph::Edge *E) {
1178 return E->Dst;
1179}
1180
1181using SplitGraphEdgeDstIterator =
1182 mapped_iterator<SplitGraph::edges_iterator, decltype(&mapEdgeToDst)>;
1183
1184} // namespace
1185
1186template <> struct GraphTraits<SplitGraph> {
1187 using NodeRef = const SplitGraph::Node *;
1188 using nodes_iterator = SplitGraph::nodes_iterator;
1189 using ChildIteratorType = SplitGraphEdgeDstIterator;
1190
1191 using EdgeRef = const SplitGraph::Edge *;
1192 using ChildEdgeIteratorType = SplitGraph::edges_iterator;
1193
1194 static NodeRef getEntryNode(NodeRef N) { return N; }
1195
1196 static ChildIteratorType child_begin(NodeRef Ref) {
1197 return {Ref->outgoing_edges().begin(), mapEdgeToDst};
1198 }
1199 static ChildIteratorType child_end(NodeRef Ref) {
1200 return {Ref->outgoing_edges().end(), mapEdgeToDst};
1201 }
1202
1203 static nodes_iterator nodes_begin(const SplitGraph &G) {
1204 return G.nodes().begin();
1205 }
1206 static nodes_iterator nodes_end(const SplitGraph &G) {
1207 return G.nodes().end();
1208 }
1209};
1210
1211template <> struct DOTGraphTraits<SplitGraph> : public DefaultDOTGraphTraits {
1212 DOTGraphTraits(bool IsSimple = false) : DefaultDOTGraphTraits(IsSimple) {}
1213
1214 static std::string getGraphName(const SplitGraph &SG) {
1215 return SG.getModule().getName().str();
1216 }
1217
1218 std::string getNodeLabel(const SplitGraph::Node *N, const SplitGraph &SG) {
1219 return N->getName().str();
1220 }
1221
1222 static std::string getNodeDescription(const SplitGraph::Node *N,
1223 const SplitGraph &SG) {
1224 std::string Result;
1225 if (N->isEntryFunctionCC())
1226 Result += "entry-fn-cc ";
1227 if (N->isNonCopyable())
1228 Result += "non-copyable ";
1229 Result += "cost:" + std::to_string(val: N->getIndividualCost());
1230 return Result;
1231 }
1232
1233 static std::string getNodeAttributes(const SplitGraph::Node *N,
1234 const SplitGraph &SG) {
1235 return N->hasAnyIncomingEdges() ? "" : "color=\"red\"";
1236 }
1237
1238 static std::string getEdgeAttributes(const SplitGraph::Node *N,
1239 SplitGraphEdgeDstIterator EI,
1240 const SplitGraph &SG) {
1241
1242 switch ((*EI.getCurrent())->Kind) {
1243 case SplitGraph::EdgeKind::DirectCall:
1244 return "";
1245 case SplitGraph::EdgeKind::IndirectCall:
1246 return "style=\"dashed\"";
1247 }
1248 llvm_unreachable("Unknown SplitGraph::EdgeKind enum");
1249 }
1250};
1251
1252//===----------------------------------------------------------------------===//
1253// Driver
1254//===----------------------------------------------------------------------===//
1255
1256namespace {
1257
1258// If we didn't externalize GVs, then local GVs need to be conservatively
1259// imported into every module (including their initializers), and then cleaned
1260// up afterwards.
1261static bool needsConservativeImport(const GlobalValue *GV) {
1262 if (const auto *Var = dyn_cast<GlobalVariable>(Val: GV))
1263 return Var->hasLocalLinkage();
1264 if (const auto *GA = dyn_cast<GlobalAlias>(Val: GV))
1265 return GA->hasLocalLinkage();
1266 return false;
1267}
1268
1269/// Prints a summary of the partition \p N, represented by module \p M, to \p
1270/// OS.
1271static void printPartitionSummary(raw_ostream &OS, unsigned N, const Module &M,
1272 unsigned PartCost, unsigned ModuleCost) {
1273 OS << "*** Partition P" << N << " ***\n";
1274
1275 for (const auto &Fn : M) {
1276 if (!Fn.isDeclaration())
1277 OS << " - [function] " << Fn.getName() << "\n";
1278 }
1279
1280 for (const auto &GV : M.globals()) {
1281 if (GV.hasInitializer())
1282 OS << " - [global] " << GV.getName() << "\n";
1283 }
1284
1285 OS << "Partition contains " << formatRatioOf(Num: PartCost, Dem: ModuleCost)
1286 << "% of the source\n";
1287}
1288
1289static void evaluateProposal(SplitProposal &Best, SplitProposal New) {
1290 SplitModuleTimer SMT("proposal_evaluation", "proposal ranking algorithm");
1291
1292 LLVM_DEBUG({
1293 New.verifyCompleteness();
1294 if (DebugProposalSearch)
1295 New.print(dbgs());
1296 });
1297
1298 const double CurBScore = Best.getBottleneckScore();
1299 const double CurCSScore = Best.getCodeSizeScore();
1300 const double NewBScore = New.getBottleneckScore();
1301 const double NewCSScore = New.getCodeSizeScore();
1302
1303 // TODO: Improve this
1304 // We can probably lower the precision of the comparison at first
1305 // e.g. if we have
1306 // - (Current): BScore: 0.489 CSCore 1.105
1307 // - (New): BScore: 0.475 CSCore 1.305
1308 // Currently we'd choose the new one because the bottleneck score is
1309 // lower, but the new one duplicates more code. It may be worth it to
1310 // discard the new proposal as the impact on build time is negligible.
1311
1312 // Compare them
1313 bool IsBest = false;
1314 if (NewBScore < CurBScore)
1315 IsBest = true;
1316 else if (NewBScore == CurBScore)
1317 IsBest = (NewCSScore < CurCSScore); // Use code size as tie breaker.
1318
1319 if (IsBest)
1320 Best = std::move(New);
1321
1322 LLVM_DEBUG(if (DebugProposalSearch) {
1323 if (IsBest)
1324 dbgs() << "[search] new best proposal!\n";
1325 else
1326 dbgs() << "[search] discarding - not profitable\n";
1327 });
1328}
1329
1330/// Trivial helper to create an identical copy of \p M.
1331static std::unique_ptr<Module> cloneAll(const Module &M) {
1332 ValueToValueMapTy VMap;
1333 return CloneModule(M, VMap, ShouldCloneDefinition: [&](const GlobalValue *GV) { return true; });
1334}
1335
1336/// Writes \p SG as a DOTGraph to \ref ModuleDotCfgDir if requested.
1337static void writeDOTGraph(const SplitGraph &SG) {
1338 if (ModuleDotCfgOutput.empty())
1339 return;
1340
1341 std::error_code EC;
1342 raw_fd_ostream OS(ModuleDotCfgOutput, EC);
1343 if (EC) {
1344 errs() << "[" DEBUG_TYPE "]: cannot open '" << ModuleDotCfgOutput
1345 << "' - DOTGraph will not be printed\n";
1346 }
1347 WriteGraph(O&: OS, G: SG, /*ShortName=*/ShortNames: false,
1348 /*Title=*/SG.getModule().getName());
1349}
1350
1351static void splitAMDGPUModule(
1352 GetTTIFn GetTTI, Module &M, unsigned NumParts,
1353 function_ref<void(std::unique_ptr<Module> MPart)> ModuleCallback) {
1354 CallGraph CG(M);
1355
1356 // Externalize functions whose address are taken.
1357 //
1358 // This is needed because partitioning is purely based on calls, but sometimes
1359 // a kernel/function may just look at the address of another local function
1360 // and not do anything (no calls). After partitioning, that local function may
1361 // end up in a different module (so it's just a declaration in the module
1362 // where its address is taken), which emits a "undefined hidden symbol" linker
1363 // error.
1364 //
1365 // Additionally, it guides partitioning to not duplicate this function if it's
1366 // called directly at some point.
1367 //
1368 // TODO: Could we be smarter about this ? This makes all functions whose
1369 // addresses are taken non-copyable. We should probably model this type of
1370 // constraint in the graph and use it to guide splitting, instead of
1371 // externalizing like this. Maybe non-copyable should really mean "keep one
1372 // visible copy, then internalize all other copies" for some functions?
1373 if (!NoExternalizeOnAddrTaken) {
1374 for (auto &Fn : M) {
1375 if (Fn.hasLocalLinkage() && Fn.hasAddressTaken()) {
1376 LLVM_DEBUG(dbgs() << "[externalize] "; Fn.printAsOperand(dbgs());
1377 dbgs() << " because its address is taken\n");
1378 Fn.externalize();
1379 }
1380 }
1381 }
1382
1383 // Externalize local GVs, which avoids duplicating their initializers, which
1384 // in turns helps keep code size in check.
1385 if (!NoExternalizeGlobals) {
1386 for (auto &GV : M.globals()) {
1387 if (GV.hasLocalLinkage())
1388 LLVM_DEBUG(dbgs() << "[externalize] GV " << GV.getName() << '\n');
1389 GV.externalize();
1390 }
1391 }
1392
1393 for (auto &GA : M.aliases()) {
1394 if (GA.hasLocalLinkage()) {
1395 LLVM_DEBUG(dbgs() << "[externalize] alias " << GA.getName() << '\n');
1396 GA.externalize();
1397 }
1398 }
1399
1400 // Start by calculating the cost of every function in the module, as well as
1401 // the module's overall cost.
1402 FunctionsCostMap FnCosts;
1403 const CostType ModuleCost = calculateFunctionCosts(GetTTI, M, CostMap&: FnCosts);
1404
1405 // Build the SplitGraph, which represents the module's functions and models
1406 // their dependencies accurately.
1407 SplitGraph SG(M, FnCosts, ModuleCost);
1408 SG.buildGraph(CG);
1409
1410 if (SG.empty()) {
1411 LLVM_DEBUG(
1412 dbgs()
1413 << "[!] no nodes in graph, input is empty - no splitting possible\n");
1414 ModuleCallback(cloneAll(M));
1415 return;
1416 }
1417
1418 LLVM_DEBUG({
1419 dbgs() << "[graph] nodes:\n";
1420 for (const SplitGraph::Node *N : SG.nodes()) {
1421 dbgs() << " - [" << N->getID() << "]: " << N->getName() << " "
1422 << (N->isGraphEntryPoint() ? "(entry)" : "") << " "
1423 << (N->isNonCopyable() ? "(noncopyable)" : "") << "\n";
1424 }
1425 });
1426
1427 writeDOTGraph(SG);
1428
1429 LLVM_DEBUG(dbgs() << "[search] testing splitting strategies\n");
1430
1431 std::optional<SplitProposal> Proposal;
1432 const auto EvaluateProposal = [&](SplitProposal SP) {
1433 SP.calculateScores();
1434 if (!Proposal)
1435 Proposal = std::move(SP);
1436 else
1437 evaluateProposal(Best&: *Proposal, New: std::move(SP));
1438 };
1439
1440 // TODO: It would be very easy to create new strategies by just adding a base
1441 // class to RecursiveSearchSplitting and abstracting it away.
1442 RecursiveSearchSplitting(SG, NumParts, EvaluateProposal).run();
1443 LLVM_DEBUG(if (Proposal) dbgs() << "[search done] selected proposal: "
1444 << Proposal->getName() << "\n";);
1445
1446 if (!Proposal) {
1447 LLVM_DEBUG(dbgs() << "[!] no proposal made, no splitting possible!\n");
1448 ModuleCallback(cloneAll(M));
1449 return;
1450 }
1451
1452 LLVM_DEBUG(Proposal->print(dbgs()););
1453
1454 std::optional<raw_fd_ostream> SummariesOS;
1455 if (!PartitionSummariesOutput.empty()) {
1456 std::error_code EC;
1457 SummariesOS.emplace(args&: PartitionSummariesOutput, args&: EC);
1458 if (EC)
1459 errs() << "[" DEBUG_TYPE "]: cannot open '" << PartitionSummariesOutput
1460 << "' - Partition summaries will not be printed\n";
1461 }
1462
1463 // One module will import all GlobalValues that are not Functions
1464 // and are not subject to conservative import.
1465 bool ImportAllGVs = true;
1466
1467 for (unsigned PID = 0; PID < NumParts; ++PID) {
1468 SplitModuleTimer SMT2("modules_creation",
1469 "creating modules for each partition");
1470 LLVM_DEBUG(dbgs() << "[split] creating new modules\n");
1471
1472 DenseSet<const Function *> FnsInPart;
1473 for (unsigned NodeID : (*Proposal)[PID].set_bits())
1474 FnsInPart.insert(V: &SG.getNode(ID: NodeID).getFunction());
1475
1476 // Don't create empty modules.
1477 if (FnsInPart.empty()) {
1478 LLVM_DEBUG(dbgs() << "[split] P" << PID
1479 << " is empty, not creating module\n");
1480 continue;
1481 }
1482
1483 ValueToValueMapTy VMap;
1484 CostType PartCost = 0;
1485 std::unique_ptr<Module> MPart(
1486 CloneModule(M, VMap, ShouldCloneDefinition: [&](const GlobalValue *GV) {
1487 // Functions go in their assigned partition.
1488 if (const auto *Fn = dyn_cast<Function>(Val: GV)) {
1489 if (FnsInPart.contains(V: Fn)) {
1490 PartCost += SG.getCost(F: *Fn);
1491 return true;
1492 }
1493 return false;
1494 }
1495
1496 // Aliases should not be separated from their underlying object.
1497 if (const auto *GA = dyn_cast<GlobalAlias>(Val: GV)) {
1498 if (const auto *Fn = dyn_cast<Function>(Val: GA->getAliaseeObject()))
1499 return FnsInPart.contains(V: Fn);
1500 }
1501
1502 // Everything else goes in the first non-empty module we create.
1503 return ImportAllGVs || needsConservativeImport(GV);
1504 }));
1505
1506 ImportAllGVs = false;
1507
1508 // Clean-up conservatively imported GVs without any users.
1509 for (auto &GV : make_early_inc_range(Range: MPart->global_values())) {
1510 if (needsConservativeImport(GV: &GV) && GV.use_empty())
1511 GV.eraseFromParent();
1512 }
1513
1514 if (SummariesOS)
1515 printPartitionSummary(OS&: *SummariesOS, N: PID, M: *MPart, PartCost, ModuleCost);
1516
1517 LLVM_DEBUG(
1518 printPartitionSummary(dbgs(), PID, *MPart, PartCost, ModuleCost));
1519
1520 ModuleCallback(std::move(MPart));
1521 }
1522}
1523} // namespace
1524
1525PreservedAnalyses AMDGPUSplitModulePass::run(Module &M,
1526 ModuleAnalysisManager &MAM) {
1527 SplitModuleTimer SMT(
1528 "total", "total pass runtime (incl. potentially waiting for lockfile)");
1529
1530 FunctionAnalysisManager &FAM =
1531 MAM.getResult<FunctionAnalysisManagerModuleProxy>(IR&: M).getManager();
1532 const auto TTIGetter = [&FAM](Function &F) -> const TargetTransformInfo & {
1533 return FAM.getResult<TargetIRAnalysis>(IR&: F);
1534 };
1535
1536 bool Done = false;
1537#ifndef NDEBUG
1538 if (UseLockFile) {
1539 SmallString<128> LockFilePath;
1540 sys::path::system_temp_directory(/*ErasedOnReboot=*/true, LockFilePath);
1541 sys::path::append(LockFilePath, "amdgpu-split-module-debug");
1542 LLVM_DEBUG(dbgs() << DEBUG_TYPE " using lockfile '" << LockFilePath
1543 << "'\n");
1544
1545 while (true) {
1546 llvm::LockFileManager Lock(LockFilePath.str());
1547 bool Owned;
1548 if (Error Err = Lock.tryLock().moveInto(Owned)) {
1549 consumeError(std::move(Err));
1550 LLVM_DEBUG(
1551 dbgs() << "[amdgpu-split-module] unable to acquire lockfile, debug "
1552 "output may be mangled by other processes\n");
1553 } else if (!Owned) {
1554 switch (Lock.waitForUnlockFor(std::chrono::seconds(90))) {
1555 case WaitForUnlockResult::Success:
1556 break;
1557 case WaitForUnlockResult::OwnerDied:
1558 continue; // try again to get the lock.
1559 case WaitForUnlockResult::Timeout:
1560 LLVM_DEBUG(
1561 dbgs()
1562 << "[amdgpu-split-module] unable to acquire lockfile, debug "
1563 "output may be mangled by other processes\n");
1564 Lock.unsafeUnlock();
1565 break; // give up
1566 }
1567 }
1568
1569 splitAMDGPUModule(TTIGetter, M, N, ModuleCallback);
1570 Done = true;
1571 break;
1572 }
1573 }
1574#endif
1575
1576 if (!Done)
1577 splitAMDGPUModule(GetTTI: TTIGetter, M, NumParts: N, ModuleCallback);
1578
1579 // We can change linkage/visibilities in the input, consider that nothing is
1580 // preserved just to be safe. This pass runs last anyway.
1581 return PreservedAnalyses::none();
1582}
1583} // namespace llvm
1584