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