[SimplifyCFG] Avoid threading loop-header branches in convergent functions

[shiltian]

SimplifyCFG can fold a conditional branch when the condition is known from a predecessor. When the destination is a loop header in a convergent function, this can change the dynamic convergence structure of the loop even though the scalar CFG rewrite is otherwise valid.

Skip this fold for loop-header branches in convergent functions so convergent control flow is preserved.

Fixes ROCM-26496.

[shiltian]

• [SimplifyCFG] Avoid threading loop-header branches in convergent functions #204958 👈
• [NFC] Pre-commit a test case for a SimplifyCFG issue #204957
• main

This stack of pull requests is managed by sgh.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Shilei Tian (shiltian)

Changes

SimplifyCFG can fold a conditional branch when the condition is known from
a predecessor. When the destination is a loop header in a convergent function,
this can change the dynamic convergence structure of the loop even though the
scalar CFG rewrite is otherwise valid.

Skip this fold for loop-header branches in convergent functions so convergent
control flow is preserved.

Fixes ROCM-26496.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+4-1)
• (modified) llvm/test/Transforms/SimplifyCFG/convergent-loop-header.ll (+6-4)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 524947dd2e95d..3166d56a4dc0b 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3713,7 +3713,10 @@ bool SimplifyCFGOpt::foldCondBranchOnValueKnownInPredecessor(CondBrInst *BI) {
   // Note: If BB is a loop header then there is a risk that threading introduces
   // a non-canonical loop by moving a back edge. So we avoid this optimization
   // for loop headers if NeedCanonicalLoop is set.
-  if (Options.NeedCanonicalLoop && is_contained(LoopHeaders, BI->getParent()))
+  // Also avoid threading loop headers in convergent functions, since changing
+  // the branch structure can change the dynamic instances of convergent ops.
+  if ((Options.NeedCanonicalLoop || BI->getFunction()->isConvergent()) &&
+      is_contained(LoopHeaders, BI->getParent()))
     return false;
 
   std::optional<bool> Result;
diff --git a/llvm/test/Transforms/SimplifyCFG/convergent-loop-header.ll b/llvm/test/Transforms/SimplifyCFG/convergent-loop-header.ll
index a988d9e469880..45ac7bef0edbb 100644
--- a/llvm/test/Transforms/SimplifyCFG/convergent-loop-header.ll
+++ b/llvm/test/Transforms/SimplifyCFG/convergent-loop-header.ll
@@ -12,13 +12,15 @@ define void @preserve_loop_header_branch(i32 %tid, ptr %ptr) convergent {
 ; CHECK-NEXT:    br i1 [[COND]], label %[[OUTER_THEN:.*]], label %[[LOOP_HEADER:.*]]
 ; CHECK:       [[OUTER_THEN]]:
 ; CHECK-NEXT:    store i32 1, ptr [[PTR]], align 4
-; CHECK-NEXT:    br label %[[LOOP_THEN:.*]]
-; CHECK:       [[LOOP_THEN]]:
-; CHECK-NEXT:    store i32 2, ptr [[PTR]], align 4
 ; CHECK-NEXT:    br label %[[LOOP_HEADER]]
 ; CHECK:       [[LOOP_HEADER]]:
+; CHECK-NEXT:    br i1 [[COND]], label %[[LOOP_THEN:.*]], label %[[LOOP_BODY:.*]]
+; CHECK:       [[LOOP_THEN]]:
+; CHECK-NEXT:    store i32 2, ptr [[PTR]], align 4
+; CHECK-NEXT:    br label %[[LOOP_BODY]]
+; CHECK:       [[LOOP_BODY]]:
 ; CHECK-NEXT:    call void @barrier()
-; CHECK-NEXT:    br i1 [[COND]], label %[[LOOP_THEN]], label %[[EXIT:.*]]
+; CHECK-NEXT:    br i1 [[COND]], label %[[LOOP_HEADER]], label %[[EXIT:.*]]
 ; CHECK:       [[EXIT]]:
 ; CHECK-NEXT:    ret void
 ;