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

[via llvm-commits]

https://github.com/modiking created https://github.com/llvm/llvm-project/pull/78648

None

>From ac28ee921ba6d77647277709774201357effe8bb Mon Sep 17 00:00:00 2001
From: modimo <modimo at fb.com>
Date: Thu, 18 Jan 2024 16:06:43 -0800
Subject: [PATCH] [LoopUnroll] Introduce UnrollMaxIterations as a hard cap on
 how many iterations we try to unroll

---
 llvm/lib/Transforms/Utils/LoopUnroll.cpp   | 12 +++++++
 llvm/test/Transforms/LoopUnroll/pr77842.ll | 37 ++++++++++++++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 llvm/test/Transforms/LoopUnroll/pr77842.ll

diff --git a/llvm/lib/Transforms/Utils/LoopUnroll.cpp b/llvm/lib/Transforms/Utils/LoopUnroll.cpp
index ee6f7b35750af0..8529fa1db18e18 100644
--- a/llvm/lib/Transforms/Utils/LoopUnroll.cpp
+++ b/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) {
diff --git a/llvm/test/Transforms/LoopUnroll/pr77842.ll b/llvm/test/Transforms/LoopUnroll/pr77842.ll
new file mode 100644
index 00000000000000..834033bbe3618e
--- /dev/null
+++ b/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"}