[LoopUnroll] Introduce PragmaUnrollFullMaxIterations as a hard cap on how many iterations we try to unroll (#78648)

Fixes [PR77842](#77842) where
UBSAN causes pragma full unroll to try and unroll INT_MAX times. This
sets a cap to make sure we don't attempt this and crash the compiler.

Testing:
ninja check-all with new test

---------

Co-authored-by: Nikita Popov <[email protected]>

Author: modiking

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

@@ -173,6 +173,10 @@ static cl::opt<unsigned>
                            cl::desc("Default threshold (max size of unrolled "
                                     "loop), used in all but O3 optimizations"));
 
+static cl::opt<unsigned> PragmaUnrollFullMaxIterations(
+    "pragma-unroll-full-max-iterations", cl::init(1'000'000), cl::Hidden,
+    cl::desc("Maximum allowed iterations to unroll under pragma unroll full."));
+
 /// A magic value for use with the Threshold parameter to indicate
 /// that the loop unroll should be performed regardless of how much
 /// code expansion would result.
@@ -776,8 +780,17 @@ shouldPragmaUnroll(Loop *L, const PragmaInfo &PInfo,
       return PInfo.PragmaCount;
   }
 
-  if (PInfo.PragmaFullUnroll && TripCount != 0)
+  if (PInfo.PragmaFullUnroll && TripCount != 0) {
+    // Certain cases with UBSAN can cause trip count to be calculated as
+    // INT_MAX, Block full unrolling at a reasonable limit so that the compiler
+    // doesn't hang trying to unroll the loop. See PR77842
+    if (TripCount > PragmaUnrollFullMaxIterations) {
+      LLVM_DEBUG(dbgs() << "Won't unroll; trip count is too large\n");
+      return std::nullopt;
+    }
+
     return TripCount;
+  }
 
   if (PInfo.PragmaEnableUnroll && !TripCount && MaxTripCount &&
       MaxTripCount <= UP.MaxUpperBound)
@@ -1282,7 +1295,7 @@ tryToUnrollLoop(Loop *L, DominatorTree &DT, LoopInfo *LI, ScalarEvolution &SE,
   }
 
   // Do not attempt partial/runtime unrolling in FullLoopUnrolling
-  if (OnlyFullUnroll && !(UP.Count >= MaxTripCount)) {
+  if (OnlyFullUnroll && (UP.Count < TripCount || UP.Count < MaxTripCount)) {
     LLVM_DEBUG(
         dbgs() << "Not attempting partial/runtime unroll in FullLoopUnroll.\n");
     return LoopUnrollResult::Unmodified;

llvm/test/Transforms/LoopUnroll/pr77842.ll

@@ -0,0 +1,54 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt -passes=loop-unroll-full -S %s | FileCheck %s
+
+; Validate that loop unroll full doesn't try to fully unroll values whose trip counts are too large.
+
+define void @foo(i64 %end) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i64 [[END:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOPHEADER:%.*]]
+; CHECK:       loopheader:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEW:%.*]], [[BACKEDGE:%.*]] ]
+; CHECK-NEXT:    [[EXIT:%.*]] = icmp eq i64 [[IV]], [[END]]
+; CHECK-NEXT:    br i1 [[EXIT]], label [[FOR_COND_CLEANUP_LOOPEXIT:%.*]], label [[CONT23:%.*]]
+; CHECK:       for.cond.cleanup.loopexit:
+; CHECK-NEXT:    ret void
+; CHECK:       cont23:
+; CHECK-NEXT:    [[EXITCOND241:%.*]] = icmp eq i64 [[IV]], 2147483647
+; CHECK-NEXT:    br i1 [[EXITCOND241]], label [[HANDLER_ADD_OVERFLOW:%.*]], label [[BACKEDGE]]
+; CHECK:       handler.add_overflow:
+; CHECK-NEXT:    unreachable
+; CHECK:       backedge:
+; CHECK-NEXT:    [[IV_NEW]] = add i64 [[IV]], 1
+; CHECK-NEXT:    br label [[LOOPHEADER]], !llvm.loop [[LOOP0:![0-9]+]]
+;
+entry:
+  br label %loopheader
+
+loopheader:
+  %iv = phi i64 [ 0, %entry ], [ %iv_new, %backedge ]
+  %exit = icmp eq i64 %iv, %end
+  br i1 %exit, label %for.cond.cleanup.loopexit, label %cont23
+
+for.cond.cleanup.loopexit:
+  ret void
+
+cont23:
+  %exitcond241 = icmp eq i64 %iv, 2147483647
+  br i1 %exitcond241, label %handler.add_overflow, label %backedge
+
+handler.add_overflow:
+  unreachable
+
+backedge: ; preds = %cont23
+  %iv_new = add i64 %iv, 1
+  br label %loopheader, !llvm.loop !0
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.full"}
+;.
+; CHECK: [[LOOP0]] = distinct !{[[LOOP0]], [[META1:![0-9]+]]}
+; CHECK: [[META1]] = !{!"llvm.loop.unroll.full"}
+;.