diff options
Diffstat (limited to 'llvm_mode/MarkNodes.cc')
| -rw-r--r-- | llvm_mode/MarkNodes.cc | 355 |
1 files changed, 355 insertions, 0 deletions
diff --git a/llvm_mode/MarkNodes.cc b/llvm_mode/MarkNodes.cc new file mode 100644 index 00000000..a156fccb --- /dev/null +++ b/llvm_mode/MarkNodes.cc @@ -0,0 +1,355 @@ +#include <algorithm> +#include <map> +#include <queue> +#include <set> +#include <vector> +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/CFG.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" +#include "llvm/Pass.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +using namespace llvm; + +DenseMap<BasicBlock *, uint32_t> LMap; +std::vector<BasicBlock *> Blocks; +std::set<uint32_t> Marked , Markabove; +std::vector< std::vector<uint32_t> > Succs , Preds; + +void reset(){ + LMap.clear(); + Blocks.clear(); + Marked.clear(); + Markabove.clear(); +} + +uint32_t start_point; + +void labelEachBlock(Function *F) { + // Fake single endpoint; + LMap[NULL] = Blocks.size(); + Blocks.push_back(NULL); + + // Assign the unique LabelID to each block; + for (auto I = F->begin(), E = F->end(); I != E; ++I) { + BasicBlock *BB = &*I; + LMap[BB] = Blocks.size(); + Blocks.push_back(BB); + } + + start_point = LMap[&F->getEntryBlock()]; +} + +void buildCFG(Function *F) { + Succs.resize( Blocks.size() ); + Preds.resize( Blocks.size() ); + for( size_t i = 0 ; i < Succs.size() ; i ++ ){ + Succs[ i ].clear(); + Preds[ i ].clear(); + } + + //uint32_t FakeID = 0; + for (auto S = F->begin(), E = F->end(); S != E; ++S) { + BasicBlock *BB = &*S; + uint32_t MyID = LMap[BB]; + //if (succ_begin(BB) == succ_end(BB)) { + //Succs[MyID].push_back(FakeID); + //Marked.insert(MyID); + //} + for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I) { + Succs[MyID].push_back(LMap[*I]); + } + } +} + +std::vector< std::vector<uint32_t> > tSuccs; +std::vector<bool> tag , indfs; + +void DFStree(size_t now_id) { + if(tag[now_id]) return; + tag[now_id]=true; + indfs[now_id]=true; + for (auto succ: tSuccs[now_id]) { + if(tag[succ] and indfs[succ]) { + Marked.insert(succ); + Markabove.insert(succ); + continue; + } + Succs[now_id].push_back(succ); + Preds[succ].push_back(now_id); + DFStree(succ); + } + indfs[now_id]=false; +} +void turnCFGintoDAG(Function *F) { + tSuccs = Succs; + tag.resize(Blocks.size()); + indfs.resize(Blocks.size()); + for (size_t i = 0; i < Blocks.size(); ++ i) { + Succs[i].clear(); + tag[i]=false; + indfs[i]=false; + } + DFStree(start_point); + for (size_t i = 0; i < Blocks.size(); ++ i) + if( Succs[i].empty() ){ + Succs[i].push_back(0); + Preds[0].push_back(i); + } +} + +uint32_t timeStamp; +namespace DominatorTree{ + std::vector< std::vector<uint32_t> > cov; + std::vector<uint32_t> dfn, nfd, par, sdom, idom, mom, mn; + + bool Compare(uint32_t u, uint32_t v) { + return dfn[u] < dfn[v]; + } + uint32_t eval(uint32_t u) { + if( mom[u] == u ) return u; + uint32_t res = eval( mom[u] ); + if(Compare(sdom[mn[mom[u]]] , sdom[mn[u]])) { + mn[u] = mn[mom[u]]; + } + return mom[u] = res; + } + + void DFS(uint32_t now) { + timeStamp += 1; + dfn[now] = timeStamp; + nfd[timeStamp - 1] = now; + for( auto succ : Succs[now] ) { + if( dfn[succ] == 0 ) { + par[succ] = now; + DFS(succ); + } + } + } + + void DominatorTree(Function *F) { + if( Blocks.empty() ) return; + uint32_t s = start_point; + + // Initialization + mn.resize(Blocks.size()); + cov.resize(Blocks.size()); + dfn.resize(Blocks.size()); + nfd.resize(Blocks.size()); + par.resize(Blocks.size()); + mom.resize(Blocks.size()); + sdom.resize(Blocks.size()); + idom.resize(Blocks.size()); + + for( uint32_t i = 0 ; i < Blocks.size() ; i ++ ) { + dfn[i] = 0; + nfd[i] = Blocks.size(); + cov[i].clear(); + idom[i] = mom[i] = mn[i] = sdom[i] = i; + } + + timeStamp = 0; + DFS(s); + + for( uint32_t i = Blocks.size() - 1 ; i >= 1u ; i -- ) { + uint32_t now = nfd[i]; + if( now == Blocks.size() ) { + continue; + } + for( uint32_t pre : Preds[ now ] ) { + if( dfn[ pre ] ) { + eval(pre); + if( Compare(sdom[mn[pre]], sdom[now]) ) { + sdom[now] = sdom[mn[pre]]; + } + } + } + cov[sdom[now]].push_back(now); + mom[now] = par[now]; + for( uint32_t x : cov[par[now]] ) { + eval(x); + if( Compare(sdom[mn[x]], par[now]) ) { + idom[x] = mn[x]; + } else { + idom[x] = par[now]; + } + } + } + + for( uint32_t i = 1 ; i < Blocks.size() ; i += 1 ) { + uint32_t now = nfd[i]; + if( now == Blocks.size() ) { + continue; + } + if(idom[now] != sdom[now]) + idom[now] = idom[idom[now]]; + } + } +}; // End of DominatorTree + +std::vector<uint32_t> Visited, InStack; +std::vector<uint32_t> TopoOrder, InDeg; +std::vector< std::vector<uint32_t> > t_Succ , t_Pred; + +void Go(uint32_t now, uint32_t tt) { + if( now == tt ) return; + Visited[now] = InStack[now] = timeStamp; + + for(uint32_t nxt : Succs[now]) { + if(Visited[nxt] == timeStamp and InStack[nxt] == timeStamp) { + Marked.insert(nxt); + } + t_Succ[now].push_back(nxt); + t_Pred[nxt].push_back(now); + InDeg[nxt] += 1; + if(Visited[nxt] == timeStamp) { + continue; + } + Go(nxt, tt); + } + + InStack[now] = 0; +} + +void TopologicalSort(uint32_t ss, uint32_t tt) { + timeStamp += 1; + + Go(ss, tt); + + TopoOrder.clear(); + std::queue<uint32_t> wait; + wait.push(ss); + while( not wait.empty() ) { + uint32_t now = wait.front(); wait.pop(); + TopoOrder.push_back(now); + for(uint32_t nxt : t_Succ[now]) { + InDeg[nxt] -= 1; + if(InDeg[nxt] == 0u) { + wait.push(nxt); + } + } + } +} + +std::vector< std::set<uint32_t> > NextMarked; +bool Indistinguish(uint32_t node1, uint32_t node2) { + if(NextMarked[node1].size() > NextMarked[node2].size()){ + uint32_t _swap = node1; + node1 = node2; + node2 = _swap; + } + for(uint32_t x : NextMarked[node1]) { + if( NextMarked[node2].find(x) != NextMarked[node2].end() ) { + return true; + } + } + return false; +} + +void MakeUniq(uint32_t now) { + bool StopFlag = false; + if (Marked.find(now) == Marked.end()) { + for(uint32_t pred1 : t_Pred[now]) { + for(uint32_t pred2 : t_Pred[now]) { + if(pred1 == pred2) continue; + if(Indistinguish(pred1, pred2)) { + Marked.insert(now); + StopFlag = true; + break; + } + } + if (StopFlag) { + break; + } + } + } + if(Marked.find(now) != Marked.end()) { + NextMarked[now].insert(now); + } else { + for(uint32_t pred : t_Pred[now]) { + for(uint32_t x : NextMarked[pred]) { + NextMarked[now].insert(x); + } + } + } +} + +void MarkSubGraph(uint32_t ss, uint32_t tt) { + TopologicalSort(ss, tt); + if(TopoOrder.empty()) return; + + for(uint32_t i : TopoOrder) { + NextMarked[i].clear(); + } + + NextMarked[TopoOrder[0]].insert(TopoOrder[0]); + for(uint32_t i = 1 ; i < TopoOrder.size() ; i += 1) { + MakeUniq(TopoOrder[i]); + } +} + +void MarkVertice(Function *F) { + uint32_t s = start_point; + + InDeg.resize(Blocks.size()); + Visited.resize(Blocks.size()); + InStack.resize(Blocks.size()); + t_Succ.resize(Blocks.size()); + t_Pred.resize(Blocks.size()); + NextMarked.resize(Blocks.size()); + + for( uint32_t i = 0 ; i < Blocks.size() ; i += 1 ) { + Visited[i] = InStack[i] = InDeg[i] = 0; + t_Succ[i].clear(); + t_Pred[i].clear(); + } + timeStamp = 0; + uint32_t t = 0; + //MarkSubGraph(s, t); + //return; + + while( s != t ) { + MarkSubGraph(DominatorTree::idom[t], t); + t = DominatorTree::idom[t]; + } + +} + +// return {marked nodes} +std::pair<std::vector<BasicBlock *>, + std::vector<BasicBlock *> >markNodes(Function *F) { + assert(F->size() > 0 && "Function can not be empty"); + + reset(); + labelEachBlock(F); + buildCFG(F); + turnCFGintoDAG(F); + DominatorTree::DominatorTree(F); + MarkVertice(F); + + std::vector<BasicBlock *> Result , ResultAbove; + for( uint32_t x : Markabove ) { + auto it = Marked.find( x ); + if( it != Marked.end() ) + Marked.erase( it ); + if( x ) + ResultAbove.push_back(Blocks[x]); + } + for( uint32_t x : Marked ) { + if (x == 0) { + continue; + } else { + Result.push_back(Blocks[x]); + } + } + + return { Result , ResultAbove }; +} |