1 //===- RegAllocPBQP.h -------------------------------------------*- C++ -*-===//
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 // This file defines the PBQPBuilder interface, for classes which build PBQP
10 // instances to represent register allocation problems, and the RegAllocPBQP
11 // interface.
12 //
13 //===----------------------------------------------------------------------===//
14
15 #ifndef LLVM_CODEGEN_REGALLOCPBQP_H
16 #define LLVM_CODEGEN_REGALLOCPBQP_H
17
18 #include "llvm/ADT/DenseMap.h"
19 #include "llvm/ADT/Hashing.h"
20 #include "llvm/CodeGen/PBQP/CostAllocator.h"
21 #include "llvm/CodeGen/PBQP/Graph.h"
22 #include "llvm/CodeGen/PBQP/Math.h"
23 #include "llvm/CodeGen/PBQP/ReductionRules.h"
24 #include "llvm/CodeGen/PBQP/Solution.h"
25 #include "llvm/CodeGen/Register.h"
26 #include "llvm/MC/MCRegister.h"
27 #include "llvm/Support/ErrorHandling.h"
28 #include <algorithm>
29 #include <cassert>
30 #include <cstddef>
31 #include <limits>
32 #include <memory>
33 #include <set>
34 #include <vector>
35
36 namespace llvm {
37
38 class FunctionPass;
39 class LiveIntervals;
40 class MachineBlockFrequencyInfo;
41 class MachineFunction;
42 class raw_ostream;
43
44 namespace PBQP {
45 namespace RegAlloc {
46
47 /// Spill option index.
getSpillOptionIdx()48 inline unsigned getSpillOptionIdx() { return 0; }
49
50 /// Metadata to speed allocatability test.
51 ///
52 /// Keeps track of the number of infinities in each row and column.
53 class MatrixMetadata {
54 public:
MatrixMetadata(const Matrix & M)55 MatrixMetadata(const Matrix& M)
56 : UnsafeRows(new bool[M.getRows() - 1]()),
57 UnsafeCols(new bool[M.getCols() - 1]()) {
58 unsigned* ColCounts = new unsigned[M.getCols() - 1]();
59
60 for (unsigned i = 1; i < M.getRows(); ++i) {
61 unsigned RowCount = 0;
62 for (unsigned j = 1; j < M.getCols(); ++j) {
63 if (M[i][j] == std::numeric_limits<PBQPNum>::infinity()) {
64 ++RowCount;
65 ++ColCounts[j - 1];
66 UnsafeRows[i - 1] = true;
67 UnsafeCols[j - 1] = true;
68 }
69 }
70 WorstRow = std::max(WorstRow, RowCount);
71 }
72 unsigned WorstColCountForCurRow =
73 *std::max_element(ColCounts, ColCounts + M.getCols() - 1);
74 WorstCol = std::max(WorstCol, WorstColCountForCurRow);
75 delete[] ColCounts;
76 }
77
78 MatrixMetadata(const MatrixMetadata &) = delete;
79 MatrixMetadata &operator=(const MatrixMetadata &) = delete;
80
getWorstRow()81 unsigned getWorstRow() const { return WorstRow; }
getWorstCol()82 unsigned getWorstCol() const { return WorstCol; }
getUnsafeRows()83 const bool* getUnsafeRows() const { return UnsafeRows.get(); }
getUnsafeCols()84 const bool* getUnsafeCols() const { return UnsafeCols.get(); }
85
86 private:
87 unsigned WorstRow = 0;
88 unsigned WorstCol = 0;
89 std::unique_ptr<bool[]> UnsafeRows;
90 std::unique_ptr<bool[]> UnsafeCols;
91 };
92
93 /// Holds a vector of the allowed physical regs for a vreg.
94 class AllowedRegVector {
95 friend hash_code hash_value(const AllowedRegVector &);
96
97 public:
98 AllowedRegVector() = default;
99 AllowedRegVector(AllowedRegVector &&) = default;
100
AllowedRegVector(const std::vector<MCRegister> & OptVec)101 AllowedRegVector(const std::vector<MCRegister> &OptVec)
102 : NumOpts(OptVec.size()), Opts(new MCRegister[NumOpts]) {
103 std::copy(OptVec.begin(), OptVec.end(), Opts.get());
104 }
105
size()106 unsigned size() const { return NumOpts; }
107 MCRegister operator[](size_t I) const { return Opts[I]; }
108
109 bool operator==(const AllowedRegVector &Other) const {
110 if (NumOpts != Other.NumOpts)
111 return false;
112 return std::equal(Opts.get(), Opts.get() + NumOpts, Other.Opts.get());
113 }
114
115 bool operator!=(const AllowedRegVector &Other) const {
116 return !(*this == Other);
117 }
118
119 private:
120 unsigned NumOpts = 0;
121 std::unique_ptr<MCRegister[]> Opts;
122 };
123
hash_value(const AllowedRegVector & OptRegs)124 inline hash_code hash_value(const AllowedRegVector &OptRegs) {
125 MCRegister *OStart = OptRegs.Opts.get();
126 MCRegister *OEnd = OptRegs.Opts.get() + OptRegs.NumOpts;
127 return hash_combine(OptRegs.NumOpts,
128 hash_combine_range(OStart, OEnd));
129 }
130
131 /// Holds graph-level metadata relevant to PBQP RA problems.
132 class GraphMetadata {
133 private:
134 using AllowedRegVecPool = ValuePool<AllowedRegVector>;
135
136 public:
137 using AllowedRegVecRef = AllowedRegVecPool::PoolRef;
138
GraphMetadata(MachineFunction & MF,LiveIntervals & LIS,MachineBlockFrequencyInfo & MBFI)139 GraphMetadata(MachineFunction &MF,
140 LiveIntervals &LIS,
141 MachineBlockFrequencyInfo &MBFI)
142 : MF(MF), LIS(LIS), MBFI(MBFI) {}
143
144 MachineFunction &MF;
145 LiveIntervals &LIS;
146 MachineBlockFrequencyInfo &MBFI;
147
setNodeIdForVReg(Register VReg,GraphBase::NodeId NId)148 void setNodeIdForVReg(Register VReg, GraphBase::NodeId NId) {
149 VRegToNodeId[VReg.id()] = NId;
150 }
151
getNodeIdForVReg(Register VReg)152 GraphBase::NodeId getNodeIdForVReg(Register VReg) const {
153 auto VRegItr = VRegToNodeId.find(VReg);
154 if (VRegItr == VRegToNodeId.end())
155 return GraphBase::invalidNodeId();
156 return VRegItr->second;
157 }
158
getAllowedRegs(AllowedRegVector Allowed)159 AllowedRegVecRef getAllowedRegs(AllowedRegVector Allowed) {
160 return AllowedRegVecs.getValue(std::move(Allowed));
161 }
162
163 private:
164 DenseMap<Register, GraphBase::NodeId> VRegToNodeId;
165 AllowedRegVecPool AllowedRegVecs;
166 };
167
168 /// Holds solver state and other metadata relevant to each PBQP RA node.
169 class NodeMetadata {
170 public:
171 using AllowedRegVector = RegAlloc::AllowedRegVector;
172
173 // The node's reduction state. The order in this enum is important,
174 // as it is assumed nodes can only progress up (i.e. towards being
175 // optimally reducible) when reducing the graph.
176 using ReductionState = enum {
177 Unprocessed,
178 NotProvablyAllocatable,
179 ConservativelyAllocatable,
180 OptimallyReducible
181 };
182
183 NodeMetadata() = default;
184
NodeMetadata(const NodeMetadata & Other)185 NodeMetadata(const NodeMetadata &Other)
186 : RS(Other.RS), NumOpts(Other.NumOpts), DeniedOpts(Other.DeniedOpts),
187 OptUnsafeEdges(new unsigned[NumOpts]), VReg(Other.VReg),
188 AllowedRegs(Other.AllowedRegs)
189 #ifndef NDEBUG
190 , everConservativelyAllocatable(Other.everConservativelyAllocatable)
191 #endif
192 {
193 if (NumOpts > 0) {
194 std::copy(&Other.OptUnsafeEdges[0], &Other.OptUnsafeEdges[NumOpts],
195 &OptUnsafeEdges[0]);
196 }
197 }
198
199 NodeMetadata(NodeMetadata &&) = default;
200 NodeMetadata& operator=(NodeMetadata &&) = default;
201
setVReg(Register VReg)202 void setVReg(Register VReg) { this->VReg = VReg; }
getVReg()203 Register getVReg() const { return VReg; }
204
setAllowedRegs(GraphMetadata::AllowedRegVecRef AllowedRegs)205 void setAllowedRegs(GraphMetadata::AllowedRegVecRef AllowedRegs) {
206 this->AllowedRegs = std::move(AllowedRegs);
207 }
getAllowedRegs()208 const AllowedRegVector& getAllowedRegs() const { return *AllowedRegs; }
209
setup(const Vector & Costs)210 void setup(const Vector& Costs) {
211 NumOpts = Costs.getLength() - 1;
212 OptUnsafeEdges = std::unique_ptr<unsigned[]>(new unsigned[NumOpts]());
213 }
214
getReductionState()215 ReductionState getReductionState() const { return RS; }
setReductionState(ReductionState RS)216 void setReductionState(ReductionState RS) {
217 assert(RS >= this->RS && "A node's reduction state can not be downgraded");
218 this->RS = RS;
219
220 #ifndef NDEBUG
221 // Remember this state to assert later that a non-infinite register
222 // option was available.
223 if (RS == ConservativelyAllocatable)
224 everConservativelyAllocatable = true;
225 #endif
226 }
227
handleAddEdge(const MatrixMetadata & MD,bool Transpose)228 void handleAddEdge(const MatrixMetadata& MD, bool Transpose) {
229 DeniedOpts += Transpose ? MD.getWorstRow() : MD.getWorstCol();
230 const bool* UnsafeOpts =
231 Transpose ? MD.getUnsafeCols() : MD.getUnsafeRows();
232 for (unsigned i = 0; i < NumOpts; ++i)
233 OptUnsafeEdges[i] += UnsafeOpts[i];
234 }
235
handleRemoveEdge(const MatrixMetadata & MD,bool Transpose)236 void handleRemoveEdge(const MatrixMetadata& MD, bool Transpose) {
237 DeniedOpts -= Transpose ? MD.getWorstRow() : MD.getWorstCol();
238 const bool* UnsafeOpts =
239 Transpose ? MD.getUnsafeCols() : MD.getUnsafeRows();
240 for (unsigned i = 0; i < NumOpts; ++i)
241 OptUnsafeEdges[i] -= UnsafeOpts[i];
242 }
243
isConservativelyAllocatable()244 bool isConservativelyAllocatable() const {
245 return (DeniedOpts < NumOpts) ||
246 (std::find(&OptUnsafeEdges[0], &OptUnsafeEdges[NumOpts], 0) !=
247 &OptUnsafeEdges[NumOpts]);
248 }
249
250 #ifndef NDEBUG
wasConservativelyAllocatable()251 bool wasConservativelyAllocatable() const {
252 return everConservativelyAllocatable;
253 }
254 #endif
255
256 private:
257 ReductionState RS = Unprocessed;
258 unsigned NumOpts = 0;
259 unsigned DeniedOpts = 0;
260 std::unique_ptr<unsigned[]> OptUnsafeEdges;
261 Register VReg;
262 GraphMetadata::AllowedRegVecRef AllowedRegs;
263
264 #ifndef NDEBUG
265 bool everConservativelyAllocatable = false;
266 #endif
267 };
268
269 class RegAllocSolverImpl {
270 private:
271 using RAMatrix = MDMatrix<MatrixMetadata>;
272
273 public:
274 using RawVector = PBQP::Vector;
275 using RawMatrix = PBQP::Matrix;
276 using Vector = PBQP::Vector;
277 using Matrix = RAMatrix;
278 using CostAllocator = PBQP::PoolCostAllocator<Vector, Matrix>;
279
280 using NodeId = GraphBase::NodeId;
281 using EdgeId = GraphBase::EdgeId;
282
283 using NodeMetadata = RegAlloc::NodeMetadata;
284 struct EdgeMetadata {};
285 using GraphMetadata = RegAlloc::GraphMetadata;
286
287 using Graph = PBQP::Graph<RegAllocSolverImpl>;
288
RegAllocSolverImpl(Graph & G)289 RegAllocSolverImpl(Graph &G) : G(G) {}
290
solve()291 Solution solve() {
292 G.setSolver(*this);
293 Solution S;
294 setup();
295 S = backpropagate(G, reduce());
296 G.unsetSolver();
297 return S;
298 }
299
handleAddNode(NodeId NId)300 void handleAddNode(NodeId NId) {
301 assert(G.getNodeCosts(NId).getLength() > 1 &&
302 "PBQP Graph should not contain single or zero-option nodes");
303 G.getNodeMetadata(NId).setup(G.getNodeCosts(NId));
304 }
305
handleRemoveNode(NodeId NId)306 void handleRemoveNode(NodeId NId) {}
handleSetNodeCosts(NodeId NId,const Vector & newCosts)307 void handleSetNodeCosts(NodeId NId, const Vector& newCosts) {}
308
handleAddEdge(EdgeId EId)309 void handleAddEdge(EdgeId EId) {
310 handleReconnectEdge(EId, G.getEdgeNode1Id(EId));
311 handleReconnectEdge(EId, G.getEdgeNode2Id(EId));
312 }
313
handleDisconnectEdge(EdgeId EId,NodeId NId)314 void handleDisconnectEdge(EdgeId EId, NodeId NId) {
315 NodeMetadata& NMd = G.getNodeMetadata(NId);
316 const MatrixMetadata& MMd = G.getEdgeCosts(EId).getMetadata();
317 NMd.handleRemoveEdge(MMd, NId == G.getEdgeNode2Id(EId));
318 promote(NId, NMd);
319 }
320
handleReconnectEdge(EdgeId EId,NodeId NId)321 void handleReconnectEdge(EdgeId EId, NodeId NId) {
322 NodeMetadata& NMd = G.getNodeMetadata(NId);
323 const MatrixMetadata& MMd = G.getEdgeCosts(EId).getMetadata();
324 NMd.handleAddEdge(MMd, NId == G.getEdgeNode2Id(EId));
325 }
326
handleUpdateCosts(EdgeId EId,const Matrix & NewCosts)327 void handleUpdateCosts(EdgeId EId, const Matrix& NewCosts) {
328 NodeId N1Id = G.getEdgeNode1Id(EId);
329 NodeId N2Id = G.getEdgeNode2Id(EId);
330 NodeMetadata& N1Md = G.getNodeMetadata(N1Id);
331 NodeMetadata& N2Md = G.getNodeMetadata(N2Id);
332 bool Transpose = N1Id != G.getEdgeNode1Id(EId);
333
334 // Metadata are computed incrementally. First, update them
335 // by removing the old cost.
336 const MatrixMetadata& OldMMd = G.getEdgeCosts(EId).getMetadata();
337 N1Md.handleRemoveEdge(OldMMd, Transpose);
338 N2Md.handleRemoveEdge(OldMMd, !Transpose);
339
340 // And update now the metadata with the new cost.
341 const MatrixMetadata& MMd = NewCosts.getMetadata();
342 N1Md.handleAddEdge(MMd, Transpose);
343 N2Md.handleAddEdge(MMd, !Transpose);
344
345 // As the metadata may have changed with the update, the nodes may have
346 // become ConservativelyAllocatable or OptimallyReducible.
347 promote(N1Id, N1Md);
348 promote(N2Id, N2Md);
349 }
350
351 private:
promote(NodeId NId,NodeMetadata & NMd)352 void promote(NodeId NId, NodeMetadata& NMd) {
353 if (G.getNodeDegree(NId) == 3) {
354 // This node is becoming optimally reducible.
355 moveToOptimallyReducibleNodes(NId);
356 } else if (NMd.getReductionState() ==
357 NodeMetadata::NotProvablyAllocatable &&
358 NMd.isConservativelyAllocatable()) {
359 // This node just became conservatively allocatable.
360 moveToConservativelyAllocatableNodes(NId);
361 }
362 }
363
removeFromCurrentSet(NodeId NId)364 void removeFromCurrentSet(NodeId NId) {
365 switch (G.getNodeMetadata(NId).getReductionState()) {
366 case NodeMetadata::Unprocessed: break;
367 case NodeMetadata::OptimallyReducible:
368 assert(OptimallyReducibleNodes.find(NId) !=
369 OptimallyReducibleNodes.end() &&
370 "Node not in optimally reducible set.");
371 OptimallyReducibleNodes.erase(NId);
372 break;
373 case NodeMetadata::ConservativelyAllocatable:
374 assert(ConservativelyAllocatableNodes.find(NId) !=
375 ConservativelyAllocatableNodes.end() &&
376 "Node not in conservatively allocatable set.");
377 ConservativelyAllocatableNodes.erase(NId);
378 break;
379 case NodeMetadata::NotProvablyAllocatable:
380 assert(NotProvablyAllocatableNodes.find(NId) !=
381 NotProvablyAllocatableNodes.end() &&
382 "Node not in not-provably-allocatable set.");
383 NotProvablyAllocatableNodes.erase(NId);
384 break;
385 }
386 }
387
moveToOptimallyReducibleNodes(NodeId NId)388 void moveToOptimallyReducibleNodes(NodeId NId) {
389 removeFromCurrentSet(NId);
390 OptimallyReducibleNodes.insert(NId);
391 G.getNodeMetadata(NId).setReductionState(
392 NodeMetadata::OptimallyReducible);
393 }
394
moveToConservativelyAllocatableNodes(NodeId NId)395 void moveToConservativelyAllocatableNodes(NodeId NId) {
396 removeFromCurrentSet(NId);
397 ConservativelyAllocatableNodes.insert(NId);
398 G.getNodeMetadata(NId).setReductionState(
399 NodeMetadata::ConservativelyAllocatable);
400 }
401
moveToNotProvablyAllocatableNodes(NodeId NId)402 void moveToNotProvablyAllocatableNodes(NodeId NId) {
403 removeFromCurrentSet(NId);
404 NotProvablyAllocatableNodes.insert(NId);
405 G.getNodeMetadata(NId).setReductionState(
406 NodeMetadata::NotProvablyAllocatable);
407 }
408
setup()409 void setup() {
410 // Set up worklists.
411 for (auto NId : G.nodeIds()) {
412 if (G.getNodeDegree(NId) < 3)
413 moveToOptimallyReducibleNodes(NId);
414 else if (G.getNodeMetadata(NId).isConservativelyAllocatable())
415 moveToConservativelyAllocatableNodes(NId);
416 else
417 moveToNotProvablyAllocatableNodes(NId);
418 }
419 }
420
421 // Compute a reduction order for the graph by iteratively applying PBQP
422 // reduction rules. Locally optimal rules are applied whenever possible (R0,
423 // R1, R2). If no locally-optimal rules apply then any conservatively
424 // allocatable node is reduced. Finally, if no conservatively allocatable
425 // node exists then the node with the lowest spill-cost:degree ratio is
426 // selected.
reduce()427 std::vector<GraphBase::NodeId> reduce() {
428 assert(!G.empty() && "Cannot reduce empty graph.");
429
430 using NodeId = GraphBase::NodeId;
431 std::vector<NodeId> NodeStack;
432
433 // Consume worklists.
434 while (true) {
435 if (!OptimallyReducibleNodes.empty()) {
436 NodeSet::iterator NItr = OptimallyReducibleNodes.begin();
437 NodeId NId = *NItr;
438 OptimallyReducibleNodes.erase(NItr);
439 NodeStack.push_back(NId);
440 switch (G.getNodeDegree(NId)) {
441 case 0:
442 break;
443 case 1:
444 applyR1(G, NId);
445 break;
446 case 2:
447 applyR2(G, NId);
448 break;
449 default: llvm_unreachable("Not an optimally reducible node.");
450 }
451 } else if (!ConservativelyAllocatableNodes.empty()) {
452 // Conservatively allocatable nodes will never spill. For now just
453 // take the first node in the set and push it on the stack. When we
454 // start optimizing more heavily for register preferencing, it may
455 // would be better to push nodes with lower 'expected' or worst-case
456 // register costs first (since early nodes are the most
457 // constrained).
458 NodeSet::iterator NItr = ConservativelyAllocatableNodes.begin();
459 NodeId NId = *NItr;
460 ConservativelyAllocatableNodes.erase(NItr);
461 NodeStack.push_back(NId);
462 G.disconnectAllNeighborsFromNode(NId);
463 } else if (!NotProvablyAllocatableNodes.empty()) {
464 NodeSet::iterator NItr =
465 std::min_element(NotProvablyAllocatableNodes.begin(),
466 NotProvablyAllocatableNodes.end(),
467 SpillCostComparator(G));
468 NodeId NId = *NItr;
469 NotProvablyAllocatableNodes.erase(NItr);
470 NodeStack.push_back(NId);
471 G.disconnectAllNeighborsFromNode(NId);
472 } else
473 break;
474 }
475
476 return NodeStack;
477 }
478
479 class SpillCostComparator {
480 public:
SpillCostComparator(const Graph & G)481 SpillCostComparator(const Graph& G) : G(G) {}
482
operator()483 bool operator()(NodeId N1Id, NodeId N2Id) {
484 PBQPNum N1SC = G.getNodeCosts(N1Id)[0];
485 PBQPNum N2SC = G.getNodeCosts(N2Id)[0];
486 if (N1SC == N2SC)
487 return G.getNodeDegree(N1Id) < G.getNodeDegree(N2Id);
488 return N1SC < N2SC;
489 }
490
491 private:
492 const Graph& G;
493 };
494
495 Graph& G;
496 using NodeSet = std::set<NodeId>;
497 NodeSet OptimallyReducibleNodes;
498 NodeSet ConservativelyAllocatableNodes;
499 NodeSet NotProvablyAllocatableNodes;
500 };
501
502 class PBQPRAGraph : public PBQP::Graph<RegAllocSolverImpl> {
503 private:
504 using BaseT = PBQP::Graph<RegAllocSolverImpl>;
505
506 public:
PBQPRAGraph(GraphMetadata Metadata)507 PBQPRAGraph(GraphMetadata Metadata) : BaseT(std::move(Metadata)) {}
508
509 /// Dump this graph to dbgs().
510 void dump() const;
511
512 /// Dump this graph to an output stream.
513 /// @param OS Output stream to print on.
514 void dump(raw_ostream &OS) const;
515
516 /// Print a representation of this graph in DOT format.
517 /// @param OS Output stream to print on.
518 void printDot(raw_ostream &OS) const;
519 };
520
solve(PBQPRAGraph & G)521 inline Solution solve(PBQPRAGraph& G) {
522 if (G.empty())
523 return Solution();
524 RegAllocSolverImpl RegAllocSolver(G);
525 return RegAllocSolver.solve();
526 }
527
528 } // end namespace RegAlloc
529 } // end namespace PBQP
530
531 /// Create a PBQP register allocator instance.
532 FunctionPass *
533 createPBQPRegisterAllocator(char *customPassID = nullptr);
534
535 } // end namespace llvm
536
537 #endif // LLVM_CODEGEN_REGALLOCPBQP_H
538