[llvm] [LoopUnroll] Allow runtime unroll with remainder for uniform trip count in convergent loops (PR #192819)
Luo Yuanke via llvm-commits
llvm-commits at lists.llvm.org
Thu Apr 23 05:42:42 PDT 2026
https://github.com/LuoYuanke updated https://github.com/llvm/llvm-project/pull/192819
>From 6d1681606f81277c651e3946732c2ad17882c659 Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Sun, 19 Apr 2026 08:55:55 +0800
Subject: [PATCH 1/5] [LoopUnroll] Allow runtime unroll with remainder for
uniform trip count in convergent loops
When a loop contains convergent operations, runtime unrolling with a
remainder is currently blocked because the remainder prelude adds a
control-flow dependency. However, if the backedge-taken count is
uniform across all threads, all threads take the same path through the
prelude, making the remainder safe.
Add divergence analysis to check if the BTC SCEV is uniform, and if so,
allow AllowRemainder and UP.Runtime even for convergent loops. This
resolves the TODO at the former line 1375.
Assisted by AI
---
llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp | 45 ++++-
llvm/lib/Transforms/Utils/LoopUnroll.cpp | 3 -
.../unroll-convergent-uniform-tripcount.ll | 186 ++++++++++++++++++
.../Transforms/LoopUnroll/NVPTX/lit.local.cfg | 2 +
.../unroll-convergent-uniform-tripcount.ll | 156 +++++++++++++++
5 files changed, 381 insertions(+), 11 deletions(-)
create mode 100644 llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll
create mode 100644 llvm/test/Transforms/LoopUnroll/NVPTX/lit.local.cfg
create mode 100644 llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 885ae1f6dc02a..cc80b29c38b00 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -32,6 +32,7 @@
#include "llvm/Analysis/ProfileSummaryInfo.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/UniformityAnalysis.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/Constant.h"
@@ -757,6 +758,19 @@ static bool hasRuntimeUnrollDisablePragma(const Loop *L) {
return getUnrollMetadataForLoop(L, "llvm.loop.unroll.runtime.disable");
}
+/// Returns true if the SCEV expression is uniform, i.e., all threads in a
+/// convergent execution agree on its value. Recursively checks operands.
+static bool isSCEVUniform(const SCEV *S, UniformityInfo &UI) {
+ if (isa<SCEVConstant>(S))
+ return true;
+ if (auto *U = dyn_cast<SCEVUnknown>(S))
+ return UI.isUniform(U->getValue());
+ for (const SCEV *Op : S->operands())
+ if (!isSCEVUniform(Op, UI))
+ return false;
+ return true;
+}
+
// If loop has an unroll_count pragma return the (necessarily
// positive) value from the pragma. Otherwise return 0.
static unsigned unrollCountPragmaValue(const Loop *L) {
@@ -1253,7 +1267,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
std::optional<bool> ProvidedAllowPeeling,
std::optional<bool> ProvidedAllowProfileBasedPeeling,
std::optional<unsigned> ProvidedFullUnrollMaxCount,
- AAResults *AA = nullptr) {
+ AAResults *AA = nullptr, UniformityInfo *UI = nullptr) {
LLVM_DEBUG(dbgs() << "Loop Unroll: F["
<< L->getHeader()->getParent()->getName() << "] Loop %"
@@ -1372,9 +1386,13 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
// is unsafe -- it adds a control-flow dependency to the convergent
// operation. Therefore restrict remainder loop (try unrolling without).
//
- // TODO: This is somewhat conservative; we could allow the remainder if the
- // trip count is uniform.
- UP.AllowRemainder &= UCE.ConvergenceAllowsRuntime;
+ // However, if the trip count is uniform across all threads, all threads
+ // will take the same path through the prelude, so the remainder is safe.
+ const SCEV *BTC = SE.getBackedgeTakenCount(L);
+ bool TripCountIsUniform = !UCE.ConvergenceAllowsRuntime && UI &&
+ !isa<SCEVCouldNotCompute>(BTC) &&
+ isSCEVUniform(BTC, *UI);
+ UP.AllowRemainder &= UCE.ConvergenceAllowsRuntime || TripCountIsUniform;
// Try to find the trip count upper bound if we cannot find the exact trip
// count.
@@ -1395,7 +1413,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
return LoopUnrollResult::Unmodified;
}
- UP.Runtime &= UCE.ConvergenceAllowsRuntime;
+ UP.Runtime &= UCE.ConvergenceAllowsRuntime || TripCountIsUniform;
if (PP.PeelCount) {
assert(UP.Count == 1 && "Cannot perform peel and unroll in the same step");
@@ -1545,6 +1563,10 @@ class LoopUnroll : public LoopPass {
const TargetTransformInfo &TTI =
getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+ UniformityInfo *UI =
+ TTI.hasBranchDivergence(&F)
+ ? &getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo()
+ : nullptr;
// For the old PM, we can't use OptimizationRemarkEmitter as an analysis
// pass. Function analyses need to be preserved across loop transformations
// but ORE cannot be preserved (see comment before the pass definition).
@@ -1556,7 +1578,8 @@ class LoopUnroll : public LoopPass {
/*OnlyFullUnroll*/ false, OnlyWhenForced, ForgetAllSCEV, ProvidedCount,
ProvidedThreshold, ProvidedAllowPartial, ProvidedRuntime,
ProvidedUpperBound, ProvidedAllowPeeling,
- ProvidedAllowProfileBasedPeeling, ProvidedFullUnrollMaxCount);
+ ProvidedAllowProfileBasedPeeling, ProvidedFullUnrollMaxCount,
+ /*AA*/ nullptr, UI);
if (Result == LoopUnrollResult::FullyUnrolled)
LPM.markLoopAsDeleted(*L);
@@ -1569,6 +1592,7 @@ class LoopUnroll : public LoopPass {
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addRequired<AssumptionCacheTracker>();
AU.addRequired<TargetTransformInfoWrapperPass>();
+ AU.addRequired<UniformityInfoWrapperPass>();
// FIXME: Loop passes are required to preserve domtree, and for now we just
// recreate dom info if anything gets unrolled.
getLoopAnalysisUsage(AU);
@@ -1583,6 +1607,7 @@ INITIALIZE_PASS_BEGIN(LoopUnroll, "loop-unroll", "Unroll loops", false, false)
INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
INITIALIZE_PASS_DEPENDENCY(LoopPass)
INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(UniformityInfoWrapperPass)
INITIALIZE_PASS_END(LoopUnroll, "loop-unroll", "Unroll loops", false, false)
Pass *llvm::createLoopUnrollPass(int OptLevel, bool OnlyWhenForced,
@@ -1702,6 +1727,10 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
AAResults &AA = AM.getResult<AAManager>(F);
+ UniformityInfo *UI = nullptr;
+ if (TTI.hasBranchDivergence(&F))
+ UI = &AM.getResult<UniformityInfoAnalysis>(F);
+
LoopAnalysisManager *LAM = nullptr;
if (auto *LAMProxy = AM.getCachedResult<LoopAnalysisManagerFunctionProxy>(F))
LAM = &LAMProxy->getManager();
@@ -1756,8 +1785,8 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
/*Count*/ std::nullopt,
/*Threshold*/ std::nullopt, UnrollOpts.AllowPartial,
UnrollOpts.AllowRuntime, UnrollOpts.AllowUpperBound, LocalAllowPeeling,
- UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount,
- &AA);
+ UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount, &AA,
+ UI);
Changed |= Result != LoopUnrollResult::Unmodified;
// The parent must not be damaged by unrolling!
diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index 3094beeb5005c..345ef94d4a44f 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -805,9 +805,6 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
return LoopUnrollResult::Unmodified;
}
- assert((!ULO.Runtime || canHaveUnrollRemainder(L)) &&
- "Can't runtime unroll if loop contains a convergent operation.");
-
bool EpilogProfitability =
UnrollRuntimeEpilog.getNumOccurrences() ? UnrollRuntimeEpilog
: isEpilogProfitable(L);
diff --git a/llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll b/llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll
new file mode 100644
index 0000000000000..acd6d285c463c
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll
@@ -0,0 +1,186 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: -p --version 6
+; RUN: opt -mtriple=amdgcn-unknown-amdhsa -mcpu=hawaii -passes=loop-unroll -unroll-runtime -unroll-allow-partial -S < %s | FileCheck %s
+
+declare void @llvm.amdgcn.s.barrier() #0
+
+; This loop has a convergent barrier and a runtime trip count that depends on
+; a uniform value (kernel argument, which is passed in SGPR). Since the trip
+; count is uniform across all threads, runtime unrolling with a remainder is
+; safe and should be performed.
+;
+define amdgpu_kernel void @runtime_unroll_uniform_trip_count(ptr addrspace(1) noalias nocapture %out, ptr addrspace(1) noalias nocapture %in, i32 %n) #1 {
+; CHECK-LABEL: define amdgpu_kernel void @runtime_unroll_uniform_trip_count(ptr addrspace(1) noalias captures(none) %out, ptr addrspace(1) noalias captures(none) %in, i32 %n) #1 {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %0 = add i32 %n, -1
+; CHECK-NEXT: %xtraiter = and i32 %n, 7
+; CHECK-NEXT: %1 = icmp ult i32 %0, 7
+; CHECK-NEXT: br i1 %1, label %for.body.epil.preheader, label %entry.new
+; CHECK: entry.new:
+; CHECK-NEXT: %unroll_iter = sub i32 %n, %xtraiter
+; CHECK-NEXT: br label %for.body
+; CHECK: for.body:
+; CHECK-NEXT: %indvars.iv = phi i32 [ 0, %entry.new ], [ %indvars.iv.next.7, %for.body ]
+; CHECK-NEXT: %sum.02 = phi i32 [ 0, %entry.new ], [ %add.7, %for.body ]
+; CHECK-NEXT: %niter = phi i32 [ 0, %entry.new ], [ %niter.next.7, %for.body ]
+; CHECK-NEXT: %arrayidx.in = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv
+; CHECK-NEXT: %arrayidx.out = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv
+; CHECK-NEXT: %load = load i32, ptr addrspace(1) %arrayidx.in, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add = add i32 %load, %sum.02
+; CHECK-NEXT: store i32 %add, ptr addrspace(1) %arrayidx.out, align 4
+; CHECK-NEXT: %indvars.iv.next = add nuw nsw i32 %indvars.iv, 1
+; CHECK-NEXT: %arrayidx.in.1 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next
+; CHECK-NEXT: %arrayidx.out.1 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next
+; CHECK-NEXT: %load.1 = load i32, ptr addrspace(1) %arrayidx.in.1, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.1 = add i32 %load.1, %add
+; CHECK-NEXT: store i32 %add.1, ptr addrspace(1) %arrayidx.out.1, align 4
+; CHECK-NEXT: %indvars.iv.next.1 = add nuw nsw i32 %indvars.iv, 2
+; CHECK-NEXT: %arrayidx.in.2 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.1
+; CHECK-NEXT: %arrayidx.out.2 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.1
+; CHECK-NEXT: %load.2 = load i32, ptr addrspace(1) %arrayidx.in.2, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.2 = add i32 %load.2, %add.1
+; CHECK-NEXT: store i32 %add.2, ptr addrspace(1) %arrayidx.out.2, align 4
+; CHECK-NEXT: %indvars.iv.next.2 = add nuw nsw i32 %indvars.iv, 3
+; CHECK-NEXT: %arrayidx.in.3 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.2
+; CHECK-NEXT: %arrayidx.out.3 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.2
+; CHECK-NEXT: %load.3 = load i32, ptr addrspace(1) %arrayidx.in.3, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.3 = add i32 %load.3, %add.2
+; CHECK-NEXT: store i32 %add.3, ptr addrspace(1) %arrayidx.out.3, align 4
+; CHECK-NEXT: %indvars.iv.next.3 = add nuw nsw i32 %indvars.iv, 4
+; CHECK-NEXT: %arrayidx.in.4 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.3
+; CHECK-NEXT: %arrayidx.out.4 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.3
+; CHECK-NEXT: %load.4 = load i32, ptr addrspace(1) %arrayidx.in.4, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.4 = add i32 %load.4, %add.3
+; CHECK-NEXT: store i32 %add.4, ptr addrspace(1) %arrayidx.out.4, align 4
+; CHECK-NEXT: %indvars.iv.next.4 = add nuw nsw i32 %indvars.iv, 5
+; CHECK-NEXT: %arrayidx.in.5 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.4
+; CHECK-NEXT: %arrayidx.out.5 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.4
+; CHECK-NEXT: %load.5 = load i32, ptr addrspace(1) %arrayidx.in.5, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.5 = add i32 %load.5, %add.4
+; CHECK-NEXT: store i32 %add.5, ptr addrspace(1) %arrayidx.out.5, align 4
+; CHECK-NEXT: %indvars.iv.next.5 = add nuw nsw i32 %indvars.iv, 6
+; CHECK-NEXT: %arrayidx.in.6 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.5
+; CHECK-NEXT: %arrayidx.out.6 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.5
+; CHECK-NEXT: %load.6 = load i32, ptr addrspace(1) %arrayidx.in.6, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.6 = add i32 %load.6, %add.5
+; CHECK-NEXT: store i32 %add.6, ptr addrspace(1) %arrayidx.out.6, align 4
+; CHECK-NEXT: %indvars.iv.next.6 = add nuw nsw i32 %indvars.iv, 7
+; CHECK-NEXT: %arrayidx.in.7 = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.next.6
+; CHECK-NEXT: %arrayidx.out.7 = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.next.6
+; CHECK-NEXT: %load.7 = load i32, ptr addrspace(1) %arrayidx.in.7, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.7 = add i32 %load.7, %add.6
+; CHECK-NEXT: store i32 %add.7, ptr addrspace(1) %arrayidx.out.7, align 4
+; CHECK-NEXT: %indvars.iv.next.7 = add i32 %indvars.iv, 8
+; CHECK-NEXT: %niter.next.7 = add i32 %niter, 8
+; CHECK-NEXT: %niter.ncmp.7 = icmp eq i32 %niter.next.7, %unroll_iter
+; CHECK-NEXT: br i1 %niter.ncmp.7, label %for.end.unr-lcssa, label %for.body
+; CHECK: for.end.unr-lcssa:
+; CHECK-NEXT: %indvars.iv.unr = phi i32 [ %indvars.iv.next.7, %for.body ]
+; CHECK-NEXT: %sum.02.unr = phi i32 [ %add.7, %for.body ]
+; CHECK-NEXT: %lcmp.mod = icmp ne i32 %xtraiter, 0
+; CHECK-NEXT: br i1 %lcmp.mod, label %for.body.epil.preheader, label %for.end
+; CHECK: for.body.epil.preheader:
+; CHECK-NEXT: %indvars.iv.epil.init = phi i32 [ 0, %entry ], [ %indvars.iv.unr, %for.end.unr-lcssa ]
+; CHECK-NEXT: %sum.02.epil.init = phi i32 [ 0, %entry ], [ %sum.02.unr, %for.end.unr-lcssa ]
+; CHECK-NEXT: %lcmp.mod1 = icmp ne i32 %xtraiter, 0
+; CHECK-NEXT: call void @llvm.assume(i1 %lcmp.mod1)
+; CHECK-NEXT: br label %for.body.epil
+; CHECK: for.body.epil:
+; CHECK-NEXT: %indvars.iv.epil = phi i32 [ %indvars.iv.next.epil, %for.body.epil ], [ %indvars.iv.epil.init, %for.body.epil.preheader ]
+; CHECK-NEXT: %sum.02.epil = phi i32 [ %add.epil, %for.body.epil ], [ %sum.02.epil.init, %for.body.epil.preheader ]
+; CHECK-NEXT: %epil.iter = phi i32 [ 0, %for.body.epil.preheader ], [ %epil.iter.next, %for.body.epil ]
+; CHECK-NEXT: %arrayidx.in.epil = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv.epil
+; CHECK-NEXT: %arrayidx.out.epil = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv.epil
+; CHECK-NEXT: %load.epil = load i32, ptr addrspace(1) %arrayidx.in.epil, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add.epil = add i32 %load.epil, %sum.02.epil
+; CHECK-NEXT: store i32 %add.epil, ptr addrspace(1) %arrayidx.out.epil, align 4
+; CHECK-NEXT: %indvars.iv.next.epil = add i32 %indvars.iv.epil, 1
+; CHECK-NEXT: %exitcond.epil = icmp eq i32 %indvars.iv.next.epil, %n
+; CHECK-NEXT: %epil.iter.next = add i32 %epil.iter, 1
+; CHECK-NEXT: %epil.iter.cmp = icmp ne i32 %epil.iter.next, %xtraiter
+; CHECK-NEXT: br i1 %epil.iter.cmp, label %for.body.epil, label %for.end.epilog-lcssa, !llvm.loop !0
+; CHECK: for.end.epilog-lcssa:
+; CHECK-NEXT: br label %for.end
+; CHECK: for.end:
+; CHECK-NEXT: ret void
+;
+entry:
+ br label %for.body
+
+for.body:
+ %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+ %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+ %arrayidx.in = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv
+ %arrayidx.out = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv
+ %load = load i32, ptr addrspace(1) %arrayidx.in
+ call void @llvm.amdgcn.s.barrier() #0
+ %add = add i32 %load, %sum.02
+ store i32 %add, ptr addrspace(1) %arrayidx.out
+ %indvars.iv.next = add i32 %indvars.iv, 1
+ %exitcond = icmp eq i32 %indvars.iv.next, %n
+ br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+ ret void
+}
+
+; This loop has a convergent barrier and a divergent trip count (derived from
+; llvm.amdgcn.workitem.id.x, which is divergent). Since the trip count is
+; divergent, runtime unrolling with a remainder must still be blocked.
+;
+define amdgpu_kernel void @runtime_unroll_divergent_trip_count(ptr addrspace(1) noalias nocapture %out, ptr addrspace(1) noalias nocapture %in) #1 {
+; CHECK-LABEL: define amdgpu_kernel void @runtime_unroll_divergent_trip_count(ptr addrspace(1) noalias captures(none) %out, ptr addrspace(1) noalias captures(none) %in) #1 {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %tid = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT: %n = add i32 %tid, 1
+; CHECK-NEXT: br label %for.body
+; CHECK: for.body:
+; CHECK-NEXT: %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+; CHECK-NEXT: %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+; CHECK-NEXT: %arrayidx.in = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv
+; CHECK-NEXT: %arrayidx.out = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv
+; CHECK-NEXT: %load = load i32, ptr addrspace(1) %arrayidx.in, align 4
+; CHECK-NEXT: call void @llvm.amdgcn.s.barrier() #4
+; CHECK-NEXT: %add = add i32 %load, %sum.02
+; CHECK-NEXT: store i32 %add, ptr addrspace(1) %arrayidx.out, align 4
+; CHECK-NEXT: %indvars.iv.next = add i32 %indvars.iv, 1
+; CHECK-NEXT: %exitcond = icmp eq i32 %indvars.iv.next, %n
+; CHECK-NEXT: br i1 %exitcond, label %for.end, label %for.body
+; CHECK: for.end:
+; CHECK-NEXT: ret void
+;
+entry:
+ %tid = call i32 @llvm.amdgcn.workitem.id.x()
+ %n = add i32 %tid, 1
+ br label %for.body
+
+for.body:
+ %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+ %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+ %arrayidx.in = getelementptr inbounds i32, ptr addrspace(1) %in, i32 %indvars.iv
+ %arrayidx.out = getelementptr inbounds i32, ptr addrspace(1) %out, i32 %indvars.iv
+ %load = load i32, ptr addrspace(1) %arrayidx.in
+ call void @llvm.amdgcn.s.barrier() #0
+ %add = add i32 %load, %sum.02
+ store i32 %add, ptr addrspace(1) %arrayidx.out
+ %indvars.iv.next = add i32 %indvars.iv, 1
+ %exitcond = icmp eq i32 %indvars.iv.next, %n
+ br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+ ret void
+}
+
+declare i32 @llvm.amdgcn.workitem.id.x() #2
+
+attributes #0 = { nounwind convergent }
+attributes #1 = { nounwind }
+attributes #2 = { nounwind readnone willreturn }
diff --git a/llvm/test/Transforms/LoopUnroll/NVPTX/lit.local.cfg b/llvm/test/Transforms/LoopUnroll/NVPTX/lit.local.cfg
new file mode 100644
index 0000000000000..0d37b86e1c8e6
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/NVPTX/lit.local.cfg
@@ -0,0 +1,2 @@
+if not "NVPTX" in config.root.targets:
+ config.unsupported = True
diff --git a/llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll b/llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll
new file mode 100644
index 0000000000000..4d7b96cfe5840
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll
@@ -0,0 +1,156 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: -p --version 6
+; RUN: opt -mtriple=nvptx64-nvidia-cuda -mcpu=sm_70 -passes=loop-unroll -unroll-runtime -unroll-allow-partial -S < %s | FileCheck %s
+
+declare i32 @llvm.nvvm.shfl.down.i32(i32, i32, i32)
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+
+; This loop has a convergent shfl_down intrinsic and a runtime trip count that
+; is uniform (kernel argument). Since the trip count is uniform across all
+; threads, runtime unrolling with a remainder is safe.
+;
+define void @runtime_unroll_uniform_trip_count(ptr %out, ptr %in, i32 %n) {
+; CHECK-LABEL: define ptx_kernel void @runtime_unroll_uniform_trip_count(ptr %out, ptr %in, i32 %n) #2 {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %0 = add i32 %n, -1
+; CHECK-NEXT: %xtraiter = and i32 %n, 3
+; CHECK-NEXT: %1 = icmp ult i32 %0, 3
+; CHECK-NEXT: br i1 %1, label %for.body.epil.preheader, label %entry.new
+; CHECK: entry.new:
+; CHECK-NEXT: %unroll_iter = sub i32 %n, %xtraiter
+; CHECK-NEXT: br label %for.body
+; CHECK: for.body:
+; CHECK-NEXT: %indvars.iv = phi i32 [ 0, %entry.new ], [ %indvars.iv.next.3, %for.body ]
+; CHECK-NEXT: %sum.02 = phi i32 [ 0, %entry.new ], [ %add.3, %for.body ]
+; CHECK-NEXT: %niter = phi i32 [ 0, %entry.new ], [ %niter.next.3, %for.body ]
+; CHECK-NEXT: %arrayidx.in = getelementptr inbounds i32, ptr %in, i32 %indvars.iv
+; CHECK-NEXT: %arrayidx.out = getelementptr inbounds i32, ptr %out, i32 %indvars.iv
+; CHECK-NEXT: %load = load i32, ptr %arrayidx.in, align 4
+; CHECK-NEXT: %shfl = call i32 @llvm.nvvm.shfl.down.i32(i32 %sum.02, i32 1, i32 31)
+; CHECK-NEXT: %add = add i32 %load, %shfl
+; CHECK-NEXT: store i32 %add, ptr %arrayidx.out, align 4
+; CHECK-NEXT: %indvars.iv.next = add nuw nsw i32 %indvars.iv, 1
+; CHECK-NEXT: %arrayidx.in.1 = getelementptr inbounds i32, ptr %in, i32 %indvars.iv.next
+; CHECK-NEXT: %arrayidx.out.1 = getelementptr inbounds i32, ptr %out, i32 %indvars.iv.next
+; CHECK-NEXT: %load.1 = load i32, ptr %arrayidx.in.1, align 4
+; CHECK-NEXT: %shfl.1 = call i32 @llvm.nvvm.shfl.down.i32(i32 %add, i32 1, i32 31)
+; CHECK-NEXT: %add.1 = add i32 %load.1, %shfl.1
+; CHECK-NEXT: store i32 %add.1, ptr %arrayidx.out.1, align 4
+; CHECK-NEXT: %indvars.iv.next.1 = add nuw nsw i32 %indvars.iv, 2
+; CHECK-NEXT: %arrayidx.in.2 = getelementptr inbounds i32, ptr %in, i32 %indvars.iv.next.1
+; CHECK-NEXT: %arrayidx.out.2 = getelementptr inbounds i32, ptr %out, i32 %indvars.iv.next.1
+; CHECK-NEXT: %load.2 = load i32, ptr %arrayidx.in.2, align 4
+; CHECK-NEXT: %shfl.2 = call i32 @llvm.nvvm.shfl.down.i32(i32 %add.1, i32 1, i32 31)
+; CHECK-NEXT: %add.2 = add i32 %load.2, %shfl.2
+; CHECK-NEXT: store i32 %add.2, ptr %arrayidx.out.2, align 4
+; CHECK-NEXT: %indvars.iv.next.2 = add nuw nsw i32 %indvars.iv, 3
+; CHECK-NEXT: %arrayidx.in.3 = getelementptr inbounds i32, ptr %in, i32 %indvars.iv.next.2
+; CHECK-NEXT: %arrayidx.out.3 = getelementptr inbounds i32, ptr %out, i32 %indvars.iv.next.2
+; CHECK-NEXT: %load.3 = load i32, ptr %arrayidx.in.3, align 4
+; CHECK-NEXT: %shfl.3 = call i32 @llvm.nvvm.shfl.down.i32(i32 %add.2, i32 1, i32 31)
+; CHECK-NEXT: %add.3 = add i32 %load.3, %shfl.3
+; CHECK-NEXT: store i32 %add.3, ptr %arrayidx.out.3, align 4
+; CHECK-NEXT: %indvars.iv.next.3 = add i32 %indvars.iv, 4
+; CHECK-NEXT: %niter.next.3 = add i32 %niter, 4
+; CHECK-NEXT: %niter.ncmp.3 = icmp eq i32 %niter.next.3, %unroll_iter
+; CHECK-NEXT: br i1 %niter.ncmp.3, label %for.end.unr-lcssa, label %for.body
+; CHECK: for.end.unr-lcssa:
+; CHECK-NEXT: %indvars.iv.unr = phi i32 [ %indvars.iv.next.3, %for.body ]
+; CHECK-NEXT: %sum.02.unr = phi i32 [ %add.3, %for.body ]
+; CHECK-NEXT: %lcmp.mod = icmp ne i32 %xtraiter, 0
+; CHECK-NEXT: br i1 %lcmp.mod, label %for.body.epil.preheader, label %for.end
+; CHECK: for.body.epil.preheader:
+; CHECK-NEXT: %indvars.iv.epil.init = phi i32 [ 0, %entry ], [ %indvars.iv.unr, %for.end.unr-lcssa ]
+; CHECK-NEXT: %sum.02.epil.init = phi i32 [ 0, %entry ], [ %sum.02.unr, %for.end.unr-lcssa ]
+; CHECK-NEXT: %lcmp.mod1 = icmp ne i32 %xtraiter, 0
+; CHECK-NEXT: call void @llvm.assume(i1 %lcmp.mod1)
+; CHECK-NEXT: br label %for.body.epil
+; CHECK: for.body.epil:
+; CHECK-NEXT: %indvars.iv.epil = phi i32 [ %indvars.iv.next.epil, %for.body.epil ], [ %indvars.iv.epil.init, %for.body.epil.preheader ]
+; CHECK-NEXT: %sum.02.epil = phi i32 [ %add.epil, %for.body.epil ], [ %sum.02.epil.init, %for.body.epil.preheader ]
+; CHECK-NEXT: %epil.iter = phi i32 [ 0, %for.body.epil.preheader ], [ %epil.iter.next, %for.body.epil ]
+; CHECK-NEXT: %arrayidx.in.epil = getelementptr inbounds i32, ptr %in, i32 %indvars.iv.epil
+; CHECK-NEXT: %arrayidx.out.epil = getelementptr inbounds i32, ptr %out, i32 %indvars.iv.epil
+; CHECK-NEXT: %load.epil = load i32, ptr %arrayidx.in.epil, align 4
+; CHECK-NEXT: %shfl.epil = call i32 @llvm.nvvm.shfl.down.i32(i32 %sum.02.epil, i32 1, i32 31)
+; CHECK-NEXT: %add.epil = add i32 %load.epil, %shfl.epil
+; CHECK-NEXT: store i32 %add.epil, ptr %arrayidx.out.epil, align 4
+; CHECK-NEXT: %indvars.iv.next.epil = add i32 %indvars.iv.epil, 1
+; CHECK-NEXT: %exitcond.epil = icmp eq i32 %indvars.iv.next.epil, %n
+; CHECK-NEXT: %epil.iter.next = add i32 %epil.iter, 1
+; CHECK-NEXT: %epil.iter.cmp = icmp ne i32 %epil.iter.next, %xtraiter
+; CHECK-NEXT: br i1 %epil.iter.cmp, label %for.body.epil, label %for.end.epilog-lcssa, !llvm.loop !0
+; CHECK: for.end.epilog-lcssa:
+; CHECK-NEXT: br label %for.end
+; CHECK: for.end:
+; CHECK-NEXT: ret void
+;
+entry:
+ br label %for.body
+
+for.body:
+ %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+ %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+ %arrayidx.in = getelementptr inbounds i32, ptr %in, i32 %indvars.iv
+ %arrayidx.out = getelementptr inbounds i32, ptr %out, i32 %indvars.iv
+ %load = load i32, ptr %arrayidx.in
+ %shfl = call i32 @llvm.nvvm.shfl.down.i32(i32 %sum.02, i32 1, i32 31)
+ %add = add i32 %load, %shfl
+ store i32 %add, ptr %arrayidx.out
+ %indvars.iv.next = add i32 %indvars.iv, 1
+ %exitcond = icmp eq i32 %indvars.iv.next, %n
+ br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+ ret void
+}
+
+; This loop has a convergent shfl_down intrinsic and a divergent trip count
+; (derived from tid.x, which is divergent on NVPTX). Since the trip count is
+; divergent, runtime unrolling with a remainder must still be blocked.
+;
+define void @runtime_unroll_divergent_trip_count(ptr %out, ptr %in) {
+; CHECK-LABEL: define ptx_kernel void @runtime_unroll_divergent_trip_count(ptr %out, ptr %in) #2 {
+; CHECK-NEXT: entry:
+; CHECK-NEXT: %tid = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+; CHECK-NEXT: %n = add i32 %tid, 1
+; CHECK-NEXT: br label %for.body
+; CHECK: for.body:
+; CHECK-NEXT: %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+; CHECK-NEXT: %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+; CHECK-NEXT: %arrayidx.in = getelementptr inbounds i32, ptr %in, i32 %indvars.iv
+; CHECK-NEXT: %arrayidx.out = getelementptr inbounds i32, ptr %out, i32 %indvars.iv
+; CHECK-NEXT: %load = load i32, ptr %arrayidx.in, align 4
+; CHECK-NEXT: %shfl = call i32 @llvm.nvvm.shfl.down.i32(i32 %sum.02, i32 1, i32 31)
+; CHECK-NEXT: %add = add i32 %load, %shfl
+; CHECK-NEXT: store i32 %add, ptr %arrayidx.out, align 4
+; CHECK-NEXT: %indvars.iv.next = add i32 %indvars.iv, 1
+; CHECK-NEXT: %exitcond = icmp eq i32 %indvars.iv.next, %n
+; CHECK-NEXT: br i1 %exitcond, label %for.end, label %for.body
+; CHECK: for.end:
+; CHECK-NEXT: ret void
+;
+entry:
+ %tid = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+ %n = add i32 %tid, 1
+ br label %for.body
+
+for.body:
+ %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+ %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
+ %arrayidx.in = getelementptr inbounds i32, ptr %in, i32 %indvars.iv
+ %arrayidx.out = getelementptr inbounds i32, ptr %out, i32 %indvars.iv
+ %load = load i32, ptr %arrayidx.in
+ %shfl = call i32 @llvm.nvvm.shfl.down.i32(i32 %sum.02, i32 1, i32 31)
+ %add = add i32 %load, %shfl
+ store i32 %add, ptr %arrayidx.out
+ %indvars.iv.next = add i32 %indvars.iv, 1
+ %exitcond = icmp eq i32 %indvars.iv.next, %n
+ br i1 %exitcond, label %for.end, label %for.body
+
+for.end:
+ ret void
+}
+
+!nvvm.annotations = !{!0, !1}
+!0 = !{ptr @runtime_unroll_uniform_trip_count, !"kernel", i32 1}
+!1 = !{ptr @runtime_unroll_divergent_trip_count, !"kernel", i32 1}
>From ac8074f360ab00672c9d8ebf0a0811bb3df1a3ba Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Wed, 22 Apr 2026 09:19:53 +0800
Subject: [PATCH 2/5] Address comments
---
.../llvm/Transforms/Utils/UnrollLoop.h | 3 +-
llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp | 43 +++++++++++--------
2 files changed, 27 insertions(+), 19 deletions(-)
diff --git a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
index 364908ca1fad5..fb994a02a87f1 100644
--- a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
+++ b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
@@ -159,7 +159,8 @@ class UnrollCostEstimator {
LLVM_ABI UnrollCostEstimator(const Loop *L, const TargetTransformInfo &TTI,
const SmallPtrSetImpl<const Value *> &EphValues,
- unsigned BEInsns);
+ unsigned BEInsns,
+ bool TripCountIsUniform = false);
/// Whether it is legal to unroll this loop.
LLVM_ABI bool canUnroll() const;
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index cc80b29c38b00..9687a9c9db140 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -688,7 +688,8 @@ static std::optional<EstimatedUnrollCost> analyzeLoopUnrollCost(
UnrollCostEstimator::UnrollCostEstimator(
const Loop *L, const TargetTransformInfo &TTI,
- const SmallPtrSetImpl<const Value *> &EphValues, unsigned BEInsns) {
+ const SmallPtrSetImpl<const Value *> &EphValues, unsigned BEInsns,
+ bool TripCountIsUniform) {
CodeMetrics Metrics;
for (BasicBlock *BB : L->blocks())
Metrics.analyzeBasicBlock(BB, TTI, EphValues, /* PrepareForLTO= */ false,
@@ -698,8 +699,9 @@ UnrollCostEstimator::UnrollCostEstimator(
Convergence = Metrics.Convergence;
LoopSize = Metrics.NumInsts;
ConvergenceAllowsRuntime =
- Metrics.Convergence != ConvergenceKind::Uncontrolled &&
- !getLoopConvergenceHeart(L);
+ (Metrics.Convergence != ConvergenceKind::Uncontrolled &&
+ !getLoopConvergenceHeart(L)) ||
+ TripCountIsUniform;
// Don't allow an estimate of size zero. This would allows unrolling of loops
// with huge iteration counts, which is a compile time problem even if it's
@@ -760,14 +762,18 @@ static bool hasRuntimeUnrollDisablePragma(const Loop *L) {
/// Returns true if the SCEV expression is uniform, i.e., all threads in a
/// convergent execution agree on its value. Recursively checks operands.
+/// Returns false if the SCEV could not be computed.
static bool isSCEVUniform(const SCEV *S, UniformityInfo &UI) {
+ if (isa<SCEVCouldNotCompute>(S))
+ return false;
if (isa<SCEVConstant>(S))
return true;
if (auto *U = dyn_cast<SCEVUnknown>(S))
return UI.isUniform(U->getValue());
- for (const SCEV *Op : S->operands())
+ for (const SCEV *Op : S->operands()) {
if (!isSCEVUniform(Op, UI))
return false;
+ }
return true;
}
@@ -1338,7 +1344,13 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
SmallPtrSet<const Value *, 32> EphValues;
CodeMetrics::collectEphemeralValues(L, &AC, EphValues);
- UnrollCostEstimator UCE(L, TTI, EphValues, UP.BEInsns);
+ // Check if the backedge-taken count is uniform before constructing UCE.
+ // This is used to allow runtime unrolling with a remainder for convergent
+ // loops when all threads agree on the trip count.
+ const SCEV *BTC = SE.getBackedgeTakenCount(L);
+ bool TripCountIsUniform = UI && isSCEVUniform(BTC, *UI);
+
+ UnrollCostEstimator UCE(L, TTI, EphValues, UP.BEInsns, TripCountIsUniform);
if (!UCE.canUnroll())
return LoopUnrollResult::Unmodified;
@@ -1385,14 +1397,9 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
// to do the first few instructions before we hit the unrolled loop
// is unsafe -- it adds a control-flow dependency to the convergent
// operation. Therefore restrict remainder loop (try unrolling without).
- //
- // However, if the trip count is uniform across all threads, all threads
- // will take the same path through the prelude, so the remainder is safe.
- const SCEV *BTC = SE.getBackedgeTakenCount(L);
- bool TripCountIsUniform = !UCE.ConvergenceAllowsRuntime && UI &&
- !isa<SCEVCouldNotCompute>(BTC) &&
- isSCEVUniform(BTC, *UI);
- UP.AllowRemainder &= UCE.ConvergenceAllowsRuntime || TripCountIsUniform;
+ // This is already accounted for in UCE.ConvergenceAllowsRuntime, which
+ // also considers whether the trip count is uniform across all threads.
+ UP.AllowRemainder &= UCE.ConvergenceAllowsRuntime;
// Try to find the trip count upper bound if we cannot find the exact trip
// count.
@@ -1413,7 +1420,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
return LoopUnrollResult::Unmodified;
}
- UP.Runtime &= UCE.ConvergenceAllowsRuntime || TripCountIsUniform;
+ UP.Runtime &= UCE.ConvergenceAllowsRuntime;
if (PP.PeelCount) {
assert(UP.Count == 1 && "Cannot perform peel and unroll in the same step");
@@ -1579,7 +1586,7 @@ class LoopUnroll : public LoopPass {
ProvidedThreshold, ProvidedAllowPartial, ProvidedRuntime,
ProvidedUpperBound, ProvidedAllowPeeling,
ProvidedAllowProfileBasedPeeling, ProvidedFullUnrollMaxCount,
- /*AA*/ nullptr, UI);
+ /*AA=*/ nullptr, UI);
if (Result == LoopUnrollResult::FullyUnrolled)
LPM.markLoopAsDeleted(*L);
@@ -1727,9 +1734,9 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
AAResults &AA = AM.getResult<AAManager>(F);
- UniformityInfo *UI = nullptr;
- if (TTI.hasBranchDivergence(&F))
- UI = &AM.getResult<UniformityInfoAnalysis>(F);
+ UniformityInfo *UI = TTI.hasBranchDivergence(&F)
+ ? &AM.getResult<UniformityInfoAnalysis>(F)
+ : nullptr;
LoopAnalysisManager *LAM = nullptr;
if (auto *LAMProxy = AM.getCachedResult<LoopAnalysisManagerFunctionProxy>(F))
>From 040ad2c6f35ddad676fa15860a233c070b33337e Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Wed, 22 Apr 2026 10:34:24 +0800
Subject: [PATCH 3/5] Address comments
---
llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp | 9 ++++-----
1 file changed, 4 insertions(+), 5 deletions(-)
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 9687a9c9db140..eec1280ee313c 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -1273,7 +1273,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
std::optional<bool> ProvidedAllowPeeling,
std::optional<bool> ProvidedAllowProfileBasedPeeling,
std::optional<unsigned> ProvidedFullUnrollMaxCount,
- AAResults *AA = nullptr, UniformityInfo *UI = nullptr) {
+ UniformityInfo *UI = nullptr, AAResults *AA = nullptr) {
LLVM_DEBUG(dbgs() << "Loop Unroll: F["
<< L->getHeader()->getParent()->getName() << "] Loop %"
@@ -1585,8 +1585,7 @@ class LoopUnroll : public LoopPass {
/*OnlyFullUnroll*/ false, OnlyWhenForced, ForgetAllSCEV, ProvidedCount,
ProvidedThreshold, ProvidedAllowPartial, ProvidedRuntime,
ProvidedUpperBound, ProvidedAllowPeeling,
- ProvidedAllowProfileBasedPeeling, ProvidedFullUnrollMaxCount,
- /*AA=*/ nullptr, UI);
+ ProvidedAllowProfileBasedPeeling, ProvidedFullUnrollMaxCount, UI);
if (Result == LoopUnrollResult::FullyUnrolled)
LPM.markLoopAsDeleted(*L);
@@ -1792,8 +1791,8 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
/*Count*/ std::nullopt,
/*Threshold*/ std::nullopt, UnrollOpts.AllowPartial,
UnrollOpts.AllowRuntime, UnrollOpts.AllowUpperBound, LocalAllowPeeling,
- UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount, &AA,
- UI);
+ UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount, UI,
+ &AA);
Changed |= Result != LoopUnrollResult::Unmodified;
// The parent must not be damaged by unrolling!
>From 6e79fa02dd8fd23ca545e8d750eb6e3605e31177 Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Thu, 23 Apr 2026 08:44:23 +0800
Subject: [PATCH 4/5] Address comments
---
llvm/include/llvm/Transforms/Utils/UnrollLoop.h | 3 +++
llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp | 1 -
2 files changed, 3 insertions(+), 1 deletion(-)
diff --git a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
index fb994a02a87f1..271bb5621c8e0 100644
--- a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
+++ b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
@@ -157,6 +157,9 @@ class UnrollCostEstimator {
ConvergenceKind Convergence;
bool ConvergenceAllowsRuntime;
+ /// \param TripCountIsUniform If true, all threads in a convergent execution
+ /// agree on the trip count, so runtime unrolling with a remainder is safe
+ /// even for loops with uncontrolled convergent operations.
LLVM_ABI UnrollCostEstimator(const Loop *L, const TargetTransformInfo &TTI,
const SmallPtrSetImpl<const Value *> &EphValues,
unsigned BEInsns,
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index eec1280ee313c..ec3e173a9f9aa 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -1349,7 +1349,6 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
// loops when all threads agree on the trip count.
const SCEV *BTC = SE.getBackedgeTakenCount(L);
bool TripCountIsUniform = UI && isSCEVUniform(BTC, *UI);
-
UnrollCostEstimator UCE(L, TTI, EphValues, UP.BEInsns, TripCountIsUniform);
if (!UCE.canUnroll())
return LoopUnrollResult::Unmodified;
>From f52fda88eb4ddd31f138e68b10bbf009a350ab27 Mon Sep 17 00:00:00 2001
From: Yuanke Luo <ykluo at birentech.com>
Date: Thu, 23 Apr 2026 20:42:04 +0800
Subject: [PATCH 5/5] Address comments
---
llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index ec3e173a9f9aa..7b17ad9ce7dc0 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -698,6 +698,9 @@ UnrollCostEstimator::UnrollCostEstimator(
NotDuplicatable = Metrics.notDuplicatable;
Convergence = Metrics.Convergence;
LoopSize = Metrics.NumInsts;
+ // Convergent operations make the remainder prelude unsafe by adding a
+ // control-flow dependency, unless the trip count is uniform per
+ // UniformityInfo, in which case all paths agree and the remainder is safe.
ConvergenceAllowsRuntime =
(Metrics.Convergence != ConvergenceKind::Uncontrolled &&
!getLoopConvergenceHeart(L)) ||
@@ -1395,9 +1398,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
// If the loop contains a convergent operation, the prelude we'd add
// to do the first few instructions before we hit the unrolled loop
// is unsafe -- it adds a control-flow dependency to the convergent
- // operation. Therefore restrict remainder loop (try unrolling without).
- // This is already accounted for in UCE.ConvergenceAllowsRuntime, which
- // also considers whether the trip count is uniform across all threads.
+ // operation. Therefore restrict remainder loop (try unrolling without).
UP.AllowRemainder &= UCE.ConvergenceAllowsRuntime;
// Try to find the trip count upper bound if we cannot find the exact trip
More information about the llvm-commits
mailing list