From c56585164b1b93717452959aad50cd3138ad8fc3 Mon Sep 17 00:00:00 2001
From: Joshua Ferguson <joshua.ferguson.273@gmail.com>
Date: Sun, 5 May 2024 00:07:26 -0400
Subject: [PATCH 1/3] adding AssumeLoopExists flag

---
 llvm/include/llvm/Analysis/ScalarEvolution.h | 3 +++
 llvm/lib/Analysis/ScalarEvolution.cpp        | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Analysis/ScalarEvolution.h b/llvm/include/llvm/Analysis/ScalarEvolution.h
index 5828cc156cc78..0399499a6224e 100644
--- a/llvm/include/llvm/Analysis/ScalarEvolution.h
+++ b/llvm/include/llvm/Analysis/ScalarEvolution.h
@@ -460,6 +460,9 @@ class ScalarEvolution {
     LoopComputable ///< The SCEV varies predictably with the loop.
   };
 
+  bool AssumeLoopFinite = false;
+  void setAssumeLoopExits();
+
   /// An enum describing the relationship between a SCEV and a basic block.
   enum BlockDisposition {
     DoesNotDominateBlock,  ///< The SCEV does not dominate the block.
diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 93f885c5d5ad8..2f8bdb3ee366d 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -509,6 +509,8 @@ const SCEV *ScalarEvolution::getVScale(Type *Ty) {
   return S;
 }
 
+void ScalarEvolution::setAssumeLoopExits() { this->AssumeLoopFinite = true; }
+
 const SCEV *ScalarEvolution::getElementCount(Type *Ty, ElementCount EC) {
   const SCEV *Res = getConstant(Ty, EC.getKnownMinValue());
   if (EC.isScalable())
@@ -7422,7 +7424,8 @@ bool ScalarEvolution::loopIsFiniteByAssumption(const Loop *L) {
   // A mustprogress loop without side effects must be finite.
   // TODO: The check used here is very conservative.  It's only *specific*
   // side effects which are well defined in infinite loops.
-  return isFinite(L) || (isMustProgress(L) && loopHasNoSideEffects(L));
+  return AssumeLoopFinite || isFinite(L) ||
+         (isMustProgress(L) && loopHasNoSideEffects(L));
 }
 
 const SCEV *ScalarEvolution::createSCEVIter(Value *V) {

From 0dfb651ca32d2e22ced31fa0416d2d4ff2b395c7 Mon Sep 17 00:00:00 2001
From: Manuel Drehwald <git@manuel.drehwald.info>
Date: Sun, 5 May 2024 00:08:54 -0400
Subject: [PATCH 2/3] add AssumeLoopExists test

---
 .../Analysis/ScalarEvolutionTest.cpp          | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index a7b3c5c404ab7..24c07096b12b3 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1085,6 +1085,36 @@ TEST_F(ScalarEvolutionsTest, SCEVComputeExpressionSize) {
   EXPECT_EQ(S2S->getExpressionSize(), 5u);
 }
 
+TEST_F(ScalarEvolutionsTest, AssumeLoopExists) {
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "define void @foo(i32 %N) { "
+      "entry: "
+      "  %cmp3 = icmp sgt i32 %N, 0 "
+      "  br i1 %cmp3, label %for.body, label %for.cond.cleanup "
+      "for.cond.cleanup: "
+      "  ret void "
+      "for.body: "
+      "  br label %for.body "
+      "} "
+      "declare i32 @llvm.loop.decrement.reg.i32.i32.i32(i32, i32) ",
+      Err, C);
+
+  ASSERT_TRUE(M && "Could not parse module?");
+  ASSERT_TRUE(!verifyModule(*M) && "Must have been well formed!");
+
+  runWithSE(*M, "foo", [&](Function &F, LoopInfo &LI, ScalarEvolution &SE) {
+    BasicBlock *L = F.begin()->getNextNode()->getNextNode();
+    auto *Loop = LI.getLoopFor(L);
+    bool IsFinite = SE.loopIsFiniteByAssumption(Loop);
+    EXPECT_FALSE(IsFinite);
+    SE.setAssumeLoopExits();
+    IsFinite = SE.loopIsFiniteByAssumption(Loop);
+    EXPECT_TRUE(IsFinite);
+  });
+}
+
 TEST_F(ScalarEvolutionsTest, SCEVLoopDecIntrinsic) {
   LLVMContext C;
   SMDiagnostic Err;

From f10bf0934cbbbdb72c7009b1b5f6b0d7d09ec678 Mon Sep 17 00:00:00 2001
From: Manuel Drehwald <git@manuel.drehwald.info>
Date: Sun, 5 May 2024 00:22:38 -0400
Subject: [PATCH 3/3] cleanup

---
 llvm/unittests/Analysis/ScalarEvolutionTest.cpp | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
index 24c07096b12b3..f74584636bd15 100644
--- a/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
+++ b/llvm/unittests/Analysis/ScalarEvolutionTest.cpp
@@ -1098,7 +1098,6 @@ TEST_F(ScalarEvolutionsTest, AssumeLoopExists) {
       "for.body: "
       "  br label %for.body "
       "} "
-      "declare i32 @llvm.loop.decrement.reg.i32.i32.i32(i32, i32) ",
       Err, C);
 
   ASSERT_TRUE(M && "Could not parse module?");