[llvm] Add loop metadata to the new branch when doing jump-threading (PR #157180)

[via llvm-commits]

https://github.com/asastry108 created https://github.com/llvm/llvm-project/pull/157180

In SimplifyCFG, in the jump threading optimization, a new branch is created.  The loop metadata was not being updated to the new branch.  The fix updates the new branch with the loop metadata.  This fixes an important regression where full unrolling was not happening on the transformed loop.

>From cdb8e08a928256c13aa93e3e9e60e94ef717a2fa Mon Sep 17 00:00:00 2001
From: asastry <asastry at nvidia.com>
Date: Thu, 28 Aug 2025 23:49:02 +0000
Subject: [PATCH] Add loop metadata to the new branch when doing jump-threading

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     |  3 ++
 .../preserve-llvm-loop-metadata-2.ll          | 30 +++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/preserve-llvm-loop-metadata-2.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 86d4750f6f000..193dad9672098 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3602,6 +3602,9 @@ foldCondBranchOnValueKnownInPredecessorImpl(BranchInst *BI, DomTreeUpdater *DTU,
     EdgeBB->setName(RealDest->getName() + ".critedge");
     EdgeBB->moveBefore(RealDest);
 
+    if (MDNode *LoopMD = BI->getMetadata(LLVMContext::MD_loop))
+      EdgeBB->getTerminator()->setMetadata(LLVMContext::MD_loop, LoopMD);
+
     // Update PHI nodes.
     addPredecessorToBlock(RealDest, EdgeBB, BB);
 
diff --git a/llvm/test/Transforms/SimplifyCFG/preserve-llvm-loop-metadata-2.ll b/llvm/test/Transforms/SimplifyCFG/preserve-llvm-loop-metadata-2.ll
new file mode 100644
index 0000000000000..9a296c928f401
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/preserve-llvm-loop-metadata-2.ll
@@ -0,0 +1,30 @@
+; RUN: opt -passes='simplifycfg' -S < %s | FileCheck %s
+
+; CHECK: br i1 %4, label %3, label %1, 
+; CHECK-SAME: llvm.loop
+
+define void @test(i32 %1 ) {
+.critedge:
+  br label %107
+
+107:                                              ; preds = %147, .critedge 
+  %111 = icmp eq i32 %1, 0
+  br i1 %111, label %112, label %156
+
+112:                                              ; preds = %107
+  br label %147
+
+147:                                              ; preds = %149, %112
+  %148 = phi i1 [ false, %149 ], [ true, %112 ]
+  br i1 %148, label %149, label %107, !llvm.loop !32
+
+149:                                              ; preds = %147
+  br label %147
+
+156:                                              ; preds = %107
+   ret void
+} 
+
+!32 = distinct !{!32, !33, !34}
+!33 = !{!"llvm.loop.unroll.enable"}
+!34 = !{!"llvm.loop.unroll.full"}