[CodeGen][MachineLoop] Fix getLoopID

llvm/lib/CodeGen/MachineLoopInfo.cpp

@@ -181,48 +181,32 @@ MachineLoopInfo::findLoopPreheader(MachineLoop *L, bool SpeculativePreheader,
 
 MDNode *MachineLoop::getLoopID() const {
   MDNode *LoopID = nullptr;
-  if (const auto *MBB = findLoopControlBlock()) {
-    // If there is a single latch block, then the metadata
-    // node is attached to its terminating instruction.
+
+  // Go through the latch blocks and check the terminator for the metadata
+  SmallVector<MachineBasicBlock *, 4> LatchesBlocks;
+  getLoopLatches(LatchesBlocks);
+  for (const auto *MBB : LatchesBlocks) {
     const auto *BB = MBB->getBasicBlock();
     if (!BB)
       return nullptr;
-    if (const auto *TI = BB->getTerminator())
-      LoopID = TI->getMetadata(LLVMContext::MD_loop);
-  } else if (const auto *MBB = getHeader()) {
-    // There seem to be multiple latch blocks, so we have to
-    // visit all predecessors of the loop header and check
-    // their terminating instructions for the metadata.
-    if (const auto *Header = MBB->getBasicBlock()) {
-      // Walk over all blocks in the loop.
-      for (const auto *MBB : this->blocks()) {
-        const auto *BB = MBB->getBasicBlock();
-        if (!BB)
-          return nullptr;
-        const auto *TI = BB->getTerminator();
-        if (!TI)
-          return nullptr;
-        MDNode *MD = nullptr;
-        // Check if this terminating instruction jumps to the loop header.
-        for (const auto *Succ : successors(TI)) {
-          if (Succ == Header) {
-            // This is a jump to the header - gather the metadata from it.
-            MD = TI->getMetadata(LLVMContext::MD_loop);
-            break;
-          }
-        }
-        if (!MD)
-          continue;
-        if (!LoopID)
-          LoopID = MD;
-        else if (MD != LoopID)
-          return nullptr;
-      }
-    }
+    const auto *TI = BB->getTerminator();
+    if (!TI)
+      return nullptr;
+    MDNode *MD = TI->getMetadata(LLVMContext::MD_loop);
+
+    if (!MD)
+      return nullptr;
+
+    if (!LoopID)
+      LoopID = MD;
+    else if (MD != LoopID)
+      return nullptr;
   }
-  if (LoopID &&
-      (LoopID->getNumOperands() == 0 || LoopID->getOperand(0) != LoopID))
-    LoopID = nullptr;
+
+  if (!LoopID || LoopID->getNumOperands() == 0 ||
+      LoopID->getOperand(0) != LoopID)
+    return nullptr;
+
   return LoopID;
 }
 

llvm/test/CodeGen/X86/code-align-loops.ll

@@ -177,6 +177,48 @@ exit:                               ; preds = %bb2, %bb3, %bb4
   ret void
 }
 
+; test5 is to check if .p2align can be correctly set on loops with a single
+; latch that's not the exiting block.
+; The test IR is generated from below simple C file:
+; $ clang -O0 -S -emit-llvm loop.c
+; $ cat loop.c
+; int test5(int n) {
+;     int i = 0;
+;     [[clang::code_align(64)]]
+;     while (i < n) {
+;         i++;
+;     }
+; }
+; CHECK-LABEL: test5:
+; ALIGN: .p2align 6
+; ALIGN-NEXT: .LBB4_1: # %while.cond
+define dso_local i32 @test5(i32 %n) #0 {
+entry:
+  %retval = alloca i32, align 4
+  %n.addr = alloca i32, align 4
+  %i = alloca i32, align 4
+  store i32 %n, ptr %n.addr, align 4
+  store i32 0, ptr %i, align 4
+  br label %while.cond
+
+while.cond:                                       ; preds = %while.body, %entry
+  %0 = load i32, ptr %i, align 4
+  %1 = load i32, ptr %n.addr, align 4
+  %cmp = icmp slt i32 %0, %1
+  br i1 %cmp, label %while.body, label %while.end
+
+while.body:                                       ; preds = %while.cond
+  %2 = load i32, ptr %i, align 4
+  %inc = add nsw i32 %2, 1
+  store i32 %inc, ptr %i, align 4
+  br label %while.cond, !llvm.loop !0
+
+while.end:                                        ; preds = %while.cond
+  %3 = load i32, ptr %retval, align 4
+  ret i32 %3
+}
+
+
 declare void @bar()
 declare void @var()