f6f3016acd5146b7de28b32f0b6eb75e04a844df
[oota-llvm.git] / lib / Transforms / Instrumentation / ProfilePaths / CombineBranch.cpp
1 //===-- CombineBranch.cpp ------------------------------------ ---*- C++ -*--=//
2 // Pass to instrument loops
3 //
4 // At every backedge, insert a counter for that backedge and a call function
5 //===----------------------------------------------------------------------===//
6
7 #include "llvm/Analysis/Dominators.h"
8 #include "llvm/Support/CFG.h"
9 #include "llvm/Constants.h"
10 #include "llvm/iMemory.h"
11 #include "llvm/GlobalVariable.h"
12 #include "llvm/DerivedTypes.h"
13 #include "llvm/iOther.h"
14 #include "llvm/iOperators.h"
15 #include "llvm/iTerminators.h"
16 #include "llvm/iPHINode.h"
17 #include "llvm/Module.h"
18 #include "llvm/Function.h"
19 #include "llvm/Pass.h"
20
21 //this is used to color vertices
22 //during DFS
23
24 enum Color{
25   WHITE,
26   GREY,
27   BLACK
28 };
29
30 namespace{
31   struct CombineBranches : public FunctionPass {
32   private:
33     //DominatorSet *DS;
34     void getBackEdgesVisit(BasicBlock *u,
35                            std::map<BasicBlock *, Color > &color,
36                            std::map<BasicBlock *, int > &d, 
37                            int &time,
38                            std::map<BasicBlock *, BasicBlock *> &be);
39     void removeRedundant(std::map<BasicBlock *, BasicBlock *> &be);
40     void getBackEdges(Function &F);
41   public:
42     bool runOnFunction(Function &F);
43   };
44   
45   RegisterOpt<CombineBranches> X("branch-combine", "Multiple backedges going to same target are merged");
46 }
47
48 //helper function to get back edges: it is called by 
49 //the "getBackEdges" function below
50 void CombineBranches::getBackEdgesVisit(BasicBlock *u,
51                        std::map<BasicBlock *, Color > &color,
52                        std::map<BasicBlock *, int > &d, 
53                        int &time,
54                        std::map<BasicBlock *, BasicBlock *> &be) {
55   
56   color[u]=GREY;
57   time++;
58   d[u]=time;
59
60   for(BasicBlock::succ_iterator vl = succ_begin(u), 
61         ve = succ_end(u); vl != ve; ++vl){
62     
63     BasicBlock *BB = *vl;
64
65     if(color[BB]!=GREY && color[BB]!=BLACK){
66       getBackEdgesVisit(BB, color, d, time, be);
67     }
68     
69     //now checking for d and f vals
70     else if(color[BB]==GREY){
71       //so v is ancestor of u if time of u > time of v
72       if(d[u] >= d[BB]){
73         //u->BB is a backedge
74         be[u] = BB;
75       }
76     }
77   }
78   color[u]=BLACK;//done with visiting the node and its neighbors
79 }
80
81 //look at all BEs, and remove all BEs that are dominated by other BE's in the
82 //set
83 void CombineBranches::removeRedundant(std::map<BasicBlock *, BasicBlock *> &be){
84   std::vector<BasicBlock *> toDelete;
85   std::map<BasicBlock *, int> seenBB;
86   
87   for(std::map<BasicBlock *, BasicBlock *>::iterator MI = be.begin(), 
88         ME = be.end(); MI != ME; ++MI){
89     
90     if(seenBB[MI->second])
91       continue;
92     
93     seenBB[MI->second] = 1;
94
95     std::vector<BasicBlock *> sameTarget;
96     sameTarget.clear();
97     
98     for(std::map<BasicBlock *, BasicBlock *>::iterator MMI = be.begin(), 
99           MME = be.end(); MMI != MME; ++MMI){
100       
101       if(MMI->first == MI->first)
102         continue;
103       
104       if(MMI->second == MI->second)
105         sameTarget.push_back(MMI->first);
106       
107     }
108     
109     //so more than one branch to same target
110     if(sameTarget.size()){
111
112       sameTarget.push_back(MI->first);
113
114       BasicBlock *newBB = new BasicBlock("newCommon", MI->first->getParent());
115       BranchInst *newBranch = new BranchInst(MI->second);
116
117       newBB->getInstList().push_back(newBranch);
118
119       std::map<PHINode *, std::vector<unsigned int> > phiMap;
120
121       
122       for(std::vector<BasicBlock *>::iterator VBI = sameTarget.begin(),
123             VBE = sameTarget.end(); VBI != VBE; ++VBI){
124
125         //std::cerr<<(*VBI)->getName()<<"\n";
126
127         BranchInst *ti = cast<BranchInst>((*VBI)->getTerminator());
128         unsigned char index = 1;
129         if(ti->getSuccessor(0) == MI->second){
130           index = 0;
131         }
132
133         ti->setSuccessor(index, newBB);
134
135         for(BasicBlock::iterator BB2Inst = MI->second->begin(), 
136               BBend = MI->second->end(); BB2Inst != BBend; ++BB2Inst){
137           
138           if (PHINode *phiInst = dyn_cast<PHINode>(BB2Inst)){
139             int bbIndex;
140             bbIndex = phiInst->getBasicBlockIndex(*VBI);
141             if(bbIndex>=0){
142               phiMap[phiInst].push_back(bbIndex);
143               //phiInst->setIncomingBlock(bbIndex, newBB); 
144             }
145           }
146         }
147       }
148
149       for(std::map<PHINode *, std::vector<unsigned int> >::iterator
150             PI = phiMap.begin(), PE = phiMap.end(); PI != PE; ++PI){
151         
152         PHINode *phiNode = new PHINode(PI->first->getType(), "phi", newBranch);
153         for(std::vector<unsigned int>::iterator II = PI->second.begin(),
154               IE = PI->second.end(); II != IE; ++II){
155           phiNode->addIncoming(PI->first->getIncomingValue(*II),
156                                PI->first->getIncomingBlock(*II));
157         }
158
159         std::vector<BasicBlock *> tempBB;
160         for(std::vector<unsigned int>::iterator II = PI->second.begin(),
161               IE = PI->second.end(); II != IE; ++II){
162           tempBB.push_back(PI->first->getIncomingBlock(*II));
163         }
164
165         for(std::vector<BasicBlock *>::iterator II = tempBB.begin(),
166               IE = tempBB.end(); II != IE; ++II){
167           PI->first->removeIncomingValue(*II);
168         }
169
170         PI->first->addIncoming(phiNode, newBB);
171       }
172       //std::cerr<<"%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n";
173       //std::cerr<<MI->second;
174       //std::cerr<<"-----------------------------------\n";
175       //std::cerr<<newBB;
176       //std::cerr<<"END%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n";
177       
178     }
179   }
180 }
181
182 //getting the backedges in a graph
183 //Its a variation of DFS to get the backedges in the graph
184 //We get back edges by associating a time
185 //and a color with each vertex.
186 //The time of a vertex is the time when it was first visited
187 //The color of a vertex is initially WHITE,
188 //Changes to GREY when it is first visited,
189 //and changes to BLACK when ALL its neighbors
190 //have been visited
191 //So we have a back edge when we meet a successor of
192 //a node with smaller time, and GREY color
193 void CombineBranches::getBackEdges(Function &F){
194   std::map<BasicBlock *, Color > color;
195   std::map<BasicBlock *, int> d;
196   std::map<BasicBlock *, BasicBlock *> be;
197   int time=0;
198   getBackEdgesVisit(F.begin(), color, d, time, be);
199
200   removeRedundant(be);
201 }
202
203 //Per function pass for inserting counters and call function
204 bool CombineBranches::runOnFunction(Function &F){
205   
206   if(F.isExternal()) {
207     return false;
208   }
209
210   //if(F.getName() == "main"){
211    // F.setName("llvm_gprof_main");
212   //}
213   
214   //std::cerr<<F;
215   //std::cerr<<"///////////////////////////////////////////////\n";
216   getBackEdges(F);
217   
218   return true;
219 }