[llvm] [JumpThread] Fix `JumpThreading` pass to skip merging when both blocks contain convergence loop/entry intrinsics. (PR #170247)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Lucie Choi (luciechoi)

<details>
<summary>Changes</summary>

Fixes a bug https://github.com/llvm/llvm-project/issues/165642. After this fix, optimization passes for the example in the bug.

[LLVM Spec](https://llvm.org/docs/ConvergentOperations.html#llvm-experimental-convergence-loop) states that only a single loop / entry convergence token can be included in a basic block.

This PR fixes the issue in `JumpThreading` pass so that when a basic block and its predecessor both contain such convergence intrinsics, it skips merging the two blocks.

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


4 Files Affected:

- (modified) llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h (+4) 
- (modified) llvm/lib/Transforms/Scalar/JumpThreading.cpp (+7) 
- (modified) llvm/lib/Transforms/Utils/BasicBlockUtils.cpp (+1-1) 
- (added) llvm/test/Transforms/JumpThreading/skip-merging-duplicate-convergence-intrinsics.ll (+60) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
index 49885b7f06a15..c951597498033 100644
--- a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
@@ -45,6 +45,10 @@ class ReturnInst;
 class TargetLibraryInfo;
 class Value;
 
+/// Check if the given basic block contains any loop or entry convergent
+/// intrinsic instructions.
+LLVM_ABI bool HasLoopOrEntryConvergenceToken(const BasicBlock *BB);
+
 /// Replace contents of every block in \p BBs with single unreachable
 /// instruction. If \p Updates is specified, collect all necessary DT updates
 /// into this vector. If \p KeepOneInputPHIs is true, one-input Phis in
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index c7d71eb5633ec..aa8c80abb66ab 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -1920,6 +1920,13 @@ bool JumpThreadingPass::maybeMergeBasicBlockIntoOnlyPred(BasicBlock *BB) {
   if (Unreachable.count(SinglePred))
     return false;
 
+  // Don't merge if both the basic block and the predecessor contain loop or
+  // entry convergent intrinsics, since there may only be one convergence token
+  // per block.
+  if (HasLoopOrEntryConvergenceToken(BB) &&
+      HasLoopOrEntryConvergenceToken(SinglePred))
+    return false;
+
   // If SinglePred was a loop header, BB becomes one.
   if (LoopHeaders.erase(SinglePred))
     LoopHeaders.insert(BB);
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 076c5da4393fc..da3e60da9df74 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -92,7 +92,7 @@ emptyAndDetachBlock(BasicBlock *BB,
          "applying corresponding DTU updates.");
 }
 
-static bool HasLoopOrEntryConvergenceToken(const BasicBlock *BB) {
+bool llvm::HasLoopOrEntryConvergenceToken(const BasicBlock *BB) {
   for (const Instruction &I : *BB) {
     const ConvergenceControlInst *CCI = dyn_cast<ConvergenceControlInst>(&I);
     if (CCI && (CCI->isLoop() || CCI->isEntry()))
diff --git a/llvm/test/Transforms/JumpThreading/skip-merging-duplicate-convergence-intrinsics.ll b/llvm/test/Transforms/JumpThreading/skip-merging-duplicate-convergence-intrinsics.ll
new file mode 100644
index 0000000000000..571b6062ecc2e
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/skip-merging-duplicate-convergence-intrinsics.ll
@@ -0,0 +1,60 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -S -passes=jump-threading | FileCheck %s
+
+define void @main(i32 %val1) #0 {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call token @llvm.experimental.convergence.entry()
+; CHECK-NEXT:    br label [[FOR_COND_I:%.*]]
+; CHECK:       for.cond.i:
+; CHECK-NEXT:    [[CMP_I:%.*]] = phi i1 [ true, [[ENTRY:%.*]] ], [ false, [[CLEANUP_I:%.*]] ]
+; CHECK-NEXT:    [[TMP1:%.*]] = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token [[TMP0]]) ]
+; CHECK-NEXT:    br i1 [[CMP_I]], label [[FOR_COND1_I:%.*]], label [[EXIT:%.*]]
+; CHECK:       for.cond1.i:
+; CHECK-NEXT:    [[CMP2_I:%.*]] = phi i1 [ true, [[FOR_COND_I]] ], [ false, [[FOR_BODY4_I:%.*]] ]
+; CHECK-NEXT:    br i1 [[CMP2_I]], label [[FOR_BODY4_I]], label [[CLEANUP_I]]
+; CHECK:       for.body4.i:
+; CHECK-NEXT:    [[CMP5_I:%.*]] = icmp eq i32 [[VAL1:%.*]], 0
+; CHECK-NEXT:    br i1 [[CMP5_I]], label [[IF_THEN_I:%.*]], label [[FOR_COND1_I]]
+; CHECK:       if.then.i:
+; CHECK-NEXT:    call void @func_test() [ "convergencectrl"(token [[TMP1]]) ]
+; CHECK-NEXT:    br label [[CLEANUP_I]]
+; CHECK:       cleanup.i:
+; CHECK-NEXT:    br label [[FOR_COND_I]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = call token @llvm.experimental.convergence.entry()
+  br label %for.cond.i
+
+for.cond.i:
+  %cmp.i = phi i1 [ true, %entry ], [ false, %cleanup.i ]
+  %1 = call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %0) ]
+  br i1 %cmp.i, label %for.cond1.i, label %exit
+
+for.cond1.i:
+  %cmp2.i = phi i1 [ true, %for.cond.i ], [ false, %for.body4.i ]
+  br i1 %cmp2.i, label %for.body4.i, label %cleanup.i
+
+for.body4.i:
+  %cmp5.i = icmp eq i32 %val1, 0
+  br i1 %cmp5.i, label %if.then.i, label %for.cond1.i
+
+if.then.i:
+  call void @func_test() [ "convergencectrl"(token %1) ]
+  br label %cleanup.i
+
+cleanup.i:
+  br label %for.cond.i
+
+exit:
+  ret void
+}
+
+declare token @llvm.experimental.convergence.entry() #0
+declare token @llvm.experimental.convergence.loop() #0
+
+declare void @func_test() #0
+
+attributes #0 = { convergent }

``````````

</details>


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