LoopOptTutorial: Basic loop optimization template [STEP3]

llvm/include/llvm/Transforms/Scalar/LoopOptTutorial.h

@@ -17,6 +17,7 @@
 #define LLVM_TRANSFORMS_SCALAR_LOOPOPTTUTORIAL_H
 
 #include "llvm/Analysis/LoopAnalysisManager.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Transforms/Utils/ValueMapper.h"
 
@@ -25,18 +26,52 @@ namespace llvm {
 class Loop;
 class LPMUpdater;
 
+/// This class splits the innermost loop in a loop nest in the middle.
 class LoopSplit {
 public:
-  LoopSplit(LoopInfo &LI) : LI(LI) {}
+  LoopSplit(LoopInfo &LI, ScalarEvolution &SE, DominatorTree &DT)
+      : LI(LI), SE(SE), DT(DT) {}
 
+  // Execute the transformation on the loop nest rooted by \p L.
   bool run(Loop &L) const;
 
 private:
-  LoopInfo &LI;
-
   /// Determines if \p L is a candidate for splitting
   bool isCandidate(const Loop &L) const;
 
+  /// Split the given loop in the middle by creating a new loop that traverse
+  /// the first half of the original iteration space and adjusting the loop
+  /// bounds of \p L to traverse the remaining half.
+  /// Note: \p L is expected to be the innermost loop in a loop nest or a top
+  /// level loop.
+  bool splitLoopInHalf(Loop &L) const;
+
+  /// Clone loop \p L and insert the cloned loop before the basic block \p
+  /// InsertBefore, \p Pred is the predecessor of \p L.
+  /// Note: \p L is expected to be the innermost loop in a loop nest or a top
+  /// level loop.
+  Loop *cloneLoop(Loop &L, BasicBlock &InsertBefore, BasicBlock &Pred) const;
+
+  /// Compute the point where to split the loop \p L. Return the instruction
+  /// calculating the split point.
+  Instruction *computeSplitPoint(const Loop &L,
+                                 Instruction *InsertBefore) const;
+
+  /// Get the latch comparison instruction of loop \p L.
+  ICmpInst *getLatchCmpInst(const Loop &L) const;
+
+  /// Update the dominator tree after cloning the loop.
+  void updateDominatorTree(const Loop &OrigLoop, const Loop &ClonedLoop,
+                           BasicBlock &InsertBefore, BasicBlock &Pred,
+                           ValueToValueMapTy &VMap) const;
+
+  // Dump the LLVM IR for function \p F.
+  void dumpFunction(const StringRef Msg, const Function &F) const;
+
+private:
+  LoopInfo &LI;
+  ScalarEvolution &SE;
+  DominatorTree &DT;
 };
 
 class LoopOptTutorialPass : public PassInfoMixin<LoopOptTutorialPass> {

llvm/lib/Transforms/Scalar/LoopOptTutorial.cpp

@@ -14,17 +14,34 @@
 //===----------------------------------------------------------------------===
 
 #include "llvm/Transforms/Scalar/LoopOptTutorial.h"
+#include "llvm/Analysis/IVDescriptors.h"
 #include "llvm/Analysis/LoopPass.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ValueMapper.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loop-opt-tutorial"
+static const char *VerboseDebug = DEBUG_TYPE "-verbose";
 
-bool LoopSplit::run(Loop &L) const {
+/// Clones a loop \p OrigLoop.  Returns the loop and the blocks in \p
+/// Blocks.
+/// Updates LoopInfo assuming the loop is dominated by block \p LoopDomBB.
+/// Insert the new blocks before block specified in \p Before.
+static Loop *myCloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
+                                      Loop *OrigLoop, ValueToValueMapTy &VMap,
+                                      const Twine &NameSuffix, LoopInfo *LI,
+                                      SmallVectorImpl<BasicBlock *> &Blocks);
+
+//===----------------------------------------------------------------------===//
+// LoopSplit implementation
+//
 
+bool LoopSplit::run(Loop &L) const {
   LLVM_DEBUG(dbgs() << "Entering " << __func__ << "\n");
 
   if (isCandidate(L))
@@ -33,8 +50,7 @@ bool LoopSplit::run(Loop &L) const {
   else
     LLVM_DEBUG(dbgs() << "Loop " << L.getName()
                       << " is not a candidate for splitting.\n");
-
-  return false;
+  return splitLoopInHalf(L);
 }
 
 bool LoopSplit::isCandidate(const Loop &L) const {
@@ -63,17 +79,187 @@ bool LoopSplit::isCandidate(const Loop &L) const {
   return true;
 }
 
+bool LoopSplit::splitLoopInHalf(Loop &L) const {
+  assert(L.isLoopSimplifyForm() && "Expecting a loop in simplify form");
+  assert(L.isSafeToClone() && "Loop is not safe to be cloned");
+  assert(L.getSubLoops().empty() && "Expecting a innermost loop");
+
+  const Function &F = *L.getHeader()->getParent();
+  LLVM_DEBUG(dbgs() << "Splitting loop " << L.getName() << "\n");
+
+  // Generate the code that computes the split point.
+  Instruction *Split =
+      computeSplitPoint(L, L.getLoopPreheader()->getTerminator());
+
+  // Split the loop preheader to create an insertion point for the cloned loop.
+  BasicBlock *Preheader = L.getLoopPreheader();
+  BasicBlock *Pred = Preheader;
+  BasicBlock *InsertBefore =
+      SplitBlock(Preheader, Preheader->getTerminator(), &DT, &LI);
+  DEBUG_WITH_TYPE(VerboseDebug,
+                  dumpFunction("After splitting preheader:\n", F););
+
+  // Clone the original loop, and insert the clone before the original loop.
+  Loop *ClonedLoop = cloneLoop(L, *InsertBefore, *Pred);
+
+  // Modify the upper bound of the cloned loop.
+  ICmpInst *LatchCmpInst = getLatchCmpInst(*ClonedLoop);
+  assert(LatchCmpInst && "Unable to find the latch comparison instruction");
+  LatchCmpInst->setOperand(1, Split);
+
+  // Modify the lower bound of the original loop.
+  PHINode *IndVar = L.getInductionVariable(SE);
+  assert(IndVar && "Unable to find the induction variable PHI node");
+  IndVar->setIncomingValueForBlock(L.getLoopPreheader(), Split);
+
+  DEBUG_WITH_TYPE(VerboseDebug,
+                  dumpFunction("After splitting the loop:\n", F););
+
+  return true;
+}
+
+Loop *LoopSplit::cloneLoop(Loop &L, BasicBlock &InsertBefore,
+                           BasicBlock &Pred) const {
+  // Clone the original loop, insert the clone before the "InsertBefore" BB.
+  Function &F = *L.getHeader()->getParent();
+  SmallVector<BasicBlock *, 4> ClonedLoopBlocks;
+  ValueToValueMapTy VMap;
+
+  // Same as cloneLoopWithPreheader but does not update the dominator tree.
+  // Use for education purposes only, use cloneLoopWithPreheader in production
+  // code.
+  Loop *NewLoop = myCloneLoopWithPreheader(&InsertBefore, &Pred, &L, VMap, "",
+                                           &LI, ClonedLoopBlocks);
+
+  assert(NewLoop && "Run ot of memory");
+  DEBUG_WITH_TYPE(VerboseDebug,
+                  dbgs() << "Create new loop: " << NewLoop->getName() << "\n";
+                  dumpFunction("After cloning loop:\n", F););
+
+  // Update instructions referencing the original loop basic blocks to
+  // reference the corresponding block in the cloned loop.
+  VMap[L.getExitBlock()] = &InsertBefore;
+  remapInstructionsInBlocks(ClonedLoopBlocks, VMap);
+  DEBUG_WITH_TYPE(VerboseDebug,
+                  dumpFunction("After instruction remapping:\n", F););
+
+  // Make the predecessor of original loop jump to the cloned loop.
+  Pred.getTerminator()->replaceUsesOfWith(&InsertBefore,
+                                          NewLoop->getLoopPreheader());
+
+  return NewLoop;
+}
+
+Instruction *LoopSplit::computeSplitPoint(const Loop &L,
+                                          Instruction *InsertBefore) const {
+  Optional<Loop::LoopBounds> Bounds = L.getBounds(SE);
+  assert(Bounds.hasValue() && "Unable to retrieve the loop bounds");
+
+  Value &IVInitialVal = Bounds->getInitialIVValue();
+  Value &IVFinalVal = Bounds->getFinalIVValue();
+  auto *Sub = BinaryOperator::Create(Instruction::Sub, &IVFinalVal,
+                                     &IVInitialVal, "", InsertBefore);
+
+  return BinaryOperator::Create(Instruction::UDiv, Sub,
+                                ConstantInt::get(IVFinalVal.getType(), 2), "",
+                                InsertBefore);
+}
+
+ICmpInst *LoopSplit::getLatchCmpInst(const Loop &L) const {
+  if (BasicBlock *Latch = L.getLoopLatch())
+    if (BranchInst *BI = dyn_cast_or_null<BranchInst>(Latch->getTerminator()))
+      if (BI->isConditional())
+        return dyn_cast<ICmpInst>(BI->getCondition());
+
+  return nullptr;
+}
+
+void LoopSplit::dumpFunction(const StringRef Msg, const Function &F) const {
+  dbgs() << Msg;
+  F.dump();
+}
+
+/// Clones a loop \p OrigLoop.  Returns the loop and the blocks in \p
+/// Blocks.
+/// Updates LoopInfo assuming the loop is dominated by block \p LoopDomBB.
+/// Insert the new blocks before block specified in \p Before.
+static Loop *myCloneLoopWithPreheader(BasicBlock *Before, BasicBlock *LoopDomBB,
+                                      Loop *OrigLoop, ValueToValueMapTy &VMap,
+                                      const Twine &NameSuffix, LoopInfo *LI,
+                                      SmallVectorImpl<BasicBlock *> &Blocks) {
+  Function *F = OrigLoop->getHeader()->getParent();
+  Loop *ParentLoop = OrigLoop->getParentLoop();
+  DenseMap<Loop *, Loop *> LMap;
+
+  Loop *NewLoop = LI->AllocateLoop();
+  LMap[OrigLoop] = NewLoop;
+  if (ParentLoop)
+    ParentLoop->addChildLoop(NewLoop);
+  else
+    LI->addTopLevelLoop(NewLoop);
+
+  BasicBlock *OrigPH = OrigLoop->getLoopPreheader();
+  assert(OrigPH && "No preheader");
+  BasicBlock *NewPH = CloneBasicBlock(OrigPH, VMap, NameSuffix, F);
+  // To rename the loop PHIs.
+  VMap[OrigPH] = NewPH;
+  Blocks.push_back(NewPH);
+
+  // Update LoopInfo.
+  if (ParentLoop)
+    ParentLoop->addBasicBlockToLoop(NewPH, *LI);
+
+  for (Loop *CurLoop : OrigLoop->getLoopsInPreorder()) {
+    Loop *&NewLoop = LMap[CurLoop];
+    if (!NewLoop) {
+      NewLoop = LI->AllocateLoop();
+
+      // Establish the parent/child relationship.
+      Loop *OrigParent = CurLoop->getParentLoop();
+      assert(OrigParent && "Could not find the original parent loop");
+      Loop *NewParentLoop = LMap[OrigParent];
+      assert(NewParentLoop && "Could not find the new parent loop");
+
+      NewParentLoop->addChildLoop(NewLoop);
+    }
+  }
+
+  for (BasicBlock *BB : OrigLoop->getBlocks()) {
+    Loop *CurLoop = LI->getLoopFor(BB);
+    Loop *&NewLoop = LMap[CurLoop];
+    assert(NewLoop && "Expecting new loop to be allocated");
+
+    BasicBlock *NewBB = CloneBasicBlock(BB, VMap, NameSuffix, F);
+    VMap[BB] = NewBB;
+
+    // Update LoopInfo.
+    NewLoop->addBasicBlockToLoop(NewBB, *LI);
+    if (BB == CurLoop->getHeader())
+      NewLoop->moveToHeader(NewBB);
+
+    Blocks.push_back(NewBB);
+  }
+
+  // Move them physically from the end of the block list.
+  F->getBasicBlockList().splice(Before->getIterator(), F->getBasicBlockList(),
+                                NewPH);
+  F->getBasicBlockList().splice(Before->getIterator(), F->getBasicBlockList(),
+                                NewLoop->getHeader()->getIterator(), F->end());
+
+  return NewLoop;
+}
+
+//===----------------------------------------------------------------------===//
+// LoopOptTutorialPass implementation
+//
+
 PreservedAnalyses LoopOptTutorialPass::run(Loop &L, LoopAnalysisManager &LAM,
                                            LoopStandardAnalysisResults &AR,
                                            LPMUpdater &U) {
-  bool Changed = false;
-
   LLVM_DEBUG(dbgs() << "Entering LoopOptTutorialPass::run\n");
   LLVM_DEBUG(dbgs() << "Loop: "; L.dump(); dbgs() << "\n");
 
-  LoopSplit LS(AR.LI);
-
-  Changed = LS.run(L);
+  bool Changed = LoopSplit(AR.LI, AR.SE, AR.DT).run(L);
 
   if (!Changed)
     return PreservedAnalyses::all();