[llvm] [LoopUnroll] Compute pragma info before loop is unrolled (PR #187034)

[Justin Fargnoli via llvm-commits]

https://github.com/justinfargnoli created https://github.com/llvm/llvm-project/pull/187034

Address https://github.com/llvm/llvm-project/pull/184529#issuecomment-4074393657

>From 4ad38d30192054b2d2c65a5face4ff9a473063a8 Mon Sep 17 00:00:00 2001
From: Justin Fargnoli <jfargnoli at nvidia.com>
Date: Tue, 17 Mar 2026 14:56:38 +0000
Subject: [PATCH] [LoopUnroll] Compute pragma info before loop is unrolled

---
 llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp |  3 +-
 llvm/test/Transforms/LoopUnroll/bbi-115997.ll | 57 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/LoopUnroll/bbi-115997.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
index 5142f90ac7079..6a6566bf6e948 100644
--- a/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp
@@ -1434,6 +1434,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
 
   // Save loop properties before it is transformed.
   MDNode *OrigLoopID = L->getLoopID();
+  UnrollPragmaInfo PInfo = UnrollPragmaInfo(L);
 
   // Unroll the loop.
   Loop *RemainderLoop = nullptr;
@@ -1477,7 +1478,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
   // If loop has an unroll count pragma or unrolled by explicitly set count
   // mark loop as unrolled to prevent unrolling beyond that requested.
   if (UnrollResult != LoopUnrollResult::FullyUnrolled &&
-      UnrollPragmaInfo(L).ExplicitUnroll)
+      PInfo.ExplicitUnroll)
     L->setLoopAlreadyUnrolled();
 
   return UnrollResult;
diff --git a/llvm/test/Transforms/LoopUnroll/bbi-115997.ll b/llvm/test/Transforms/LoopUnroll/bbi-115997.ll
new file mode 100644
index 0000000000000..6b15202eedce2
--- /dev/null
+++ b/llvm/test/Transforms/LoopUnroll/bbi-115997.ll
@@ -0,0 +1,57 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=loop-unroll -S | FileCheck %s
+
+; Verify that llvm.loop.unroll.disable is added after partial unrolling
+; with a pragma unroll count, even when the last latch's exit branch is
+; folded into an unconditional branch by UnrollLoop.
+;
+; The loop has trip count 5 and unroll count 2. After unrolling, the
+; second copy's exit is folded (since iteration 2,4,... cannot be the
+; last), and the conditional branch is replaced with an unconditional
+; one that does not carry !llvm.loop metadata. setLoopAlreadyUnrolled
+; must still mark the loop to prevent further unrolling.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @test(ptr %p) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV_NEXT_1:%.*]], %[[LOOP_1:.*]] ]
+; CHECK-NEXT:    [[PTR:%.*]] = getelementptr i32, ptr [[P]], i64 [[IV]]
+; CHECK-NEXT:    store i32 0, ptr [[PTR]], align 4
+; CHECK-NEXT:    [[IV_NEXT:%.*]] = add nuw nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i64 [[IV_NEXT]], 5
+; CHECK-NEXT:    br i1 [[CMP]], label %[[LOOP_1]], label %[[EXIT:.*]]
+; CHECK:       [[LOOP_1]]:
+; CHECK-NEXT:    [[PTR_1:%.*]] = getelementptr i32, ptr [[P]], i64 [[IV_NEXT]]
+; CHECK-NEXT:    store i32 0, ptr [[PTR_1]], align 4
+; CHECK-NEXT:    [[IV_NEXT_1]] = add nuw nsw i64 [[IV]], 2
+; CHECK-NEXT:    br label %[[LOOP]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %ptr = getelementptr i32, ptr %p, i64 %iv
+  store i32 0, ptr %ptr
+  %iv.next = add nuw nsw i64 %iv, 1
+  %cmp = icmp ne i64 %iv.next, 5
+  br i1 %cmp, label %loop, label %exit, !llvm.loop !0
+
+exit:
+  ret void
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.count", i32 2}
+;.
+; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]]}
+; CHECK: [[META1]] = !{!"llvm.loop.unroll.disable"}
+;.