[LoopInfo] Don't bail out early in getLoopID

[loladiro]

Summary

As far as I understand getLoopID needs to look at the intersection of

1. BasicBlocks in the loop and
2. Predecessors of the loop header

and check each of those terminators for a !loop metadata node. However,
what getLoopID is currently doing is iterate through the basic blocks
and require that every one of them have an edge back to the loop header
with the appropriate metadata. This seems incorrect and causes situations
in which a LoopID set with setLoopID cannot be subsequently retrieved with
getLoopID.

[fhahn]

Could you add a test for that? e.g in unittests/Analysis/LoopInfoTest.cpp

lib/Analysis/LoopInfo.cpp

223–225

Shouldn't getLoopLatches return exactly the nodes we are interested in (predecessors of the header in the loop) [1]? IIUC using it would simplify the code even further.

[1] https://llvm.org/doxygen/classllvm_1_1LoopBase.html#abbefed3df8ce43cf3f0c2f7fd004f6a4

[loladiro]

Updated to use getLoopLatches and add a test.

[fhahn]

Thanks for adding the test case! Just a few small comments left.

lib/Analysis/LoopInfo.cpp

221

Please reflow this comment

224

Usually we don't use this->

256

Usually we don't use this->.

257

No braces for one stmt blocks

unittests/Analysis/LoopInfoTest.cpp

108

not: capitialize first letter of variable names?

110

We can get all the info using LoopInfo I think. An assert that we have 2 latches should be enough?

133

Could you extend this test case to also check the cases where the latches contain differing ids and where one latch does not contain a loop id, while the other does? It should be possible to just change the attached loop ids I think.

[vtjnash]

@loladiro bump. It sounds like this was approved, but it doesn't apply cleanly now.

[loladiro]

Looks like this got fixed by @jdoerfert in rL341926 with a near identical patch and test. No rebase required.

[loladiro]

Actually, looks like the change to setLoopID is still relevant. I'll submit a new revision with just that change.

lib/Analysis/LoopInfo.cpp

 
 MDNode *Loop::getLoopID() const {
   MDNode *LoopID = nullptr;
   if (BasicBlock *Latch = getLoopLatch()) {
     LoopID = Latch->getTerminator()->getMetadata(LLVMContext::MD_loop);
   } else {
     assert(!getLoopLatch() &&
            "The loop should have no single latch at this point");
-    // Go through each predecessor of the loop header and check the
+    // Go through each loop latch and check the
     // terminator for the metadata.
-    BasicBlock *H = getHeader();
-    for (BasicBlock *BB : this->blocks()) {
+    SmallVector<BasicBlock *, 4> LoopLatches;
+    this->getLoopLatches(LoopLatches);
+    for (BasicBlock *BB : LoopLatches) {
       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;
-        }
-      }
+      MDNode *MD = TI->getMetadata(LLVMContext::MD_loop);
       if (!MD)
         return nullptr;
 
       if (!LoopID)
         LoopID = MD;
       else if (MD != LoopID)
         return nullptr;
     }
void Loop::setLoopID(MDNode *LoopID) const {
 
   if (BasicBlock *Latch = getLoopLatch()) {
     Latch->getTerminator()->setMetadata(LLVMContext::MD_loop, LoopID);
     return;
   }
 
   assert(!getLoopLatch() &&
          "The loop should have no single latch at this point");
-  BasicBlock *H = getHeader();
-  for (BasicBlock *BB : this->blocks()) {
-    TerminatorInst *TI = BB->getTerminator();
-    for (BasicBlock *Successor : TI->successors()) {
-      if (Successor == H)
-        TI->setMetadata(LLVMContext::MD_loop, LoopID);
-    }
+  SmallVector<BasicBlock *, 4> LoopLatches;
+  this->getLoopLatches(LoopLatches);
+  for (BasicBlock *BB : LoopLatches) {
+    BB->getTerminator()->setMetadata(LLVMContext::MD_loop, LoopID);
   }
 }
 
 void Loop::setLoopAlreadyUnrolled() {
   MDNode *LoopID = getLoopID();
   // First remove any existing loop unrolling metadata.
   SmallVector<Metadata *, 4> MDs;
   // Reserve first location for self reference to the LoopID metadata node.

unittests/Analysis/LoopInfoTest.cpp

runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
     }
 
     // We must have successfully found and set the loop id in the
     // only latch the loop has.
     EXPECT_TRUE(loopIDFoundAndSet);
   });
 }
 
+TEST(LoopInfoTest, LoopWithMultipleLatches) {
+  const char *ModuleStr =
+      "define void @foo(i1 %cond, i32 %n) {\n"
+      "entry:\n"
+      "  br i1 undef, label %header, label %for.end\n"
+      "header:\n"
+      "  %i.0 = phi i32 [ 0, %entry ], [ %inc, %latch1 ], [ %inc, %latch2 ]\n"
+      "  %cmp = icmp slt i32 %i.0, %n\n"
+      "  br i1 %cmp, label %for.inc, label %for.end\n"
+      "for.inc:\n"
+      "  %inc = add nsw i32 %i.0, 1\n"
+      "  br i1 %cond, label %latch1, label %latch2\n"
+      "latch1:\n"
+      "  br label %header\n"
+      "latch2:\n"
+      "  br label %header\n"
+      "for.end:\n"
+      "  ret void\n"
+      "}\n";
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleStr);
+
+  runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
+    BasicBlock *header = nullptr, *latch1 = nullptr, *latch2 = nullptr;
+    // Iterate over the function and find our interersting blocks
+    for (BasicBlock &BB : F) {
+      if (BB.getName() == "header")
+        header = &BB;
+      else if (BB.getName() == "latch1")
+        latch1 = &BB;
+      else if (BB.getName() == "latch2")
+        latch2 = &BB;
+    }
+    assert(header && latch1 && latch2);
+    Loop *L = LI.getLoopFor(header);
+    EXPECT_TRUE(LI.isLoopHeader(header));
+    // We start out with no loop ID
+    EXPECT_EQ(L->getLoopID(), nullptr);
+    // Now set a loop ID - This will add the metadata to each
+    // loop latch.
+    MDNode *LoopID = MDNode::get(Context, ArrayRef<Metadata *>(NULL));
+    LoopID->replaceOperandWith(0, LoopID);
+    L->setLoopID(LoopID);
+    // Make sure this took effect properly and can be retrieved with
+    EXPECT_EQ(latch1->getTerminator()->getMetadata(LLVMContext::MD_loop),
+              LoopID);
+    EXPECT_EQ(latch2->getTerminator()->getMetadata(LLVMContext::MD_loop),
+              LoopID);
+    EXPECT_EQ(L->getLoopID(), LoopID);
+  });
+}
+
 TEST(LoopInfoTest, PreorderTraversals) {
   const char *ModuleStr = "define void @f() {\n"
                           "entry:\n"
                           "  br label %loop.0\n"
                           "loop.0:\n"
                           "  br i1 undef, label %loop.0.0, label %loop.1\n"
                           "loop.0.0:\n"
                           "  br i1 undef, label %loop.0.0, label %loop.0.1\n"