[OpenMPIRBuilder] Implicitly defined control blocks. NFC.

[Meinersbur]

Summary

Do not explicitly store the BasicBlocks for Preheader, Body and After inside CanonicalLoopInfo, but look the up when needed using their position relative to the other loop control blocks. By definition, instructions inside these are not managed by CanonicalLoopInfo (except terminator for Preheader) hence it makes sense to think of them as connections to the CanonicalLoopInfo instead of part of the CanonicalLoopInfo itself.

In particular for Preheader, it makes using SplitBasicBlock easier since inserting control flow at an InsertPoint may otherwise require updating the CanonicalLoopInfo's Preheader because the branch that jumps to the header is moved to another BasicBlock.

[Meinersbur]

Rebase

[ftynse]

LGTM

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

 /// TODO: Provide mechanisms for exception handling and cancellation points.
 ///
 /// Defined outside OpenMPIRBuilder because nested classes cannot be
 /// forward-declared, e.g. to avoid having to include the entire OMPIRBuilder.h.
 class CanonicalLoopInfo {
   friend class OpenMPIRBuilder;
 
 private:
-  BasicBlock *Preheader = nullptr;
   BasicBlock *Header = nullptr;
   BasicBlock *Cond = nullptr;
-  BasicBlock *Body = nullptr;
   BasicBlock *Latch = nullptr;
   BasicBlock *Exit = nullptr;
-  BasicBlock *After = nullptr;
 
   /// Add the control blocks of this loop to \p BBs.
   ///
   /// This does not include any block from the body, including the one returned
   /// by getBody().
   ///
   /// FIXME: This currently includes the Preheader and After blocks even though
   /// their content is (mostly) not under CanonicalLoopInfo's control.
   /// Re-evaluated whether this makes sense.
   void collectControlBlocks(SmallVectorImpl<BasicBlock *> &BBs);
 
 public:
   /// Returns whether this object currently represents the IR of a loop. If
   /// returning false, it may have been consumed by a loop transformation or not
   /// been intialized. Do not use in this case;
   bool isValid() const { return Header; }
 
   /// The preheader ensures that there is only a single edge entering the loop.
   /// Code that must be execute before any loop iteration can be emitted here,
   /// such as computing the loop trip count and begin lifetime markers. Code in
   /// the preheader is not considered part of the canonical loop.
-  BasicBlock *getPreheader() const {
-    assert(isValid() && "Requires a valid canonical loop");
-    return Preheader;
-  }
+  BasicBlock *getPreheader() const;
 
   /// The header is the entry for each iteration. In the canonical control flow,
   /// it only contains the PHINode for the induction variable.
   BasicBlock *getHeader() const {
     assert(isValid() && "Requires a valid canonical loop");
     return Header;
   }
 
   /// The condition block computes whether there is another loop iteration. If
   /// yes, branches to the body; otherwise to the exit block.
   BasicBlock *getCond() const {
     assert(isValid() && "Requires a valid canonical loop");
     return Cond;
   }
 
   /// The body block is the single entry for a loop iteration and not controlled
   /// by CanonicalLoopInfo. It can contain arbitrary control flow but must
   /// eventually branch to the \p Latch block.
   BasicBlock *getBody() const {
     assert(isValid() && "Requires a valid canonical loop");
-    return Body;
+    return cast<BranchInst>(Cond->getTerminator())->getSuccessor(0);
   }
 
   /// Reaching the latch indicates the end of the loop body code. In the
   /// canonical control flow, it only contains the increment of the induction
   /// variable.
   BasicBlock *getLatch() const {
     assert(isValid() && "Requires a valid canonical loop");
     return Latch;
   }
 
   /// Reaching the exit indicates no more iterations are being executed.
   BasicBlock *getExit() const {
     assert(isValid() && "Requires a valid canonical loop");
     return Exit;
   }
 
   /// The after block is intended for clean-up code such as lifetime end
   /// markers. It is separate from the exit block to ensure, analogous to the
   /// preheader, it having just a single entry edge and being free from PHI
   /// nodes should there be multiple loop exits (such as from break
   /// statements/cancellations).
   BasicBlock *getAfter() const {
     assert(isValid() && "Requires a valid canonical loop");
-    return After;
+    return Exit->getSingleSuccessor();
   }
 
   /// Returns the llvm::Value containing the number of loop iterations. It must
   /// be valid in the preheader and always interpreted as an unsigned integer of
   /// any bit-width.
   Value *getTripCount() const {
     assert(isValid() && "Requires a valid canonical loop");
     Instruction *CmpI = &Cond->front();
public:
   Type *getIndVarType() const {
     assert(isValid() && "Requires a valid canonical loop");
     return getIndVar()->getType();
   }
 
   /// Return the insertion point for user code before the loop.
   OpenMPIRBuilder::InsertPointTy getPreheaderIP() const {
     assert(isValid() && "Requires a valid canonical loop");
+    BasicBlock *Preheader = getPreheader();
     return {Preheader, std::prev(Preheader->end())};
   };
 
   /// Return the insertion point for user code in the body.
   OpenMPIRBuilder::InsertPointTy getBodyIP() const {
     assert(isValid() && "Requires a valid canonical loop");
+    BasicBlock *Body = getBody();
     return {Body, Body->begin()};
   };
 
   /// Return the insertion point for user code after the loop.
   OpenMPIRBuilder::InsertPointTy getAfterIP() const {
     assert(isValid() && "Requires a valid canonical loop");
+    BasicBlock *After = getAfter();
     return {After, After->begin()};
   };
 
   Function *getFunction() const {
     assert(isValid() && "Requires a valid canonical loop");
     return Header->getParent();
   }
 

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

CanonicalLoopInfo *OpenMPIRBuilder::createLoopSkeleton(
 
   Builder.SetInsertPoint(Exit);
   Builder.CreateBr(After);
 
   // Remember and return the canonical control flow.
   LoopInfos.emplace_front();
   CanonicalLoopInfo *CL = &LoopInfos.front();
 
-  CL->Preheader = Preheader;
   CL->Header = Header;
   CL->Cond = Cond;
-  CL->Body = Body;
   CL->Latch = Latch;
   CL->Exit = Exit;
-  CL->After = After;
 
 #ifndef NDEBUG
   CL->assertOK();
 #endif
   return CL;
 }
 
 CanonicalLoopInfo *
 OpenMPIRBuilder::createCanonicalLoop(const LocationDescription &Loc,
                                      LoopBodyGenCallbackTy BodyGenCB,
                                      Value *TripCount, const Twine &Name) {
   BasicBlock *BB = Loc.IP.getBlock();
   BasicBlock *NextBB = BB->getNextNode();
 
   CanonicalLoopInfo *CL = createLoopSkeleton(Loc.DL, TripCount, BB->getParent(),
                                              NextBB, NextBB, Name);
   BasicBlock *After = CL->getAfter();
 
   // If location is not set, don't connect the loop.
   if (updateToLocation(Loc)) {
     // Split the loop at the insertion point: Branch to the preheader and move
     // every following instruction to after the loop (the After BB). Also, the
     // new successor is the loop's after block.
-    Builder.CreateBr(CL->Preheader);
+    Builder.CreateBr(CL->getPreheader());
     After->getInstList().splice(After->begin(), BB->getInstList(),
                                 Builder.GetInsertPoint(), BB->end());
     After->replaceSuccessorsPhiUsesWith(BB, After);
   }
 
   // Emit the body content. We do it after connecting the loop to the CFG to
   // avoid that the callback encounters degenerate BBs.
   BodyGenCB(CL->getBodyIP(), CL->getIndVar());
if (NumLoops == 1)
     return Loops.front();
 
   CanonicalLoopInfo *Outermost = Loops.front();
   CanonicalLoopInfo *Innermost = Loops.back();
   BasicBlock *OrigPreheader = Outermost->getPreheader();
   BasicBlock *OrigAfter = Outermost->getAfter();
   Function *F = OrigPreheader->getParent();
 
+  // Loop control blocks that may become orphaned later.
+  SmallVector<BasicBlock *, 12> OldControlBBs;
+  OldControlBBs.reserve(6 * Loops.size());
+  for (CanonicalLoopInfo *Loop : Loops)
+    Loop->collectControlBlocks(OldControlBBs);
+
   // Setup the IRBuilder for inserting the trip count computation.
   Builder.SetCurrentDebugLocation(DL);
   if (ComputeIP.isSet())
     Builder.restoreIP(ComputeIP);
   else
     Builder.restoreIP(Outermost->getPreheaderIP());
 
   // Derive the collapsed' loop trip count.
OpenMPIRBuilder::collapseLoops(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
   redirectTo(Outermost->getPreheader(), Result->getPreheader(), DL);
   redirectTo(Result->getAfter(), Outermost->getAfter(), DL);
 
   // Replace the input loop indvars with the derived ones.
   for (size_t i = 0; i < NumLoops; ++i)
     Loops[i]->getIndVar()->replaceAllUsesWith(NewIndVars[i]);
 
   // Remove unused parts of the input loops.
-  SmallVector<BasicBlock *, 12> OldControlBBs;
-  OldControlBBs.reserve(6 * Loops.size());
-  for (CanonicalLoopInfo *Loop : Loops)
-    Loop->collectControlBlocks(OldControlBBs);
   removeUnusedBlocksFromParent(OldControlBBs);
 
   for (CanonicalLoopInfo *L : Loops)
     L->invalidate();
 
 #ifndef NDEBUG
   Result->assertOK();
 #endif
OpenMPIRBuilder::tileLoops(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
   assert(NumLoops >= 1 && "At least one loop to tile required");
 
   CanonicalLoopInfo *OutermostLoop = Loops.front();
   CanonicalLoopInfo *InnermostLoop = Loops.back();
   Function *F = OutermostLoop->getBody()->getParent();
   BasicBlock *InnerEnter = InnermostLoop->getBody();
   BasicBlock *InnerLatch = InnermostLoop->getLatch();
 
+  // Loop control blocks that may become orphaned later.
+  SmallVector<BasicBlock *, 12> OldControlBBs;
+  OldControlBBs.reserve(6 * Loops.size());
+  for (CanonicalLoopInfo *Loop : Loops)
+    Loop->collectControlBlocks(OldControlBBs);
+
   // Collect original trip counts and induction variable to be accessible by
   // index. Also, the structure of the original loops is not preserved during
   // the construction of the tiled loops, so do it before we scavenge the BBs of
   // any original CanonicalLoopInfo.
   SmallVector<Value *, 4> OrigTripCounts, OrigIndVars;
   for (CanonicalLoopInfo *L : Loops) {
     assert(L->isValid() && "All input loops must be valid canonical loops");
     OrigTripCounts.push_back(L->getTripCount());
for (int i = 0; i < NumLoops; ++i) {
     Value *Scale =
         Builder.CreateMul(Size, FloorLoop->getIndVar(), {}, /*HasNUW=*/true);
     Value *Shift =
         Builder.CreateAdd(Scale, TileLoop->getIndVar(), {}, /*HasNUW=*/true);
     OrigIndVar->replaceAllUsesWith(Shift);
   }
 
   // Remove unused parts of the original loops.
-  SmallVector<BasicBlock *, 12> OldControlBBs;
-  OldControlBBs.reserve(6 * Loops.size());
-  for (CanonicalLoopInfo *Loop : Loops)
-    Loop->collectControlBlocks(OldControlBBs);
   removeUnusedBlocksFromParent(OldControlBBs);
 
   for (CanonicalLoopInfo *L : Loops)
     L->invalidate();
 
 #ifndef NDEBUG
   for (CanonicalLoopInfo *GenL : Result)
     GenL->assertOK();

 
 void CanonicalLoopInfo::collectControlBlocks(
     SmallVectorImpl<BasicBlock *> &BBs) {
   // We only count those BBs as control block for which we do not need to
   // reverse the CFG, i.e. not the loop body which can contain arbitrary control
   // flow. For consistency, this also means we do not add the Body block, which
   // is just the entry to the body code.
   BBs.reserve(BBs.size() + 6);
-  BBs.append({Preheader, Header, Cond, Latch, Exit, After});
+  BBs.append({getPreheader(), Header, Cond, Latch, Exit, getAfter()});
+}
+
+BasicBlock *CanonicalLoopInfo::getPreheader() const {
+  assert(isValid() && "Requires a valid canonical loop");
+  for (BasicBlock *Pred : predecessors(Header)) {
+    if (Pred != Latch)
+      return Pred;
+  }
+  llvm_unreachable("Missing preheader");
 }
 
 void CanonicalLoopInfo::assertOK() const {
 #ifndef NDEBUG
   // No constraints if this object currently does not describe a loop.
   if (!isValid())
     return;
 
+  BasicBlock *Preheader = getPreheader();
+  BasicBlock *Body = getBody();
+  BasicBlock *After = getAfter();
+
   // Verify standard control-flow we use for OpenMP loops.
   assert(Preheader);
   assert(isa<BranchInst>(Preheader->getTerminator()) &&
          "Preheader must terminate with unconditional branch");
   assert(Preheader->getSingleSuccessor() == Header &&
          "Preheader must jump to header");
 
   assert(Header);
#ifndef NDEBUG
   assert(CmpI->getOperand(0) == IndVar &&
          "Exit condition must compare the induction variable");
   assert(CmpI->getOperand(1) == TripCount &&
          "Exit condition must compare with the trip count");
 #endif
 }
 
 void CanonicalLoopInfo::invalidate() {
-  Preheader = nullptr;
   Header = nullptr;
   Cond = nullptr;
-  Body = nullptr;
   Latch = nullptr;
   Exit = nullptr;
-  After = nullptr;
 }