Search for all predecessors of loop header while getting loop metadata

[hiraditya]

Summary

The current implementation of getLoopID seems incorrect as it returns nullptr only after looking into the successors of first block in the loop.
In the setLoopID, there is no need to iterate over all the basic blocks of loop as loop metadata is only associated with predecessor (s) of loop header.

[hiraditya]

Ping

[hiraditya]

Ping

[mzolotukhin]

llvm/lib/Analysis/LoopInfo.cpp

216–224

This comment is stale now.

249

Are you sure that we cannot have multiple branches? I think it's true for well-formed loops, but might be incorrect in a general case.

[hiraditya]

llvm/lib/Analysis/LoopInfo.cpp

249

I'll remove this assert. THanks,

[hiraditya]

Addressed Michael's comments.

[mzolotukhin]

llvm/lib/Analysis/LoopInfo.cpp

214–215

AFAIU, this->contains would be as expensive, as iterating through all blocks (as we did in the original version). Moreover, now we'll do it several times while in the original code we did that only once. Are you sure it's worth it?

The issue you mentioned - getLoopID returns nullptr only after looking into the successors of the first block in the loop - still looks relevant though, so some fix is indeed needed.

244–245

Same question about the cost of this->contains versus original iterating over all blocks.

[hiraditya]

llvm/lib/Analysis/LoopInfo.cpp

214–215

AFAIU, this->contains would be as expensive, as iterating through all blocks

Only for loops with few basic blocks, because DenseBlockSet is a SmallPtrSet, whereas in the previous implementation it would iterate through all the basic blocks and all the successors of each BB in the loop.

llvm/lib/Analysis/LoopInfo.cpp

 
 MDNode *Loop::getLoopID() const {
   MDNode *LoopID = nullptr;
   if (isLoopSimplifyForm()) {
     LoopID = getLoopLatch()->getTerminator()->getMetadata(LLVMContext::MD_loop);
   } else {
     // Go through each predecessor of the loop header and check the
     // terminator for the metadata.
-    BasicBlock *H = getHeader();
-    for (BasicBlock *BB : this->blocks()) {
-      TerminatorInst *TI = BB->getTerminator();
-      MDNode *MD = nullptr;
-
-      // Check if this terminator branches to the loop header.
-      for (BasicBlock *Successor : TI->successors()) {
-        if (Successor == H) {
-          MD = TI->getMetadata(LLVMContext::MD_loop);
-          break;
-        }
-      }
-      if (!MD)
-        return nullptr;
+    for (BasicBlock *BB : predecessors(getHeader())) {
+      if (!this->contains(BB))
+        continue;
 
+      TerminatorInst *TI = BB->getTerminator();
+      if (MDNode *MD = TI->getMetadata(LLVMContext::MD_loop)) {
         if (!LoopID)
           LoopID = MD;
-      else if (MD != LoopID)
+        else if (MD != LoopID) // Multiple MD_loop found => corrupt metadata.
           return nullptr;
       }
     }
+  }
   if (!LoopID || LoopID->getNumOperands() == 0 ||
       LoopID->getOperand(0) != LoopID)
     return nullptr;
   return LoopID;
 }
 
 void Loop::setLoopID(MDNode *LoopID) const {
   assert(LoopID && "Loop ID should not be null");
   assert(LoopID->getNumOperands() > 0 && "Loop ID needs at least one operand");
   assert(LoopID->getOperand(0) == LoopID && "Loop ID should refer to itself");
 
   if (isLoopSimplifyForm()) {
     getLoopLatch()->getTerminator()->setMetadata(LLVMContext::MD_loop, LoopID);
     return;
   }
 
-  BasicBlock *H = getHeader();
-  for (BasicBlock *BB : this->blocks()) {
+  // Set metadata on all the back-edges.
+  for (BasicBlock *BB : predecessors(getHeader())) {
+    if (!this->contains(BB))
+      continue;
+
     TerminatorInst *TI = BB->getTerminator();
-    for (BasicBlock *Successor : TI->successors()) {
-      if (Successor == H)
-        TI->setMetadata(LLVMContext::MD_loop, LoopID);
-    }
-  }
-}
 
 bool Loop::isAnnotatedParallel() const {
   MDNode *DesiredLoopIdMetadata = getLoopID();
 
   if (!DesiredLoopIdMetadata)
       return false;
 
   // The loop branch contains the parallel loop metadata. In order to ensure