Implement doLoopSplit()

* not tested yet.

src/BatchMaker.cpp

@@ -27,6 +27,8 @@ using namespace llvm;
 
 namespace tas {
 
+// TODO Use return val array
+
 /// Old Function Prototype:
 ///    Ret Fn(Type1 A, Type2 B BATCH_ARG, Type3 C BATCH_ARG);
 /// BatchForm Fn Prototype:

src/ForLoop.cpp

@@ -9,24 +9,21 @@ using namespace llvm;
 namespace tas {
 
 TASForLoop::TASForLoop(LLVMContext & Ctx, BasicBlock * Prev,
-    BasicBlock * Next, const std::string & Name, Function * F, llvm::Value * TC, AllocaInst * IP)
+    BasicBlock * Next, const std::string & Name, Function * F,
+    llvm::Value * TC, AllocaInst * IP)
   : F(F), Name (std::move(Name)), TripCount(TC), IdxVarPtr(IP)
 {
   addEmptyLoop(Ctx, Prev, Next);
 }
 
-void TASForLoop::addEmptyLoop(LLVMContext & Ctx, BasicBlock * Prev, BasicBlock * Next) {
+void TASForLoop::addEmptyLoop(LLVMContext & Ctx, BasicBlock * Prev,
+                              BasicBlock * Next) {
   PreHeader = BasicBlock::Create(Ctx, Name + ".preheader", F, Next);
   Header = BasicBlock::Create(Ctx, Name + ".header", F, Next);
   Latch = BasicBlock::Create(Ctx, Name + ".latch", F, Next);
 
   // Update phi node edge if any
-  IRBuilder<> Builder(Next);
-  auto * PN = &*(Next->phis().begin());
-  if (PN)
-    PN->addIncoming(Builder.getInt32(0), Header);
-
-  Builder.SetInsertPoint(PreHeader);
+  IRBuilder<> Builder(PreHeader);
   Builder.CreateStore(Builder.getInt32(0), IdxVarPtr);
   Builder.CreateBr(Header);
 
@@ -38,7 +35,7 @@ void TASForLoop::addEmptyLoop(LLVMContext & Ctx, BasicBlock * Prev, BasicBlock *
   Builder.SetInsertPoint(Header);
   IndexVar = Builder.CreateLoad(IdxVarPtr);
   auto * icmp = Builder.CreateICmpSLT(IndexVar, TripCount, "loop-predicate");
-  
+
   // Stitch entry point in control flow.
   if (Prev) {
     if (Prev->getTerminator()->getNumOperands() == 3) 
@@ -57,7 +54,8 @@ void TASForLoop::setLoopBody(BasicBlock * BodyBB) {
   BodyBB->getTerminator()->setSuccessor(0, Latch);
 }
 
-void TASForLoop::setLoopBody(BasicBlock * EntryBB, BasicBlock * ExitingBB) {
+void TASForLoop::setLoopBody(BasicBlock * EntryBB,
+                             BasicBlock * ExitingBB) {
   EntryBody = EntryBB;
   ExitingBody = ExitingBB;
   Header->getTerminator()->setSuccessor(0, EntryBody);

src/LoopSplitter.cpp

@@ -17,7 +17,7 @@ using namespace llvm;
 
 namespace tas {
 
-void LoopSplitter::addAdapterBasicBlocks(Instruction * SP) {
+void LoopSplitter::addAdapterBasicBlocks(Instruction * SP, Value * Idx) {
   auto TopHalf = SP->getParent();
   auto BottomHalf = TopHalf->splitBasicBlock(SP->getNextNode());
 
@@ -26,18 +26,22 @@ void LoopSplitter::addAdapterBasicBlocks(Instruction * SP) {
 
   BranchInst::Create(DistBB, CollectBB);
   setSuccessor(TopHalf, CollectBB);
+  setSuccessor(CollectBB, DistBB);
   LoopSplitEdgeBlocks.push_back(CollectBB);
 
-  // General case: Find all edges from basic blocks above tophalf and connect it to
-  // CollectBB, use switch Inst
+  // General case: Find all edges from basic blocks above tophalf
+  // and connect it to CollectBB, use switch Inst
   IRBuilder<> Builder(&F->getEntryBlock().front());
-  auto BrTargetAlloca = Builder.CreateAlloca(Builder.getInt32Ty());
+  auto BrTgtArray = createArray(F, Builder.getInt32Ty(), 32 /*XXX Max Batch size*/);
+  //auto BrTargetAlloca = Builder.CreateAlloca(Builder.getInt32Ty());
 
   Builder.SetInsertPoint(DistBB);
-  auto BrVal = Builder.CreateLoad(BrTargetAlloca);
+  auto BrValPtr = Builder.CreateGEP(BrTgtArray, {Builder.getInt64(0), Idx});
+  auto BrVal = Builder.CreateLoad(BrValPtr);
   auto SwitchI = Builder.CreateSwitch(BrVal, BottomHalf);
 
-  // XXX We assume now CFG we have is the one after block predication transformation.
+  // XXX We assume now CFG we have is the one after block
+  // predication transformation.
   SmallVector<BasicBlock *, 4> DivergeBlocks;
   SmallVector<std::pair<Value *, BasicBlock *>, 4> TargetBlocks;
   DivergeBlocks.push_back(TopHalf->getUniquePredecessor());
@@ -50,13 +54,14 @@ void LoopSplitter::addAdapterBasicBlocks(Instruction * SP) {
     auto TgtBB = TermI->getSuccessor(1);
     auto TgtBBVal = BBToId[TgtBB];
 
-    Builder.CreateStore(TgtBBVal, BrTargetAlloca);
+    auto BrValPtr = Builder.CreateGEP(BrTgtArray, {Builder.getInt64(0), Idx});
+    Builder.CreateStore(TgtBBVal, BrValPtr);
     TermI->setSuccessor(1, CollectBB);
     SwitchI->addCase(TgtBBVal, TgtBB);
   }
 }
 
-bool LoopSplitter::prepareForLoopSplit(Function *F, Stats & stat) {
+bool LoopSplitter::prepareForLoopSplit(Function *F, Stats & stat, Value * Idx) {
   auto AnnotatedVars = detectExpPtrVars(F);
   auto VarUsePoints = detectExpPtrUses(AnnotatedVars);
 
@@ -64,10 +69,12 @@ bool LoopSplitter::prepareForLoopSplit(Function *F, Stats & stat) {
   unsigned i = 0;
   IRBuilder<> Builder(F->getContext());
   for_each(*F,
-      [&] (const auto & BB) { BBToId.insert(std::make_pair(&BB, Builder.getInt32(++i))); });
+      [&] (const auto & BB) {
+      BBToId.insert(std::make_pair(&BB, Builder.getInt32(++i)));
+    });
 
   for_each(VarUsePoints,
-      [&] (auto & VarUse) { addAdapterBasicBlocks(VarUse); });
+      [&] (auto & VarUse) { addAdapterBasicBlocks(VarUse, Idx); });
 
   for_each(VarUsePoints,
       [&] (auto & VarUse) { insertLLVMPrefetchIntrinsic(F, VarUse); });
@@ -77,17 +84,70 @@ bool LoopSplitter::prepareForLoopSplit(Function *F, Stats & stat) {
   return VarUsePoints.size() != 0;
 }
 
+Value * getLoopTripCount(Loop * L0) {
+  auto Header = L0->getHeader();
+  auto Cond = cast<BranchInst>(Header->getTerminator())->getCondition();
+  return cast<LoadInst>(cast<ICmpInst>(Cond)->getOperand(1))->getOperand(0);
+}
+
+void doLoopSplit(Function * F, Loop * L0, BasicBlock * SplitBlock) {
+  auto OldHeader = L0->getHeader();
+  auto ExitBlock = OldHeader->getTerminator()->getSuccessor(1);
+  auto DistBB = SplitBlock->getUniqueSuccessor();
+
+  BasicBlock * EntryBlock = nullptr;
+  for (auto BB : predecessors(OldHeader)) {
+    if (BB != L0->getLoopLatch())
+      EntryBlock = BB;
+  }
+
+  auto NewPreHeader = BasicBlock::Create(F->getContext(), "preheader", F);
+  auto NewHeader = BasicBlock::Create(F->getContext(), "header", F);
+  auto NewLatch = BasicBlock::Create(F->getContext(), "latch", F);
+
+  auto IndexVar = getLoopIndexVar(L0);
+  auto TripCount = cast<AllocaInst>(getLoopTripCount(L0));
+  errs() << *IndexVar << "  " << *TripCount << "\n";
+
+  IRBuilder<> Builder(NewPreHeader);
+  Builder.CreateStore(Builder.getInt32(0), IndexVar);
+  Builder.CreateBr(NewHeader);
+
+  Builder.SetInsertPoint(NewLatch);
+  auto BI = Builder.CreateBr(NewHeader);
+  addIncrementIndexOp(IndexVar, BI);
+
+  Builder.SetInsertPoint(NewHeader);
+  auto IndexVarVal = Builder.CreateLoad(IndexVar);
+  auto TripCountVal = Builder.CreateLoad(TripCount);
+  auto * icmp = Builder.CreateICmpSLT(IndexVarVal, TripCountVal,
+                                      "loop-predicate");
+  auto NewEntryBB = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(0);
+  Builder.CreateCondBr(icmp, NewEntryBB, L0->getHeader());
+
+  SplitBlock->getTerminator()->setSuccessor(0, NewLatch);
+  OldHeader->getTerminator()->setSuccessor(0, DistBB);
+  OldHeader->getTerminator()->setSuccessor(1, ExitBlock);
+  setSuccessor(EntryBlock, NewPreHeader);
+}
+
 bool LoopSplitter::run() {
-  bool changed = prepareForLoopSplit(F, stat);
-  if (!changed) return false;
 
   // If no loops, we are done.
   if (LI->begin() == LI->end()) return false;
 
+
   // XXX Assume only one loop for now.
   auto L0 = *LI->begin();
+  auto IndexVar = getLoopIndexVar(L0);
+
+  bool changed = prepareForLoopSplit(F, stat, IndexVar);
+  if (!changed) return false;
+
+  auto & SplitBB = LoopSplitEdgeBlocks.front();
+  doLoopSplit(F, L0, SplitBB);
 
-  //assert (L0->getLoopPreheader() && "Loop doesn't have preheader!");
+  F->print(errs());
 
   return true;
 }

src/LoopSplitter.h

@@ -37,14 +37,14 @@ class LoopSplitter {
 
   llvm::DenseMap<const llvm::BasicBlock *, llvm::ConstantInt *> BBToId;
   llvm::SmallVector<llvm::BasicBlock *, 4> LoopSplitEdgeBlocks;
-  bool prepareForLoopSplit(llvm::Function * F, Stats & stat);
+  bool prepareForLoopSplit(llvm::Function * F, Stats & stat, llvm::Value * Idx);
   public:
   LoopSplitter(llvm::Function * F_, llvm::LoopInfo * LI_)
     : F(F_), LI(LI_) {}
 
   bool run();
   Stats & getStats() { return stat; }
-  void addAdapterBasicBlocks(llvm::Instruction * SP);
+  void addAdapterBasicBlocks(llvm::Instruction * SP, llvm::Value * Idx);
 };
 
 } // tas namespace