[LoopUnroll] Introduce UnrollMaxIterations as a hard cap on how many iterations we try to unroll

Author: modiking

llvm/lib/Transforms/Utils/LoopUnroll.cpp

@@ -109,6 +109,10 @@ UnrollVerifyLoopInfo("unroll-verify-loopinfo", cl::Hidden,
 #endif
                     );
 
+static cl::opt<unsigned>
+    UnrollMaxIterations("unroll-max-iterations", cl::init(1'000'000),
+                        cl::Hidden,
+                        cl::desc("Maximum allowed iterations to unroll."));
 
 /// Check if unrolling created a situation where we need to insert phi nodes to
 /// preserve LCSSA form.
@@ -453,6 +457,14 @@ LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
     }
   }
 
+  // Certain cases with UBSAN can cause trip count to be calculated as INT_MAX,
+  // Block unrolling at a reasonable limit so that the compiler doesn't hang
+  // trying to unroll the loop. See PR77842
+  if (ULO.Count > UnrollMaxIterations) {
+    LLVM_DEBUG(dbgs() << "Won't unroll; trip count is too large\n");
+    return LoopUnrollResult::Unmodified;
+  }
+
   using namespace ore;
   // Report the unrolling decision.
   if (CompletelyUnroll) {

llvm/test/Transforms/LoopUnroll/pr77842.ll

@@ -0,0 +1,37 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=loop-unroll -disable-output -debug-only=loop-unroll %s 2>&1 | FileCheck %s
+
+; Validate that loop unroll full doesn't try to fully unroll values whose trip counts are too large.
+
+; CHECK: Exiting block %cont23: TripCount=2147483648, TripMultiple=0, BreakoutTrip=0
+; CHECK-NEXT: Won't unroll; trip count is too large
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-redhat-linux-gnu"
+
+define void @foo(i64 %end) {
+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"}