[MachineLoopInfo] Fix getLoopID to handle multi latches.

This patch also fixed CodegenPrepare to reserve loop metadata
when merging blocks.

This fixes issue #102632

Author: FreddyLeaf

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -1183,6 +1183,12 @@ void CodeGenPrepare::eliminateMostlyEmptyBlock(BasicBlock *BB) {
     }
   }
 
+  // Reserve loop Metadata.
+  if (BI->hasMetadata(LLVMContext::MD_loop)) {
+    for (auto *Pred : predecessors(BB))
+      Pred->getTerminator()->copyMetadata(*BI, LLVMContext::MD_loop);
+  }
+
   // The PHIs are now updated, change everything that refers to BB to use
   // DestBB and remove BB.
   BB->replaceAllUsesWith(DestBB);

llvm/lib/CodeGen/MachineLoopInfo.cpp

@@ -211,8 +211,6 @@ MDNode *MachineLoop::getLoopID() const {
             break;
           }
         }
-        if (!MD)
-          return nullptr;
         if (!LoopID)
           LoopID = MD;
         else if (MD != LoopID)

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

@@ -96,6 +96,56 @@ for.body5:                                        ; preds = %for.body, %for.body
   br i1 %exitcond16.not, label %for.cond.cleanup4, label %for.body5, !llvm.loop !2
 }
 
+; test3 is to check if .p2align can be correctly set on loops with multi latches.
+; The test IR is generated from below simple C file:
+; $ clang -O0 -S -emit-llvm loop.c
+; $ cat loop.c
+; int test3() {
+;     int i = 0;
+;     [[clang::code_align(32)]]
+;     while (i < 10) {
+;         if (i % 2) {
+;             continue;
+;         }
+;         i++;
+;     }
+; }
+; CHECK-LABEL: test3_multilatch:
+; ALIGN: .p2align 6, 0x90
+; ALIGN-NEXT: .LBB2_1: # %while.cond
+define dso_local i32 @test3_multilatch() #0 {
+entry:
+  %retval = alloca i32, align 4
+  %i = alloca i32, align 4
+  store i32 0, ptr %retval, align 4
+  store i32 0, ptr %i, align 4
+  br label %while.cond
+
+while.cond:                                       ; preds = %if.end, %if.then, %entry
+  %0 = load i32, ptr %i, align 4
+  %cmp = icmp slt i32 %0, 10
+  br i1 %cmp, label %while.body, label %while.end
+
+while.body:                                       ; preds = %while.cond
+  %1 = load i32, ptr %i, align 4
+  %rem = srem i32 %1, 2
+  %tobool = icmp ne i32 %rem, 0
+  br i1 %tobool, label %if.then, label %if.end
+
+if.then:                                          ; preds = %while.body
+  br label %while.cond, !llvm.loop !0
+
+if.end:                                           ; preds = %while.body
+  %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()