Added an additional PBQP problem builder which adds coalescing costs (both between...
[oota-llvm.git] / include / llvm / CodeGen / RegAllocPBQP.h
1 //===-- RegAllocPBQP.h ------------------------------------------*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file defines the PBQPBuilder interface, for classes which build PBQP
11 // instances to represent register allocation problems, and the RegAllocPBQP
12 // interface.
13 //
14 //===----------------------------------------------------------------------===//
15
16 #ifndef LLVM_CODEGEN_REGALLOCPBQP_H
17 #define LLVM_CODEGEN_REGALLOCPBQP_H
18
19 #include "llvm/ADT/DenseMap.h"
20 #include "llvm/CodeGen/MachineFunctionPass.h"
21 #include "llvm/CodeGen/PBQP/Graph.h"
22 #include "llvm/CodeGen/PBQP/Solution.h"
23
24 #include <map>
25
26 namespace llvm {
27
28   class LiveInterval;
29   class MachineFunction;
30   class MachineLoopInfo;
31
32   /// This class wraps up a PBQP instance representing a register allocation
33   /// problem, plus the structures necessary to map back from the PBQP solution
34   /// to a register allocation solution. (i.e. The PBQP-node <--> vreg map,
35   /// and the PBQP option <--> storage location map).
36
37   class PBQPRAProblem {
38   public:
39
40     typedef SmallVector<unsigned, 16> AllowedSet;
41
42     PBQP::Graph& getGraph() { return graph; }
43
44     const PBQP::Graph& getGraph() const { return graph; }
45
46     /// Record the mapping between the given virtual register and PBQP node,
47     /// and the set of allowed pregs for the vreg.
48     ///
49     /// If you are extending
50     /// PBQPBuilder you are unlikely to need this: Nodes and options for all
51     /// vregs will already have been set up for you by the base class. 
52     template <typename AllowedRegsItr>
53     void recordVReg(unsigned vreg, PBQP::Graph::NodeItr node,
54                     AllowedRegsItr arBegin, AllowedRegsItr arEnd) {
55       assert(node2VReg.find(node) == node2VReg.end() && "Re-mapping node.");
56       assert(vreg2Node.find(vreg) == vreg2Node.end() && "Re-mapping vreg.");
57       assert(allowedSets[vreg].empty() && "vreg already has pregs.");
58
59       node2VReg[node] = vreg;
60       vreg2Node[vreg] = node;
61       std::copy(arBegin, arEnd, std::back_inserter(allowedSets[vreg]));
62     }
63
64     /// Get the virtual register corresponding to the given PBQP node.
65     unsigned getVRegForNode(PBQP::Graph::ConstNodeItr node) const;
66
67     /// Get the PBQP node corresponding to the given virtual register.
68     PBQP::Graph::NodeItr getNodeForVReg(unsigned vreg) const;
69
70     /// Returns true if the given PBQP option represents a physical register,
71     /// false otherwise.
72     bool isPRegOption(unsigned vreg, unsigned option) const {
73       // At present we only have spills or pregs, so anything that's not a
74       // spill is a preg. (This might be extended one day to support remat).
75       return !isSpillOption(vreg, option);
76     }
77
78     /// Returns true if the given PBQP option represents spilling, false
79     /// otherwise.
80     bool isSpillOption(unsigned vreg, unsigned option) const {
81       // We hardcode option zero as the spill option.
82       return option == 0;
83     }
84
85     /// Returns the allowed set for the given virtual register.
86     const AllowedSet& getAllowedSet(unsigned vreg) const;
87
88     /// Get PReg for option.
89     unsigned getPRegForOption(unsigned vreg, unsigned option) const;
90
91   private:
92
93     typedef std::map<PBQP::Graph::ConstNodeItr, unsigned,
94                      PBQP::NodeItrComparator>  Node2VReg;
95     typedef DenseMap<unsigned, PBQP::Graph::NodeItr> VReg2Node;
96     typedef std::map<unsigned, AllowedSet> AllowedSetMap;
97
98     PBQP::Graph graph;
99     Node2VReg node2VReg;
100     VReg2Node vreg2Node;
101
102     AllowedSetMap allowedSets;
103     
104   };
105
106   /// Builds PBQP instances to represent register allocation problems. Includes
107   /// spill, interference and coalescing costs by default. You can extend this
108   /// class to support additional constraints for your architecture.
109   class PBQPBuilder {
110   private:
111     PBQPBuilder(const PBQPBuilder&) {}
112     void operator=(const PBQPBuilder&) {}
113   public:
114
115     typedef std::set<unsigned> RegSet;
116  
117     /// Default constructor.
118     PBQPBuilder() {}
119
120     /// Clean up a PBQPBuilder.
121     virtual ~PBQPBuilder() {}
122
123     /// Build a PBQP instance to represent the register allocation problem for
124     /// the given MachineFunction.
125     virtual std::auto_ptr<PBQPRAProblem> build(
126                                               MachineFunction *mf,
127                                               const LiveIntervals *lis,
128                                               const MachineLoopInfo *loopInfo,
129                                               const RegSet &vregs);
130   private:
131
132     void addSpillCosts(PBQP::Vector &costVec, PBQP::PBQPNum spillCost);
133
134     void addInterferenceCosts(PBQP::Matrix &costMat,
135                               const PBQPRAProblem::AllowedSet &vr1Allowed,
136                               const PBQPRAProblem::AllowedSet &vr2Allowed,
137                               const TargetRegisterInfo *tri);
138   };
139
140   /// Extended builder which adds coalescing constraints to a problem.
141   class PBQPBuilderWithCoalescing : public PBQPBuilder {
142   public:
143  
144     /// Build a PBQP instance to represent the register allocation problem for
145     /// the given MachineFunction.
146     virtual std::auto_ptr<PBQPRAProblem> build(
147                                               MachineFunction *mf,
148                                               const LiveIntervals *lis,
149                                               const MachineLoopInfo *loopInfo,
150                                               const RegSet &vregs);   
151
152   private:
153
154     void addPhysRegCoalesce(PBQP::Vector &costVec, unsigned pregOption,
155                             PBQP::PBQPNum benefit);
156
157     void addVirtRegCoalesce(PBQP::Matrix &costMat,
158                             const PBQPRAProblem::AllowedSet &vr1Allowed,
159                             const PBQPRAProblem::AllowedSet &vr2Allowed,
160                             PBQP::PBQPNum benefit);
161   };
162
163   ///
164   /// PBQP based allocators solve the register allocation problem by mapping
165   /// register allocation problems to Partitioned Boolean Quadratic
166   /// Programming problems.
167   class RegAllocPBQP : public MachineFunctionPass {
168   public:
169
170     static char ID;
171
172     /// Construct a PBQP register allocator.
173     RegAllocPBQP(std::auto_ptr<PBQPBuilder> b) : MachineFunctionPass(ID), builder(b) {}
174
175     /// Return the pass name.
176     virtual const char* getPassName() const {
177       return "PBQP Register Allocator";
178     }
179
180     /// PBQP analysis usage.
181     virtual void getAnalysisUsage(AnalysisUsage &au) const;
182
183     /// Perform register allocation
184     virtual bool runOnMachineFunction(MachineFunction &MF);
185
186   private:
187
188     typedef std::map<const LiveInterval*, unsigned> LI2NodeMap;
189     typedef std::vector<const LiveInterval*> Node2LIMap;
190     typedef std::vector<unsigned> AllowedSet;
191     typedef std::vector<AllowedSet> AllowedSetMap;
192     typedef std::pair<unsigned, unsigned> RegPair;
193     typedef std::map<RegPair, PBQP::PBQPNum> CoalesceMap;
194     typedef std::vector<PBQP::Graph::NodeItr> NodeVector;
195     typedef std::set<unsigned> RegSet;
196
197
198     std::auto_ptr<PBQPBuilder> builder;
199
200     MachineFunction *mf;
201     const TargetMachine *tm;
202     const TargetRegisterInfo *tri;
203     const TargetInstrInfo *tii;
204     const MachineLoopInfo *loopInfo;
205     MachineRegisterInfo *mri;
206     RenderMachineFunction *rmf;
207
208     LiveIntervals *lis;
209     LiveStacks *lss;
210     VirtRegMap *vrm;
211
212     LI2NodeMap li2Node;
213     Node2LIMap node2LI;
214     AllowedSetMap allowedSets;
215     RegSet vregsToAlloc, emptyIntervalVRegs;
216     NodeVector problemNodes;
217
218
219     /// Builds a PBQP cost vector.
220     template <typename RegContainer>
221     PBQP::Vector buildCostVector(unsigned vReg,
222                                  const RegContainer &allowed,
223                                  const CoalesceMap &cealesces,
224                                  PBQP::PBQPNum spillCost) const;
225
226     /// \brief Builds a PBQP interference matrix.
227     ///
228     /// @return Either a pointer to a non-zero PBQP matrix representing the
229     ///         allocation option costs, or a null pointer for a zero matrix.
230     ///
231     /// Expects allowed sets for two interfering LiveIntervals. These allowed
232     /// sets should contain only allocable registers from the LiveInterval's
233     /// register class, with any interfering pre-colored registers removed.
234     template <typename RegContainer>
235     PBQP::Matrix* buildInterferenceMatrix(const RegContainer &allowed1,
236                                           const RegContainer &allowed2) const;
237
238     ///
239     /// Expects allowed sets for two potentially coalescable LiveIntervals,
240     /// and an estimated benefit due to coalescing. The allowed sets should
241     /// contain only allocable registers from the LiveInterval's register
242     /// classes, with any interfering pre-colored registers removed.
243     template <typename RegContainer>
244     PBQP::Matrix* buildCoalescingMatrix(const RegContainer &allowed1,
245                                         const RegContainer &allowed2,
246                                         PBQP::PBQPNum cBenefit) const;
247
248     /// \brief Finds coalescing opportunities and returns them as a map.
249     ///
250     /// Any entries in the map are guaranteed coalescable, even if their
251     /// corresponding live intervals overlap.
252     CoalesceMap findCoalesces();
253
254     /// \brief Finds the initial set of vreg intervals to allocate.
255     void findVRegIntervalsToAlloc();
256
257     /// \brief Constructs a PBQP problem representation of the register
258     /// allocation problem for this function.
259     ///
260     /// @return a PBQP solver object for the register allocation problem.
261     PBQP::Graph constructPBQPProblem();
262
263     /// \brief Adds a stack interval if the given live interval has been
264     /// spilled. Used to support stack slot coloring.
265     void addStackInterval(const LiveInterval *spilled,MachineRegisterInfo* mri);
266
267     /// \brief Given a solved PBQP problem maps this solution back to a register
268     /// assignment.
269     bool mapPBQPToRegAlloc(const PBQP::Solution &solution);
270
271     /// \brief Given a solved PBQP problem maps this solution back to a register
272     /// assignment.
273     bool mapPBQPToRegAlloc2(const PBQPRAProblem &problem,
274                             const PBQP::Solution &solution);
275
276     /// \brief Postprocessing before final spilling. Sets basic block "live in"
277     /// variables.
278     void finalizeAlloc() const;
279
280   };
281
282 }
283
284 #endif /* LLVM_CODEGEN_REGALLOCPBQP_H */