[llvm] a09e84b - [SimpleLoopUnswitch] Tweak token live-out behavior (#217916)

via llvm-commits llvm-commits at lists.llvm.org
Fri Aug 21 10:16:41 PDT 2026


Author: Keno Fischer
Date: 2026-08-21T19:16:36+02:00
New Revision: a09e84b4da7c8907cd5757e36bae3e93dfc524ee

URL: https://github.com/llvm/llvm-project/commit/a09e84b4da7c8907cd5757e36bae3e93dfc524ee
DIFF: https://github.com/llvm/llvm-project/commit/a09e84b4da7c8907cd5757e36bae3e93dfc524ee.diff

LOG: [SimpleLoopUnswitch] Tweak token live-out behavior (#217916)

The previous check (moved into Loop::isSafeToCloneConditionally in the
preceding NFC commit) rejected any token instruction with a use outside
its defining basic block. This is both too strict and too lax:

- Too strict: cloning only requires forming phis for values that are
live-out of the *loop*. Token uses that cross blocks but stay within the
loop are cloned along with their defs and need no phis, so they do not
inhibit non-trivial unswitching.

- Too lax: `isTokenTy()` misses token-like target extension types
(`isTokenLikeTy()`), which cannot be used in phi nodes either. A
live-out token-like value would have made unswitching produce invalid
IR.

Switch the check to reject exactly the token-like values that are
live-out of the loop, reusing the same use-scanning logic as
`isBlockInLCSSAForm` (`loopContainsUser`): a use in a phi is treated as
occurring in the corresponding incoming block, and users in unreachable
blocks are ignored. This requires threading a DominatorTree into the
check.

Split out of #151062 by review request.

Added: 
    llvm/test/Transforms/SimpleLoopUnswitch/token-uses.ll

Modified: 
    llvm/include/llvm/Analysis/LoopInfo.h
    llvm/lib/Analysis/LoopInfo.cpp
    llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Analysis/LoopInfo.h b/llvm/include/llvm/Analysis/LoopInfo.h
index bbac76c85f1ed..64caab77893e6 100644
--- a/llvm/include/llvm/Analysis/LoopInfo.h
+++ b/llvm/include/llvm/Analysis/LoopInfo.h
@@ -329,10 +329,11 @@ class LLVM_ABI Loop : public LoopBase<BasicBlock, Loop> {
   bool isSafeToClone() const;
 
   /// Like `isSafeToClone`, but for transformations where the cloned loop
-  /// bodies may run conditionally. This additionally rejects convergent
-  /// calls (which must not gain new control dependencies) and token values
-  /// that are used outside their defining block.
-  bool isSafeToCloneConditionally() const;
+  /// bodies may run conditionally. This additionally checks that we may form
+  /// phis for all values that are live-out from the loop (in particular that
+  /// no token-like values are live-out) and that there are no convergent calls
+  /// (which must not gain new control dependencies).
+  bool isSafeToCloneConditionally(const DominatorTree &DT) const;
 
   /// Returns true if the loop is annotated parallel.
   ///

diff  --git a/llvm/lib/Analysis/LoopInfo.cpp b/llvm/lib/Analysis/LoopInfo.cpp
index ab535e302b305..613e9887f3c45 100644
--- a/llvm/lib/Analysis/LoopInfo.cpp
+++ b/llvm/lib/Analysis/LoopInfo.cpp
@@ -530,14 +530,22 @@ bool Loop::isSafeToClone() const {
   return true;
 }
 
-bool Loop::isSafeToCloneConditionally() const {
+bool Loop::isSafeToCloneConditionally(const DominatorTree &DT) const {
   if (!isSafeToClone())
     return false;
 
   for (BasicBlock *BB : this->blocks()) {
     for (Instruction &I : *BB) {
-      if (I.getType()->isTokenTy() && I.isUsedOutsideOfBlock(BB))
-        return false;
+      // Token-like values cannot be used in PHI nodes, so cloning is only
+      // possible if all their uses are contained in the loop. Uses within
+      // the loop (even across blocks) are fine: cloning only requires
+      // forming phis for values that are live-out of the loop.
+      if (I.getType()->isTokenLikeTy()) {
+        for (const Use &U : I.uses()) {
+          if (!loopContainsUser(*this, *BB, U, DT))
+            return false;
+        }
+      }
       if (auto *CB = dyn_cast<CallBase>(&I)) {
         assert(!CB->cannotDuplicate() && "Checked by isSafeToClone().");
         if (CB->isConvergent())

diff  --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 619e48f0ce509..ee7998f5bdcc8 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -3358,8 +3358,9 @@ static bool collectUnswitchCandidatesWithInjections(
   return Found;
 }
 
-static bool isSafeForNoNTrivialUnswitching(Loop &L, LoopInfo &LI) {
-  if (!L.isSafeToCloneConditionally())
+static bool isSafeForNoNTrivialUnswitching(const DominatorTree &DT, Loop &L,
+                                           LoopInfo &LI) {
+  if (!L.isSafeToCloneConditionally(DT))
     return false;
 
   // Check if there are irreducible CFG cycles in this loop. If so, we cannot
@@ -3706,7 +3707,7 @@ static bool unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI,
     return false;
 
   // Perform legality checks.
-  if (!isSafeForNoNTrivialUnswitching(L, LI))
+  if (!isSafeForNoNTrivialUnswitching(DT, L, LI))
     return false;
 
   // For non-trivial unswitching, because it often creates new loops, we rely on

diff  --git a/llvm/test/Transforms/SimpleLoopUnswitch/token-uses.ll b/llvm/test/Transforms/SimpleLoopUnswitch/token-uses.ll
new file mode 100644
index 0000000000000..0f13e41851258
--- /dev/null
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/token-uses.ll
@@ -0,0 +1,179 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>)' -S < %s | FileCheck %s
+
+declare token @llvm.source_token()
+declare void @llvm.sink_token(token)
+declare void @a()
+declare void @b()
+
+; A token that is used in a 
diff erent block than its definition, but where all
+; uses are still inside the loop, does not prevent non-trivial unswitching:
+; cloning the loop only requires forming phis for values that are live-out.
+define void @token_used_across_blocks(i32 %n, i1 %cond) {
+; CHECK-LABEL: define void @token_used_across_blocks(
+; CHECK-SAME: i32 [[N:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[COND_FR:%.*]] = freeze i1 [[COND]]
+; CHECK-NEXT:    br i1 [[COND_FR]], label %[[ENTRY_SPLIT_US:.*]], label %[[ENTRY_SPLIT:.*]]
+; CHECK:       [[ENTRY_SPLIT_US]]:
+; CHECK-NEXT:    br label %[[LOOP_US:.*]]
+; CHECK:       [[LOOP_US]]:
+; CHECK-NEXT:    [[I_US:%.*]] = phi i32 [ 0, %[[ENTRY_SPLIT_US]] ], [ [[I_NEXT_US:%.*]], %[[LATCH_US:.*]] ]
+; CHECK-NEXT:    [[TOKEN_US:%.*]] = call token @llvm.source_token()
+; CHECK-NEXT:    br label %[[LEFT_US:.*]]
+; CHECK:       [[LEFT_US]]:
+; CHECK-NEXT:    call void @a()
+; CHECK-NEXT:    br label %[[LATCH_US]]
+; CHECK:       [[LATCH_US]]:
+; CHECK-NEXT:    call void @llvm.sink_token(token [[TOKEN_US]])
+; CHECK-NEXT:    [[I_NEXT_US]] = add i32 [[I_US]], 1
+; CHECK-NEXT:    [[EXIT_COND_US:%.*]] = icmp slt i32 [[I_NEXT_US]], [[N]]
+; CHECK-NEXT:    br i1 [[EXIT_COND_US]], label %[[LOOP_US]], label %[[EXIT_SPLIT_US:.*]]
+; CHECK:       [[EXIT_SPLIT_US]]:
+; CHECK-NEXT:    br label %[[EXIT:.*]]
+; CHECK:       [[ENTRY_SPLIT]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, %[[ENTRY_SPLIT]] ], [ [[I_NEXT:%.*]], %[[LATCH:.*]] ]
+; CHECK-NEXT:    [[TOKEN:%.*]] = call token @llvm.source_token()
+; CHECK-NEXT:    br label %[[RIGHT:.*]]
+; CHECK:       [[RIGHT]]:
+; CHECK-NEXT:    call void @b()
+; CHECK-NEXT:    br label %[[LATCH]]
+; CHECK:       [[LATCH]]:
+; CHECK-NEXT:    call void @llvm.sink_token(token [[TOKEN]])
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[EXIT_COND:%.*]] = icmp slt i32 [[I_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[EXIT_COND]], label %[[LOOP]], label %[[EXIT_SPLIT:.*]]
+; CHECK:       [[EXIT_SPLIT]]:
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %latch ]
+  %token = call token @llvm.source_token()
+  br i1 %cond, label %left, label %right
+
+left:
+  call void @a()
+  br label %latch
+
+right:
+  call void @b()
+  br label %latch
+
+latch:
+  call void @llvm.sink_token(token %token)
+  %i.next = add i32 %i, 1
+  %exit.cond = icmp slt i32 %i.next, %n
+  br i1 %exit.cond, label %loop, label %exit
+
+exit:
+  ret void
+}
+
+; A token that is live-out of the loop still prevents unswitching: we cannot
+; form a phi for it in the exit block.
+define void @token_live_out(i32 %n, i1 %cond) {
+; CHECK-LABEL: define void @token_live_out(
+; CHECK-SAME: i32 [[N:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LATCH:.*]] ]
+; CHECK-NEXT:    [[TOKEN:%.*]] = call token @llvm.source_token()
+; CHECK-NEXT:    br i1 [[COND]], label %[[LEFT:.*]], label %[[RIGHT:.*]]
+; CHECK:       [[LEFT]]:
+; CHECK-NEXT:    call void @a()
+; CHECK-NEXT:    br label %[[LATCH]]
+; CHECK:       [[RIGHT]]:
+; CHECK-NEXT:    call void @b()
+; CHECK-NEXT:    br label %[[LATCH]]
+; CHECK:       [[LATCH]]:
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[EXIT_COND:%.*]] = icmp slt i32 [[I_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[EXIT_COND]], label %[[LOOP]], label %[[EXIT:.*]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    call void @llvm.sink_token(token [[TOKEN]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %latch ]
+  %token = call token @llvm.source_token()
+  br i1 %cond, label %left, label %right
+
+left:
+  call void @a()
+  br label %latch
+
+right:
+  call void @b()
+  br label %latch
+
+latch:
+  %i.next = add i32 %i, 1
+  %exit.cond = icmp slt i32 %i.next, %n
+  br i1 %exit.cond, label %loop, label %exit
+
+exit:
+  call void @llvm.sink_token(token %token)
+  ret void
+}
+
+; Same for live-out token-like target extension types, which cannot be used in
+; phi nodes either.
+define void @token_like_live_out(i32 %n, i1 %cond, ptr %p) {
+; CHECK-LABEL: define void @token_like_live_out(
+; CHECK-SAME: i32 [[N:%.*]], i1 [[COND:%.*]], ptr [[P:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, %[[ENTRY]] ], [ [[I_NEXT:%.*]], %[[LATCH:.*]] ]
+; CHECK-NEXT:    [[HANDLE:%.*]] = load target("dx.RawBuffer", i32, 1, 0), ptr [[P]], align 8
+; CHECK-NEXT:    br i1 [[COND]], label %[[LEFT:.*]], label %[[RIGHT:.*]]
+; CHECK:       [[LEFT]]:
+; CHECK-NEXT:    call void @a()
+; CHECK-NEXT:    br label %[[LATCH]]
+; CHECK:       [[RIGHT]]:
+; CHECK-NEXT:    call void @b()
+; CHECK-NEXT:    br label %[[LATCH]]
+; CHECK:       [[LATCH]]:
+; CHECK-NEXT:    [[I_NEXT]] = add i32 [[I]], 1
+; CHECK-NEXT:    [[EXIT_COND:%.*]] = icmp slt i32 [[I_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[EXIT_COND]], label %[[LOOP]], label %[[EXIT:.*]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    store target("dx.RawBuffer", i32, 1, 0) [[HANDLE]], ptr [[P]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %i.next, %latch ]
+  %handle = load target("dx.RawBuffer", i32, 1, 0), ptr %p
+  br i1 %cond, label %left, label %right
+
+left:
+  call void @a()
+  br label %latch
+
+right:
+  call void @b()
+  br label %latch
+
+latch:
+  %i.next = add i32 %i, 1
+  %exit.cond = icmp slt i32 %i.next, %n
+  br i1 %exit.cond, label %loop, label %exit
+
+exit:
+  store target("dx.RawBuffer", i32, 1, 0) %handle, ptr %p
+  ret void
+}


        


More information about the llvm-commits mailing list