[llvm] [LoopUnroll] Allow runtime unroll with remainder for uniform trip count in convergent loops (PR #192819)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Apr 18 18:00:00 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-nvptx
Author: Luo Yuanke (LuoYuanke)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/192819.diff
6 Files Affected:
- (modified) llvm/include/llvm/Transforms/Utils/UnrollLoop.h (+1)
- (modified) llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp (+41-7)
- (modified) llvm/lib/Transforms/Utils/LoopUnroll.cpp (+1-1)
- (added) llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll (+70)
- (added) llvm/test/Transforms/LoopUnroll/NVPTX/lit.local.cfg (+2)
- (added) llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll (+68)
``````````diff
diff --git a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
index 364908ca1fad5..4cf361b64f0ca 100644
--- a/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
+++ b/llvm/include/llvm/Transforms/Utils/UnrollLoop.h
@@ -80,6 +80,7 @@ struct UnrollLoopOptions {
unsigned SCEVExpansionBudget;
bool RuntimeUnrollMultiExit = false;
bool AddAdditionalAccumulators = false;
+ bool CanHaveUnrollRemainder = true;
};
LLVM_ABI LoopUnrollResult UnrollLoop(Loop *L, UnrollLoopOptions ULO,
diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 885ae1f6dc02a..d70e1e6e4c953 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -21,6 +21,7 @@
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/UniformityAnalysis.h"
#include "llvm/Analysis/BlockFrequencyInfo.h"
#include "llvm/Analysis/CodeMetrics.h"
#include "llvm/Analysis/LoopAnalysisManager.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,8 @@ 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 +1387,15 @@ 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.
+ bool TripCountIsUniform = false;
+ if (!UCE.ConvergenceAllowsRuntime && UI) {
+ if (const SCEV *BTC = SE.getBackedgeTakenCount(L);
+ !isa<SCEVCouldNotCompute>(BTC))
+ TripCountIsUniform = 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 +1416,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");
@@ -1452,6 +1473,8 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
ULO.SCEVExpansionBudget = UP.SCEVExpansionBudget;
ULO.RuntimeUnrollMultiExit = UP.RuntimeUnrollMultiExit;
ULO.AddAdditionalAccumulators = UP.AddAdditionalAccumulators;
+ ULO.CanHaveUnrollRemainder =
+ UCE.ConvergenceAllowsRuntime || TripCountIsUniform;
LoopUnrollResult UnrollResult = UnrollLoop(
L, ULO, LI, &SE, &DT, &AC, &TTI, &ORE, PreserveLCSSA, &RemainderLoop, AA);
if (UnrollResult == LoopUnrollResult::Unmodified)
@@ -1545,6 +1568,10 @@ class LoopUnroll : public LoopPass {
const TargetTransformInfo &TTI =
getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+ UniformityInfo *UI = nullptr;
+ if (TTI.hasBranchDivergence(&F)) {
+ UI = &getAnalysis<UniformityInfoWrapperPass>().getUniformityInfo();
+ }
// 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 +1583,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 +1597,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 +1612,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 +1732,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();
@@ -1757,7 +1791,7 @@ PreservedAnalyses LoopUnrollPass::run(Function &F,
/*Threshold*/ std::nullopt, UnrollOpts.AllowPartial,
UnrollOpts.AllowRuntime, UnrollOpts.AllowUpperBound, LocalAllowPeeling,
UnrollOpts.AllowProfileBasedPeeling, UnrollOpts.FullUnrollMaxCount,
- &AA);
+ &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..6089ed18d4873 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
@@ -805,7 +805,7 @@ llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
return LoopUnrollResult::Unmodified;
}
- assert((!ULO.Runtime || canHaveUnrollRemainder(L)) &&
+ assert((!ULO.Runtime || ULO.CanHaveUnrollRemainder) &&
"Can't runtime unroll if loop contains a convergent operation.");
bool EpilogProfitability =
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..339edb35f066c
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/AMDGPU/unroll-convergent-uniform-tripcount.ll
@@ -0,0 +1,70 @@
+; 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.
+;
+; CHECK-LABEL: @runtime_unroll_uniform_trip_count(
+; CHECK: call void @llvm.amdgcn.s.barrier()
+; CHECK: call void @llvm.amdgcn.s.barrier()
+; CHECK: epil
+define amdgpu_kernel void @runtime_unroll_uniform_trip_count(ptr addrspace(1) noalias nocapture %out, ptr addrspace(1) noalias nocapture %in, i32 %n) #1 {
+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.
+;
+; CHECK-LABEL: @runtime_unroll_divergent_trip_count(
+; CHECK: call void @llvm.amdgcn.s.barrier()
+; CHECK-NOT: call void @llvm.amdgcn.s.barrier()
+; CHECK-NOT: epil
+define amdgpu_kernel void @runtime_unroll_divergent_trip_count(ptr addrspace(1) noalias nocapture %out, ptr addrspace(1) noalias nocapture %in) #1 {
+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..c1d0018ff3f5b
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/NVPTX/unroll-convergent-uniform-tripcount.ll
@@ -0,0 +1,68 @@
+; 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.
+;
+; CHECK-LABEL: @runtime_unroll_uniform_trip_count(
+; CHECK: call i32 @llvm.nvvm.shfl.down.i32
+; CHECK: call i32 @llvm.nvvm.shfl.down.i32
+; CHECK: epil
+define void @runtime_unroll_uniform_trip_count(ptr %out, ptr %in, i32 %n) {
+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.
+;
+; CHECK-LABEL: @runtime_unroll_divergent_trip_count(
+; CHECK: call i32 @llvm.nvvm.shfl.down.i32
+; CHECK-NOT: call i32 @llvm.nvvm.shfl.down.i32
+; CHECK-NOT: epil
+define void @runtime_unroll_divergent_trip_count(ptr %out, ptr %in) {
+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}
``````````
</details>
https://github.com/llvm/llvm-project/pull/192819
More information about the llvm-commits
mailing list