LoopSplitter.cpp 6.46 KB
Newer Older
1
#include "LoopSplitter.h"
guruhegde's avatar
guruhegde committed
2
#include "ForLoop.h"
guruhegde's avatar
guruhegde committed
3
#include "ForLoopV2.h"
guruhegde's avatar
guruhegde committed
4
#include "Util.h"
guruhegde's avatar
guruhegde committed
5 6

#include <llvm/ADT/SmallVector.h>
guruhegde's avatar
guruhegde committed
7
#include <llvm/ADT/STLExtras.h>
guruhegde's avatar
guruhegde committed
8 9
#include <llvm/IR/CFG.h>
#include <llvm/IR/Dominators.h>
guruhegde's avatar
guruhegde committed
10 11 12
#include <llvm/IR/IRBuilder.h>

#include <iostream>
guruhegde's avatar
guruhegde committed
13
#include <vector>
guruhegde's avatar
guruhegde committed
14 15
#include <string>

guruhegde's avatar
guruhegde committed
16
using namespace std;
guruhegde's avatar
guruhegde committed
17 18 19 20 21 22
using namespace llvm;

#define DEBUG_TYPE "tas-batch-process"

namespace tas {

guruhegde's avatar
guruhegde committed
23
void LoopSplitter::addAdapterBasicBlocks(Loop * L, Instruction * SP, Value * Idx) {
guruhegde's avatar
guruhegde committed
24
  auto TopHalf = SP->getParent();
guruhegde's avatar
guruhegde committed
25
  auto BottomHalf = TopHalf->splitBasicBlock(SP);
guruhegde's avatar
guruhegde committed
26 27 28 29 30 31

  auto CollectBB = BasicBlock::Create(F->getContext(), "collector", F, BottomHalf);
  auto DistBB = BasicBlock::Create(F->getContext(), "distributor", F, BottomHalf);

  BranchInst::Create(DistBB, CollectBB);
  setSuccessor(TopHalf, CollectBB);
guruhegde's avatar
guruhegde committed
32
  setSuccessor(CollectBB, DistBB);
guruhegde's avatar
guruhegde committed
33 34 35
  LoopSplitEdgeBlocks.push_back(CollectBB);

  IRBuilder<> Builder(&F->getEntryBlock().front());
guruhegde's avatar
guruhegde committed
36
  auto BrTgtArray = createArray(F, Builder.getInt32Ty(), 32 /*XXX Max Batch size*/);
guruhegde's avatar
guruhegde committed
37 38

  Builder.SetInsertPoint(DistBB);
guruhegde's avatar
guruhegde committed
39
  auto IdxVal = Builder.CreateLoad(Idx);
guruhegde's avatar
guruhegde committed
40
  auto IdxVal64 = Builder.CreateSExtOrBitCast(IdxVal, Builder.getInt64Ty());
guruhegde's avatar
guruhegde committed
41
  auto BrValPtr = Builder.CreateGEP(BrTgtArray, {Builder.getInt64(0), IdxVal64});
guruhegde's avatar
guruhegde committed
42
  auto BrVal = Builder.CreateLoad(BrValPtr);
43
  SwitchI = Builder.CreateSwitch(BrVal, BottomHalf);
guruhegde's avatar
guruhegde committed
44

guruhegde's avatar
guruhegde committed
45
  writeToAsmFile(*F->getParent());
guruhegde's avatar
guruhegde committed
46

guruhegde's avatar
guruhegde committed
47
  SmallVector<BasicBlock *, 4> DivergeBlocks;
guruhegde's avatar
guruhegde committed
48 49 50
  auto DivergeBlock = TopHalf->getUniquePredecessor();
  if (DivergeBlock != L->getHeader()) {
    DivergeBlocks.push_back(DivergeBlock);
guruhegde's avatar
guruhegde committed
51
  }
guruhegde's avatar
guruhegde committed
52 53 54 55 56 57 58

  for (auto & DivergeBB : DivergeBlocks) {
    auto TermI = cast<BranchInst>(DivergeBB->getTerminator());
    auto Cond = TermI->getCondition();

    Builder.SetInsertPoint(TermI);

guruhegde's avatar
guruhegde committed
59 60 61 62 63 64 65
    auto FalseBB = TermI->getSuccessor(1);
    auto TgtBBVal = Builder.CreateSelect(Cond,
                                        BBToId[TermI->getSuccessor(0)],
                                        BBToId[FalseBB]);

    auto IdxVal = Builder.CreateLoad(Idx);
    auto IdxVal64 = Builder.CreateSExtOrBitCast(IdxVal, Builder.getInt64Ty());
guruhegde's avatar
guruhegde committed
66
    auto BrValPtr = Builder.CreateGEP(BrTgtArray, {Builder.getInt64(0), IdxVal64});
guruhegde's avatar
guruhegde committed
67
    Builder.CreateStore(TgtBBVal, BrValPtr);
guruhegde's avatar
guruhegde committed
68
    TermI->setSuccessor(1, CollectBB);
guruhegde's avatar
guruhegde committed
69
    SwitchI->addCase(BBToId[FalseBB], FalseBB);
guruhegde's avatar
guruhegde committed
70 71 72
  }
}

guruhegde's avatar
guruhegde committed
73 74
bool LoopSplitter::prepareForLoopSplit(Function *F, Loop * L0, Stats & stat) {
  auto Idx = getLoopIndexVar(L0);
75 76
  auto AnnotatedVars = detectExpPtrVars(F);
  auto VarUsePoints = detectExpPtrUses(AnnotatedVars);
guruhegde's avatar
guruhegde committed
77 78 79 80

  // Add unique id to each basic block
  unsigned i = 0;
  IRBuilder<> Builder(F->getContext());
guruhegde's avatar
guruhegde committed
81 82 83 84
  auto SetIntValForBB = [&] (const auto & BB) {
                        BBToId.insert(std::make_pair(&BB, Builder.getInt32(++i)));
  };
  for_each(*F, SetIntValForBB);
guruhegde's avatar
guruhegde committed
85

86 87
  for_each(VarUsePoints,
      [&] (auto & VarUse) { insertLLVMPrefetchIntrinsic(F, VarUse); });
guruhegde's avatar
guruhegde committed
88

guruhegde's avatar
guruhegde committed
89
  for_each(VarUsePoints,
guruhegde's avatar
guruhegde committed
90
      [&] (auto & VarUse) { addAdapterBasicBlocks(L0, VarUse, Idx); });
guruhegde's avatar
guruhegde committed
91
  
guruhegde's avatar
guruhegde committed
92 93
  stat.AnnotatedVarsSize = AnnotatedVars.size();
  stat.VarUsePointsSize = VarUsePoints.size();
guruhegde's avatar
guruhegde committed
94
  return VarUsePoints.size() != 0;
guruhegde's avatar
guruhegde committed
95 96
}

guruhegde's avatar
guruhegde committed
97
void LoopSplitter::doLoopSplit(Function * F, Loop * L0, BasicBlock * SplitBlock) {
guruhegde's avatar
guruhegde committed
98
  auto OldHeader = L0->getHeader();
guruhegde's avatar
guruhegde committed
99 100 101 102
  auto OldEntry = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(0);
  auto PreLoopBB = getPreLoopBlock(L0);
  auto PostLoopBB = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(1);
  auto MidBlock = SplitBlock->getUniqueSuccessor();
guruhegde's avatar
guruhegde committed
103

guruhegde's avatar
guruhegde committed
104
  auto LBT = LoopBodyTraverser(L0);
guruhegde's avatar
guruhegde committed
105

guruhegde's avatar
guruhegde committed
106
  // Collect Blocks in range [OldEntry, MidBlock)
guruhegde's avatar
guruhegde committed
107
  vector<BasicBlock *> TopLoopBlocks;
guruhegde's avatar
guruhegde committed
108
  LBT.traverse(TopLoopBlocks, OldEntry, MidBlock);
guruhegde's avatar
guruhegde committed
109

guruhegde's avatar
guruhegde committed
110
  // Collect Blocks in range [MidBlock, Latch)
guruhegde's avatar
guruhegde committed
111
  vector<BasicBlock *> BottomLoopBlocks;
guruhegde's avatar
guruhegde committed
112
  LBT.traverse(BottomLoopBlocks, MidBlock, L0->getLoopLatch());
guruhegde's avatar
guruhegde committed
113

guruhegde's avatar
guruhegde committed
114 115
  auto BottomLoop = IRLoop();
  BottomLoop.extractLoopSkeleton(L0);
guruhegde's avatar
guruhegde committed
116

guruhegde's avatar
guruhegde committed
117
  auto TopLoop = IRLoop();
guruhegde's avatar
guruhegde committed
118
  TopLoop.constructEmptyLoop(getLoopTripCount(L0), PreLoopBB ? PreLoopBB : BottomLoop.getHeader());
guruhegde's avatar
guruhegde committed
119

guruhegde's avatar
guruhegde committed
120 121
  TopLoop.setLoopBlocks(TopLoopBlocks);
  BottomLoop.setLoopBlocks(BottomLoopBlocks);
guruhegde's avatar
guruhegde committed
122

guruhegde's avatar
guruhegde committed
123
  setSuccessor(PreLoopBB, TopLoop.getPreHeader());
guruhegde's avatar
guruhegde committed
124 125 126 127
  // TODO Setting index value to 0 when entering new loop.
  BasicBlock * BLoopEntry = BottomLoop.getPreHeader() == &F->getEntryBlock() ?
                            BottomLoop.getHeader() : BottomLoop.getPreHeader();
  setSuccessor(TopLoop.getHeader(), BLoopEntry, 1); 
guruhegde's avatar
guruhegde committed
128
  setSuccessor(BottomLoop.getHeader(), PostLoopBB, 1);
129

guruhegde's avatar
guruhegde committed
130
  /*
131 132 133 134 135 136 137 138
  // If FalseBB is terminating instruction, use latch block as target instead.
  SmallVector<BasicBlock *, 4> Returns;
  getReturnBlocks(F, Returns);
  for (auto & Case : SwitchI->cases()) {
    if (find(Returns, Case.getCaseSuccessor()) != Returns.end()) {
      Case.setSuccessor(L0->getLoopLatch());
    }
  }
guruhegde's avatar
guruhegde committed
139 140
  */
}
guruhegde's avatar
guruhegde committed
141

guruhegde's avatar
guruhegde committed
142
/*
guruhegde's avatar
guruhegde committed
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
  DenseSet<BasicBlock *> Blocks;
  auto TruePathBB = NewHeader->getTerminator()->getSuccessor(0);
  visitSuccessor(Blocks, TruePathBB, NewLatch);

  auto BB = Blocks.begin();
  auto BE = Blocks.end();
  auto IB = (*BB)->begin();
  auto IE = (*BB)->end();
  while (BB != BE) {
    while (IB != IE) {
      auto UB = (*IB).user_begin();
      auto UE = (*IB).user_end();
      while (UB != UE) {
        auto * U = *UB++;
        if (auto * Inst = dyn_cast<Instruction>(U)) {
          if (Blocks.find(Inst->getParent()) == Blocks.end()) {
            auto arrayPtr = createArray(F, (*IB).getType(), 32);

            Builder.SetInsertPoint((*IB).getNextNode());
            IndexVarVal = Builder.CreateLoad(IndexVar);
            auto ptr = Builder.CreateGEP(arrayPtr, {Builder.getInt64(0), IndexVarVal});
            auto str = Builder.CreateStore(&*IB, ptr);
            errs() << *ptr << "\n" << *str << "\n"; 
          }
        }
      }
      ++IB;
    }
    ++BB;
  }
guruhegde's avatar
guruhegde committed
173
}
guruhegde's avatar
guruhegde committed
174
*/
guruhegde's avatar
guruhegde committed
175

guruhegde's avatar
guruhegde committed
176
bool LoopSplitter::run() {
guruhegde's avatar
guruhegde committed
177 178 179 180 181
  // If no loops, we are done.
  if (LI->begin() == LI->end()) return false;

  // XXX Assume only one loop for now.
  auto L0 = *LI->begin();
guruhegde's avatar
guruhegde committed
182

guruhegde's avatar
guruhegde committed
183 184 185
  ExitBlock = L0->getExitBlock();
  assert (ExitBlock && "Loop must have a single exit block!");

guruhegde's avatar
guruhegde committed
186 187 188
  auto ParentLoop = IRLoop();
  ParentLoop.analyze(L0);

guruhegde's avatar
guruhegde committed
189
  bool changed = prepareForLoopSplit(F, L0, stat);
guruhegde's avatar
guruhegde committed
190 191
  if (!changed) return false;

guruhegde's avatar
guruhegde committed
192 193 194 195 196 197 198
  DominatorTree DT (*F);
  LoopInfo NewLI (DT);
  L0 = *NewLI.begin();
  if (NewLI.begin() == NewLI.end()) {
    errs() << "Loop info lost, something wrong with preparation\n";
    return false;
  }
guruhegde's avatar
guruhegde committed
199
  auto & SplitBB = LoopSplitEdgeBlocks.front();
guruhegde's avatar
guruhegde committed
200 201 202 203

  doLoopSplit(F, L0, SplitBB);

  assert(stat.AnnotatedVarsSize == 1 && "Atleast one annotation must be there!");
guruhegde's avatar
guruhegde committed
204 205 206 207 208

  return true;
}

}