[llvm] [JumpThreading] Don't thread into blocks with convergent operations (PR #213870)

[via llvm-commits]

https://github.com/39ali created https://github.com/llvm/llvm-project/pull/213870

Jump threading refuses to duplicate a block containing convergent calls
(`getJumpThreadDuplicationCost` returns infinite cost), but nothing prevents
it from rerouting an edge so that a block containing convergent operations
gains a new entry edge from a different control-flow region: `threadEdge` and
`threadThroughTwoBasicBlocks` duplicate only the (cost-checked) threaded
blocks and then wire the clones directly into `SuccBB`, whose contents are
never examined.

For convergent operations the destination's entry edges are semantically
load-bearing: targets using the token-less `convergent` attribute (NVPTX)
derive warp reconvergence points from the CFG merge structure. The new edge
can bypass the merge point that was the only valid reconvergence location
between a divergent branch and the convergent operation, so the operation
executes with a partially converged warp.

This miscompiles real code: Triton kernels through plain
`opt -passes='default<O3>'` end up executing `llvm.nvvm.ldmatrix.sync` with a
partially converged warp.

Fix: refuse `tryThreadEdge` / `maybethreadThroughTwoBasicBlocks` when the
threading destination contains a convergent operation.

Fixes #213868

>From 1a81627ad20ae6b27742e35ac41eb3e61cf904d0 Mon Sep 17 00:00:00 2001
From: 39ali <alimilhim5 at gmail.com>
Date: Tue, 4 Aug 2026 11:06:32 +0300
Subject: [PATCH] [JumpThreading] Don't thread into blocks with convergent
 operations

---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp  | 24 +++++
 .../JumpThreading/thread-into-convergent.ll   | 91 +++++++++++++++++++
 2 files changed, 115 insertions(+)
 create mode 100644 llvm/test/Transforms/JumpThreading/thread-into-convergent.ll

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 7a0542aac83fc..1ad312110b3b5 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -2147,6 +2147,14 @@ void JumpThreadingPass::cloneInstructions(ValueToValueMapTy &ValueMapping,
   }
 }
 
+/// Return true if BB contains a convergent operation.
+static bool blockContainsConvergentOp(const BasicBlock *BB) {
+  return llvm::any_of(*BB, [](const Instruction &I) {
+    const auto *CB = dyn_cast<CallBase>(&I);
+    return CB && CB->isConvergent();
+  });
+}
+
 /// Attempt to thread through two successive basic blocks.
 bool JumpThreadingPass::maybethreadThroughTwoBasicBlocks(BasicBlock *BB,
                                                          Value *Cond) {
@@ -2267,6 +2275,14 @@ bool JumpThreadingPass::maybethreadThroughTwoBasicBlocks(BasicBlock *BB,
     return false;
   }
 
+  // Don't create a new entry edge into a block with convergent operations.
+  if (blockContainsConvergentOp(SuccBB)) {
+    LLVM_DEBUG(dbgs() << "  Not threading through BB '" << BB->getName()
+                      << "' to dest BB '" << SuccBB->getName()
+                      << "' - destination contains convergent operations!\n");
+    return false;
+  }
+
   // Compute the cost of duplicating BB and PredBB.
   unsigned BBCost = getJumpThreadDuplicationCost(
       TTI, BB, BB->getTerminator(), BBDupThreshold);
@@ -2389,6 +2405,14 @@ bool JumpThreadingPass::tryThreadEdge(
     return false;
   }
 
+  // Don't create a new entry edge into a block with convergent operations.
+  if (blockContainsConvergentOp(SuccBB)) {
+    LLVM_DEBUG(dbgs() << "  Not threading across BB '" << BB->getName()
+                      << "' to dest BB '" << SuccBB->getName()
+                      << "' - destination contains convergent operations!\n");
+    return false;
+  }
+
   unsigned JumpThreadCost = getJumpThreadDuplicationCost(
       TTI, BB, BB->getTerminator(), BBDupThreshold);
   if (JumpThreadCost > BBDupThreshold) {
diff --git a/llvm/test/Transforms/JumpThreading/thread-into-convergent.ll b/llvm/test/Transforms/JumpThreading/thread-into-convergent.ll
new file mode 100644
index 0000000000000..9d1adc7bb314c
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/thread-into-convergent.ll
@@ -0,0 +1,91 @@
+; RUN: opt -S -passes=jump-threading < %s | FileCheck %s
+
+; Jump threading must not create a new entry edge into a block that contains
+; convergent operations: the new edge changes the control-flow paths by which
+; threads reach the convergent call, and targets using the token-less
+; `convergent` attribute  derive warp reconvergence points from
+; the CFG merge structure. Threading past the %retest merge leaves no valid
+; reconvergence point between the divergence introduced in %A/%C and the
+; convergent call, so the call executes with a partially converged warp.
+
+declare void @convergent_op() convergent
+declare void @plain_op()
+
+; %t and %f are opaque (runtime-uniform) conditions. On the %skip path %t is
+; known false, so jump threading wants to duplicate %mid+%retest and route
+; that path directly into %tail — a new entry edge into a convergent block.
+; That must be refused. The %ret_body branch in between keeps the remaining
+; structure from folding away.
+
+define void @dont_thread_into_convergent(i1 %t, i1 %f) {
+; CHECK-LABEL: @dont_thread_into_convergent(
+; CHECK: retest:
+; CHECK: br i1 %t, label %C, label %tail
+; CHECK: tail:
+; CHECK-NOT: tail.thread
+; CHECK: call void @convergent_op()
+entry:
+  br i1 %t, label %A, label %skip
+
+A:
+  call void @plain_op()
+  br label %mid
+
+skip:
+  br label %mid
+
+mid:
+  br i1 %f, label %ret_body, label %retest
+
+ret_body:
+  call void @plain_op()
+  ret void
+
+retest:
+  br i1 %t, label %C, label %tail
+
+C:
+  call void @plain_op()
+  br label %tail
+
+tail:
+  call void @convergent_op()
+  ret void
+}
+
+; Same shape without the convergent call: threading fires and %tail2 gets the
+; threaded entry edge. Guards against the new check being overly broad.
+
+define void @thread_into_plain(i1 %t, i1 %f) {
+; CHECK-LABEL: @thread_into_plain(
+; CHECK: mid2.thread:
+; CHECK: tail2:
+; CHECK-SAME: preds = %mid2.thread
+entry:
+  br i1 %t, label %A2, label %skip2
+
+A2:
+  call void @plain_op()
+  br label %mid2
+
+skip2:
+  br label %mid2
+
+mid2:
+  br i1 %f, label %ret_body2, label %retest2
+
+ret_body2:
+  call void @plain_op()
+  ret void
+
+retest2:
+  br i1 %t, label %C2, label %tail2
+
+C2:
+  call void @plain_op()
+  br label %tail2
+
+tail2:
+  call void @plain_op()
+  ret void
+}