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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (asastry108)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/157180.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+3) 
- (added) llvm/test/Transforms/SimplifyCFG/preserve-llvm-loop-metadata-2.ll (+30) 


``````````diff
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"}

``````````

</details>


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