Implement IRLoop class

src/CMakeLists.txt

@@ -3,6 +3,7 @@ add_library(tasopt_src
             BlockPredication.cpp
             CacheUsageInfo.cpp
             ForLoop.cpp
+            ForLoopV2.cpp
             LoopSplitter.cpp
             PacketPathAnalysis.cpp
             ToolUtil.cpp

src/ForLoop.cpp

@@ -22,7 +22,6 @@ void TASForLoop::addEmptyLoop(LLVMContext & Ctx, BasicBlock * Prev,
   Header = BasicBlock::Create(Ctx, Name + ".header", F, Next);
   Latch = BasicBlock::Create(Ctx, Name + ".latch", F, Next);
 
-  // Update phi node edge if any
   IRBuilder<> Builder(PreHeader);
   Builder.CreateStore(Builder.getInt32(0), IdxVarPtr);
   Builder.CreateBr(Header);

src/ForLoop.h

@@ -35,10 +35,6 @@ public:
   void setLoopBody(llvm::BasicBlock * BodyBB);
   void setLoopBody(llvm::BasicBlock * EntryBB, llvm::BasicBlock * ExitingBB);
 
-  llvm::BasicBlock * getLatchBlock() {
-    return Latch;
-  }
-
   llvm::BasicBlock * getLatch() {
     return Latch;
   }
@@ -51,10 +47,6 @@ public:
     return IndexVar;
   }
 
-  llvm::Value * getIndexVariable64Bit() {
-    return IndexVar;
-  }
-
   llvm::BasicBlock * getBody() {
     if (EntryBody != ExitingBody)
       return nullptr;
@@ -64,6 +56,10 @@ public:
   static unsigned getLoopTripCount() {
     return BATCH_SIZE;
   }
+
+  llvm::Value * getLoopTripCountValue() {
+    return TripCount;
+  }
 };
 
 }

src/ForLoopV2.cpp

+#include "ForLoopV2.h"
+#include "Util.h"
+
+#include <llvm/IR/IRBuilder.h>
+
+using namespace llvm;
+
+namespace tas {
+
+void IRLoop::extractLoopSkeleton(Loop * L) {
+  Header = L->getHeader();
+  Latch = L->getLoopLatch();
+  IdxAlloca = getLoopIndexVar(L);
+  PreHeader = L->getLoopPreheader();
+}
+
+void IRLoop::analyze(Loop * L) {
+  extractLoopSkeleton(L);
+  
+  // Copy loop body
+  for (auto & BB : L->getBlocks()) {
+    Blocks.push_back(BB);
+  }
+}
+
+void IRLoop::constructEmptyLoop(Value * TripCount, BasicBlock * ExitBlock) {
+  auto & Ctx = ExitBlock->getContext();
+  const auto & F = ExitBlock->getParent();
+
+  Latch = BasicBlock::Create(Ctx, "latch", F);
+  Header = BasicBlock::Create(Ctx, "header", F, Latch);
+  PreHeader = BasicBlock::Create(Ctx, "preheader", F, Header);
+  EmptyBody = BasicBlock::Create(Ctx, "emptyblock", F, Latch);
+
+  // Create Loop Index Variable.
+  IRBuilder<> Builder(Ctx);
+  Builder.SetInsertPoint(&F->getEntryBlock().front());
+  IdxAlloca = Builder.CreateAlloca(Builder.getInt32Ty());
+
+  // Set Index variable to 0 in preheader block.
+  Builder.SetInsertPoint(PreHeader);
+  Builder.CreateStore(Builder.getInt32(0), IdxAlloca);
+  Builder.CreateBr(Header);
+
+  // Populate latch block
+  Builder.SetInsertPoint(Latch);
+  auto BackEdge = Builder.CreateBr(Header);
+  addIncrementIndexOp(IdxAlloca, BackEdge);
+
+  // Branch Empty block to latch
+  Builder.SetInsertPoint(EmptyBody);
+  Builder.CreateBr(Latch);
+
+  // Populate header block
+  Builder.SetInsertPoint(Header);
+  auto IdxVal = Builder.CreateLoad(IdxAlloca);
+  auto * Icmp = Builder.CreateICmpSLT(IdxVal, TripCount, "loop-predicate");
+  Builder.CreateCondBr(Icmp, EmptyBody, ExitBlock);
+}
+
+void IRLoop::setLoopBlocks(SmallVectorImpl<BasicBlock *> & Blocks) {
+  for_each(Blocks, [&] (auto & BB) { Blocks.push_back(BB); });
+  assert (!Blocks.empty() && "Blocks can't be empty");
+  setSuccessor(Header, Blocks.front()); // True Path
+  assert (Latch && "Latch is NULL");
+  setSuccessor(Blocks.back(), Latch);
+}
+
+void LoopBodyTraverser::traverse(SmallVectorImpl<BasicBlock *> & Blocks,
+                                 BasicBlock * Start, BasicBlock * End) {
+  if (Start == End) return;
+  Blocks.push_back(Start);
+  for (auto * BB : successors(Start)){
+    // Don't travel out of loop
+    //if (std::find(ExitBlocks.begin(), ExitBlocks.end(), BB) != Blocks.end()) continue;
+    traverse(Blocks, BB, End);
+  }
+}
+
+}

src/ForLoopV2.h

+#ifndef TAS_FORLOOPV2_H
+#define TAS_FORLOOPV2_H
+
+#include <llvm/Analysis/LoopInfo.h>
+#include <llvm/IR/Function.h>
+#include <llvm/IR/Instructions.h>
+#include <llvm/IR/Module.h>
+#include <llvm/IR/Value.h>
+
+namespace tas {
+
+// This is similar to Loop object available in LoopInfo object
+// but simple.
+class IRLoop {
+  llvm::BasicBlock * PreHeader;
+  llvm::BasicBlock * Header;
+  llvm::BasicBlock * Latch;
+  llvm::BasicBlock * EmptyBody; // Useful for creating empty loop.
+  llvm::AllocaInst * IdxAlloca;
+  llvm::SmallVector<llvm::BasicBlock *, 4> Blocks;
+
+public:
+  IRLoop() = default;
+
+  void analyze(llvm::Loop * L);
+  void constructEmptyLoop(llvm::Value * TripCount,
+                          llvm::BasicBlock * InsertAfter);
+  void extractLoopSkeleton(llvm::Loop * L);
+  void setLoopBlocks(llvm::SmallVectorImpl<llvm::BasicBlock *> & Blocks);
+
+  llvm::BasicBlock * getPreHeader() {
+    return PreHeader;
+  }
+
+  llvm::BasicBlock * getHeader() {
+    return Header;
+  }
+};
+
+class LoopBodyTraverser {
+  llvm::SmallVector<llvm::BasicBlock *, 4> ExitBlocks;
+
+public:
+  LoopBodyTraverser(llvm::Loop * L) {
+    L->getExitBlocks(ExitBlocks);
+  }
+
+  void traverse(llvm::SmallVectorImpl<llvm::BasicBlock *> & Blocks,
+                llvm::BasicBlock * Start, llvm::BasicBlock * End);
+
+  void printExitBlocks() {
+    llvm::errs() << "Exit blocks = ";
+    for (auto & BB : ExitBlocks) {
+      BB->printAsOperand(llvm::errs());
+      llvm::errs() << " ";
+    }
+    llvm::errs() << "\n";
+  }
+};
+
+}
+
+#endif

src/LoopSplitter.cpp

 #include "LoopSplitter.h"
 #include "ForLoop.h"
+#include "ForLoopV2.h"
 #include "Util.h"
 
 #include <llvm/ADT/SmallVector.h>
@@ -17,7 +18,7 @@ using namespace llvm;
 
 namespace tas {
 
-void LoopSplitter::addAdapterBasicBlocks(Instruction * SP, Value * Idx) {
+void LoopSplitter::addAdapterBasicBlocks(Loop * L, Instruction * SP, Value * Idx) {
   auto TopHalf = SP->getParent();
   auto BottomHalf = TopHalf->splitBasicBlock(SP);
 
@@ -41,10 +42,24 @@ void LoopSplitter::addAdapterBasicBlocks(Instruction * SP, Value * Idx) {
   auto BrVal = Builder.CreateLoad(BrValPtr);
   SwitchI = Builder.CreateSwitch(BrVal, BottomHalf);
 
-  // XXX We assume now CFG we have is the one after block
-  // predication transformation.
+  // Handle diverge blocks
+  auto OldHeader = L->getHeader();
+  auto OldEntry = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(0);
+
+  SmallVector<BasicBlock *, 4> TopBlocks;
+  SmallVector<BasicBlock *, 4> BottomBlocks;
+  LoopBodyTraverser LBT (L);
+  LBT.traverse(TopBlocks, OldEntry, DistBB);
+  LBT.traverse(BottomBlocks, DistBB, L->getLoopLatch());
+
   SmallVector<BasicBlock *, 4> DivergeBlocks;
-  DivergeBlocks.push_back(TopHalf->getUniquePredecessor());
+  for (auto & BB : TopBlocks) {
+    for (auto * Succ : successors(BB)) {
+      Succ->printAsOperand(errs());
+      if (std::find(BottomBlocks.begin(), BottomBlocks.end(), Succ) != BottomBlocks.end())
+        DivergeBlocks.push_back(BB);
+    }
+  }
 
   for (auto & DivergeBB : DivergeBlocks) {
     auto TermI = cast<BranchInst>(DivergeBB->getTerminator());
@@ -69,86 +84,58 @@ void LoopSplitter::addAdapterBasicBlocks(Instruction * SP, Value * Idx) {
 bool LoopSplitter::prepareForLoopSplit(Function *F, Loop * L0, Stats & stat) {
   auto Idx = getLoopIndexVar(L0);
   auto AnnotatedVars = detectExpPtrVars(F);
-  errs() << "AnnotatedVars = " << AnnotatedVars.size() << "  " << *AnnotatedVars.front() << "\n";
   auto VarUsePoints = detectExpPtrUses(AnnotatedVars);
 
   // Add unique id to each basic block
   unsigned i = 0;
   IRBuilder<> Builder(F->getContext());
-  for_each(*F,
-      [&] (const auto & BB) {
-      BBToId.insert(std::make_pair(&BB, Builder.getInt32(++i)));
-    });
+  auto SetIntValForBB = [&] (const auto & BB) {
+                        BBToId.insert(std::make_pair(&BB, Builder.getInt32(++i)));
+  };
+  for_each(*F, SetIntValForBB);
 
   for_each(VarUsePoints,
       [&] (auto & VarUse) { insertLLVMPrefetchIntrinsic(F, VarUse); });
 
   for_each(VarUsePoints,
-      [&] (auto & VarUse) { addAdapterBasicBlocks(VarUse, Idx); });
+      [&] (auto & VarUse) { addAdapterBasicBlocks(L0, VarUse, Idx); });
   
   stat.AnnotatedVarsSize = AnnotatedVars.size();
   stat.VarUsePointsSize = VarUsePoints.size();
   return VarUsePoints.size() != 0;
 }
 
-Value * getLoopTripCount(Loop * L0) {
-  auto Header = L0->getHeader();
-  auto Cond = cast<BranchInst>(Header->getTerminator())->getCondition();
-  return cast<ICmpInst>(Cond)->getOperand(1);
-}
-
-void visitSuccessor(DenseSet<BasicBlock *> & Blocks, BasicBlock * CurBlock,
-                    BasicBlock * EndBlock) {
-  for (auto * BB : successors(CurBlock)){
-    if (BB == EndBlock) return;
-    Blocks.insert(BB);
-    visitSuccessor(Blocks, BB, EndBlock);
-  }
-}
-
 void LoopSplitter::doLoopSplit(Function * F, Loop * L0, BasicBlock * SplitBlock) {
   auto OldHeader = L0->getHeader();
-  auto DistBB = SplitBlock->getUniqueSuccessor();
-
-  BasicBlock * EntryBlock = nullptr;
-  for (auto BB : predecessors(OldHeader)) {
-    if (BB != L0->getLoopLatch())
-      EntryBlock = BB;
-  }
+  auto OldEntry = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(0);
+  auto PreLoopBB = getPreLoopBlock(L0);
+  auto PostLoopBB = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(1);
+  auto MidBlock = SplitBlock->getUniqueSuccessor();
 
-  auto OldPreHeader = BasicBlock::Create(F->getContext(), "oldpreheader", F);
-  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 LBT = LoopBodyTraverser(L0);
 
-  auto IndexVar = getLoopIndexVar(L0);
-  auto TripCount = getLoopTripCount(L0);
+  // Collect Blocks in range [OldEntry, MidBlock)
+  SmallVector<BasicBlock *, 4> TopLoopBlocks;
+  LBT.traverse(TopLoopBlocks, OldEntry, MidBlock);
 
-  IRBuilder<> Builder(NewPreHeader);
-  Builder.CreateStore(Builder.getInt32(0), IndexVar);
-  Builder.CreateBr(NewHeader);
+  // Collect Blocks in range [MidBlock, Latch)
+  SmallVector<BasicBlock *, 4> BottomLoopBlocks;
+  LBT.traverse(BottomLoopBlocks, MidBlock, L0->getLoopLatch());
 
-  Builder.SetInsertPoint(NewLatch);
-  auto BI = Builder.CreateBr(NewHeader);
-  addIncrementIndexOp(IndexVar, BI);
+  auto BottomLoop = IRLoop();
+  BottomLoop.extractLoopSkeleton(L0);
 
-  Builder.SetInsertPoint(NewHeader);
-  auto IndexVarVal = Builder.CreateLoad(IndexVar);
-  auto * icmp = Builder.CreateICmpSLT(IndexVarVal, TripCount,
-                                      "loop-predicate");
-  auto NewEntryBB = cast<BranchInst>(OldHeader->getTerminator())->getSuccessor(0);
-  Builder.CreateCondBr(icmp, NewEntryBB, OldPreHeader);
+  auto TopLoop = IRLoop();
+  TopLoop.constructEmptyLoop(getLoopTripCount(L0), BottomLoop.getHeader());
 
-  SplitBlock->getTerminator()->setSuccessor(0, NewLatch);
-  OldHeader->getTerminator()->setSuccessor(0, DistBB);
-  OldHeader->getTerminator()->setSuccessor(1, ExitBlock);
-  setSuccessor(EntryBlock, NewPreHeader);
+  TopLoop.setLoopBlocks(TopLoopBlocks);
+  BottomLoop.setLoopBlocks(BottomLoopBlocks);
 
-  Builder.SetInsertPoint(OldPreHeader);
-  errs() << *IndexVar << "\n";
-  Builder.CreateStore(Builder.getInt32(0), IndexVar);
-  Builder.CreateBr(OldHeader);
+  setSuccessor(PreLoopBB, TopLoop.getPreHeader());
+  setSuccessor(TopLoop.getHeader(), BottomLoop.getHeader(), 1); 
+  setSuccessor(BottomLoop.getHeader(), PostLoopBB, 1);
 
+  /*
   // If FalseBB is terminating instruction, use latch block as target instead.
   SmallVector<BasicBlock *, 4> Returns;
   getReturnBlocks(F, Returns);
@@ -157,7 +144,10 @@ void LoopSplitter::doLoopSplit(Function * F, Loop * L0, BasicBlock * SplitBlock)
       Case.setSuccessor(L0->getLoopLatch());
     }
   }
+  */
+}
 
+/*
   DenseSet<BasicBlock *> Blocks;
   auto TruePathBB = NewHeader->getTerminator()->getSuccessor(0);
   visitSuccessor(Blocks, TruePathBB, NewLatch);
@@ -189,6 +179,7 @@ void LoopSplitter::doLoopSplit(Function * F, Loop * L0, BasicBlock * SplitBlock)
     ++BB;
   }
 }
+*/
 
 bool LoopSplitter::run() {
   // If no loops, we are done.
@@ -200,12 +191,23 @@ bool LoopSplitter::run() {
   ExitBlock = L0->getExitBlock();
   assert (ExitBlock && "Loop must have a single exit block!");
 
-  errs() << "Running prepareForLoopSplit\n";
+  auto ParentLoop = IRLoop();
+  ParentLoop.analyze(L0);
+
   bool changed = prepareForLoopSplit(F, L0, stat);
   if (!changed) return false;
 
+//  F->print(errs());
+
+  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;
+  }
   auto & SplitBB = LoopSplitEdgeBlocks.front();
-  doLoopSplit(F, L0, SplitBB);
+//  doLoopSplit(F, L0, SplitBB);
 
   return true;
 }

src/LoopSplitter.h

@@ -48,7 +48,7 @@ public:
 
   bool run();
   Stats & getStats() { return stat; }
-  void addAdapterBasicBlocks(llvm::Instruction * SP, llvm::Value * Idx);
+  void addAdapterBasicBlocks(llvm::Loop * L0, llvm::Instruction * SP, llvm::Value * Idx);
   void doLoopSplit(llvm::Function * F, llvm::Loop * L0, llvm::BasicBlock * SplitBlock);
 };
 

src/Util.cpp

@@ -37,7 +37,7 @@ namespace tas {
 // Annotation list structure in IR: stored as array of struct
 // ConstantArray : [Size x struct]
 // struct def: [Function  ptr, GlobalVariable ptr, GlobalVariable ptr, i32]
-// @llvm.global.annotations = [N x {i8*, i8*, i8*, i32}]
+// @llvm.global.annotations = [N x {i8*, i8*, *, i32}]
 // N - number of global annotations in a module
 // Struct members details:
 // i8* - Function pointer
@@ -257,8 +257,6 @@ unsigned getTypeSizeInBits(Type * Ty) {
   return Total;
 }
 
-
-
 Value * createArray(Function * F, Type * Ty, unsigned size) {
   // Allocate temporary array
   IRBuilder<> Builder (F->getContext());
@@ -290,7 +288,6 @@ unsigned getGEPIndex(const GetElementPtrInst * GEP) {
   return FieldIdx;
 }
 
-
 StoreInst * findFirstUseInStoreInst(Value * V) {
   // User list is inverted, first user is at the end of the list.
   // Loop below finds first user which is StoreInst.
@@ -474,4 +471,31 @@ AllocaInst * getLoopIndexVar(Loop * L) {
   return Index;
 }
 
+void visitSuccessor(SmallVectorImpl<BasicBlock *> & Blocks, BasicBlock * StartBlock,
+                    BasicBlock * EndBlock) {
+  for (auto * BB : successors(StartBlock)){
+    if (BB == EndBlock) return;
+    Blocks.push_back(BB);
+    visitSuccessor(Blocks, BB, EndBlock);
+  }
+}
+
+Value * getLoopTripCount(Loop * L0) {
+  auto Header = L0->getHeader();
+  auto Cond = cast<BranchInst>(Header->getTerminator())->getCondition();
+  return cast<ICmpInst>(Cond)->getOperand(1);
+}
+
+BasicBlock * getPreLoopBlock(Loop * L) {
+  BasicBlock * PreLoopBB = L->getLoopPreheader();
+  if (!PreLoopBB) {
+    for (auto * BB : predecessors(L->getHeader())) {
+      if (BB == L->getLoopLatch()) continue;
+      if (L->contains(BB)) continue;
+      PreLoopBB = BB;
+    }
+  }
+  return PreLoopBB;
+}
+
 } // namespace tas

src/Util.h

@@ -44,7 +44,6 @@ unsigned getTypeSizeInBits(llvm::Type * Ty);
 
 llvm::SmallVector<llvm::Value *, 4> detectExpPtrVars(llvm::Function * F);
 
-
 llvm::SmallVector<llvm::LoadInst *, 4>
 detectExpPtrUses(llvm::SmallVectorImpl<llvm::Value *> & AnnotatedVars);
 
@@ -105,12 +104,17 @@ auto findFirstUseOfValueInInstType(llvm::Value * V) {
   return LU;
 }
 
+llvm::AllocaInst * getLoopIndexVar(llvm::Loop * L);
 
+const llvm::LoadInst * findEarliestPointerDerefInstruction(llvm::Value * V);
 
-llvm::AllocaInst * getLoopIndexVar(llvm::Loop * L);
+void visitSuccessor(llvm::SmallVectorImpl<llvm::BasicBlock *> & Blocks,
+                    llvm::BasicBlock * CurBlock,
+                    llvm::BasicBlock * EndBlock);
 
+llvm::Value * getLoopTripCount(llvm::Loop * L0);
 
-const llvm::LoadInst * findEarliestPointerDerefInstruction(llvm::Value * V);
+llvm::BasicBlock * getPreLoopBlock(llvm::Loop * L);
 
 } // namespace tas