[llvm] [MachineLoopInfo] Fix getLoopID to handle multi latches. (PR #106195)

[Freddy Ye via llvm-commits]

https://github.com/FreddyLeaf updated https://github.com/llvm/llvm-project/pull/106195

>From b52dc77e9d8935870d12b2e984e0c53b5cac5517 Mon Sep 17 00:00:00 2001
From: Freddy Ye <freddy.ye at intel.com>
Date: Wed, 21 Aug 2024 11:12:57 +0800
Subject: [PATCH 1/2] [MachineLoopInfo] Fix getLoopID to handle multi latches.

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

This fixes issue #102632
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp       | 6 ++++++
 llvm/lib/CodeGen/MachineLoopInfo.cpp      | 2 --
 llvm/test/CodeGen/X86/code-align-loops.ll | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index bf48c1fdab0ff0..c63a265fe27789 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/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);
diff --git a/llvm/lib/CodeGen/MachineLoopInfo.cpp b/llvm/lib/CodeGen/MachineLoopInfo.cpp
index 88ba10fbe29a91..b82b41fac662d7 100644
--- a/llvm/lib/CodeGen/MachineLoopInfo.cpp
+++ b/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)
diff --git a/llvm/test/CodeGen/X86/code-align-loops.ll b/llvm/test/CodeGen/X86/code-align-loops.ll
index 616478993a51da..40ad3b4c8d600e 100644
--- a/llvm/test/CodeGen/X86/code-align-loops.ll
+++ b/llvm/test/CodeGen/X86/code-align-loops.ll
@@ -111,7 +111,7 @@ for.body5:                                        ; preds = %for.body, %for.body
 ;     }
 ; }
 ; CHECK-LABEL: test3_multilatch:
-; ALIGN: .p2align 4, 0x90
+; ALIGN: .p2align 6, 0x90
 ; ALIGN-NEXT: .LBB2_1: # %while.cond
 define dso_local i32 @test3_multilatch() #0 {
 entry:

>From f58d536ebc343af5ec663d0d2fa3add0e0cb86de Mon Sep 17 00:00:00 2001
From: Freddy Ye <freddy.ye at intel.com>
Date: Wed, 28 Aug 2024 15:39:37 +0800
Subject: [PATCH 2/2] address comments.

---
 llvm/lib/CodeGen/CodeGenPrepare.cpp       |  2 +-
 llvm/lib/CodeGen/MachineLoopInfo.cpp      |  2 ++
 llvm/test/CodeGen/X86/code-align-loops.ll | 35 +++++++++++++++++++++--
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index c63a265fe27789..da6c758d53d487 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -1183,7 +1183,7 @@ void CodeGenPrepare::eliminateMostlyEmptyBlock(BasicBlock *BB) {
     }
   }
 
-  // Reserve loop Metadata.
+  // Preserve loop Metadata.
   if (BI->hasMetadata(LLVMContext::MD_loop)) {
     for (auto *Pred : predecessors(BB))
       Pred->getTerminator()->copyMetadata(*BI, LLVMContext::MD_loop);
diff --git a/llvm/lib/CodeGen/MachineLoopInfo.cpp b/llvm/lib/CodeGen/MachineLoopInfo.cpp
index b82b41fac662d7..b0d74ecd6a8549 100644
--- a/llvm/lib/CodeGen/MachineLoopInfo.cpp
+++ b/llvm/lib/CodeGen/MachineLoopInfo.cpp
@@ -211,6 +211,8 @@ MDNode *MachineLoop::getLoopID() const {
             break;
           }
         }
+        if (!MD)
+          continue;
         if (!LoopID)
           LoopID = MD;
         else if (MD != LoopID)
diff --git a/llvm/test/CodeGen/X86/code-align-loops.ll b/llvm/test/CodeGen/X86/code-align-loops.ll
index 40ad3b4c8d600e..164df14df4173f 100644
--- a/llvm/test/CodeGen/X86/code-align-loops.ll
+++ b/llvm/test/CodeGen/X86/code-align-loops.ll
@@ -96,8 +96,8 @@ 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:
+; test3 and test4 is to check if .p2align can be correctly set on loops with
+; multi latches. The IR is generated from below simple C file:
 ; $ clang -O0 -S -emit-llvm loop.c
 ; $ cat loop.c
 ; int test3() {
@@ -146,6 +146,37 @@ while.end:                                        ; preds = %while.cond
   ret i32 %3
 }
 
+; CHECK-LABEL: test4_multilatch:
+; ALIGN: .p2align 6, 0x90
+; ALIGN-NEXT: .LBB3_4: # %bb4
+define void @test4_multilatch(i32 %a, i32 %b, i32 %c, i32 %d) nounwind {
+entry:
+  br label %bb1
+
+bb1:                               ; preds = %bb2, %bb4, %entry
+  call void @bar()
+  %cmp3 = icmp sgt i32 %c, 10
+  br i1 %cmp3, label %bb3, label %bb4
+
+bb2:                                ; preds = %bb3
+  call void @bar()
+  %cmp1 = icmp sgt i32 %a, 11
+  br i1 %cmp1, label %bb1, label %exit, !llvm.loop !0
+
+bb3:                                ; preds = %bb1
+  call void @bar()
+  %cmp2 = icmp sgt i32 %b, 12
+  br i1 %cmp2, label %bb2, label %exit
+
+bb4:                                ; preds = %bb1
+  call void @bar()
+  %cmp4 = icmp sgt i32 %d, 14
+  br i1 %cmp4, label %bb1, label %exit
+
+exit:                               ; preds = %bb2, %bb3, %bb4
+  ret void
+}
+
 declare void @bar()
 declare void @var()