[SimplifyCFG] Add loop metadata to the new branch when doing jump-threading

[asastry108]

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.

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: None (asastry108)

Changes

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

[asastry108]

Is there any further change I need to do ?

[asastry108]

Gentle ping to know if anything further needed from me

[asastry108]

Can someone update me if this merge request is ready for merge or more work is needed ?

[XChy]

LGTM

[asastry108]

Hi, Who will be merging this change ?
It is a first time for me and I need help to get it in.

[nikic]

It's not really obvious to me that copying the metadata in this way is correct. The loop metadata is attached to the loop latch. If we're threading the latch, doesn't it depend on the threaded successor whether the result is another latch or not?

In your specific example, it seems like we actually end with the loop metadata attached to a terminator that is a latch for two different loops, so that can probably end up applying the metadata to the wrong loop? (Previously it was applied to the outer loop, but now it also applies to the inner?)

[XChy]

In your specific example, it seems like we actually end with the loop metadata attached to a terminator that is a latch for two different loops, so that can probably end up applying the metadata to the wrong loop? (Previously it was applied to the outer loop, but now it also applies to the inner?)

Sorry for misreading the LangRef. I was not aware of applying the metadata to both loops. I agree it's incorrect too.

If we're threading the latch, doesn't it depend on the threaded successor whether the result is another latch or not?

Looks like we cannot fix metadata without loop info here. And there may be other problems, such as whether the metadata on these latches is not persistent for the same loop.