[llvm] [SCEVExpander] Relax hoisting condition for AddRec start (PR #75916)

[via llvm-commits]

https://github.com/averine created https://github.com/llvm/llvm-project/pull/75916

Don't require the existence of a preheader to hoist an AddRec start value: a unique loop predecessor which isLegalToHoistInto should be enough.

>From 3350c5ec136b406a717d00b2482ddf8cba478e6d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexandre=20V=C3=A9rine?= <“verine.alexandre at proton.me”>
Date: Tue, 19 Dec 2023 10:23:08 +0100
Subject: [PATCH] [SCEVExpander] Relax hoisting condition for AddRec start

Don't require the existence of a preheader to hoist an
AddRec start value: a unique loop predecessor which
isLegalToHoistInto should be enough.
---
 .../Utils/ScalarEvolutionExpander.cpp         | 13 ++++---
 .../scevexpander-in-non-simplified-loop.ll    | 39 +++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/IndVarSimplify/scevexpander-in-non-simplified-loop.ll

diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
index cd3ac317cd238e..f4be4c94925e62 100644
--- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
+++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
@@ -980,11 +980,14 @@ SCEVExpander::getAddRecExprPHILiterally(const SCEVAddRecExpr *Normalized,
   PostIncLoopSet SavedPostIncLoops = PostIncLoops;
   PostIncLoops.clear();
 
-  // Expand code for the start value into the loop preheader.
-  assert(L->getLoopPreheader() &&
-         "Can't expand add recurrences without a loop preheader!");
-  Value *StartV =
-      expand(Normalized->getStart(), L->getLoopPreheader()->getTerminator());
+  // Expand code for the start value into the loop predecessor. The loop is not
+  // necessarily in Loop Simplify Form, so assert it's legal to do so.
+  auto *LP = L->getLoopPredecessor();
+  assert(LP && LP->isLegalToHoistInto() &&
+         "Can't expand add recurrences without hoisting into the loop "
+         "predecessor.");
+
+  Value *StartV = expand(Normalized->getStart(), LP->getTerminator());
 
   // StartV must have been be inserted into L's preheader to dominate the new
   // phi.
diff --git a/llvm/test/Transforms/IndVarSimplify/scevexpander-in-non-simplified-loop.ll b/llvm/test/Transforms/IndVarSimplify/scevexpander-in-non-simplified-loop.ll
new file mode 100644
index 00000000000000..7b546ab7d5d329
--- /dev/null
+++ b/llvm/test/Transforms/IndVarSimplify/scevexpander-in-non-simplified-loop.ll
@@ -0,0 +1,39 @@
+; RUN: opt -passes=indvars < %s
+
+target datalayout = "e-n32:64"
+
+declare void @g(i64)
+
+define void @f(ptr %block_address_arg)  {
+
+; This loop isn't in Simplified Form. Ensure SCEV expander doesn't hit an assert
+; when visiting it.
+f_entry:
+  indirectbr ptr %block_address_arg, [label %f.exit, label %f_for.cond]
+
+f_for.cond:                       ; preds = %f_for.body, %f_entry
+  %i = phi i32 [ %j, %f_for.body ], [ 0, %f_entry ]
+  %cmp = icmp ult i32 %i, 64
+  br i1 %cmp, label %f_for.body, label %f_for.end
+
+f_for.body:                       ; preds = %f_for.cond
+  %j = add nuw nsw i32 %i, 1
+  br label %f_for.cond
+
+; Indvars pass visits this loop, since it's in Simplified Form.
+; Because of the integer extension in that loop's body, SCEV expander may also
+; visit the first loop.
+f_for.end:                        ; preds = %f_for.cond
+  %k = phi i32 [ %i, %f_for.cond ]
+  br label %f_for2.body
+
+f_for2.body:                       ; preds = %f_for2.body, %f_for.end
+  %l = phi i32 [ %k, %f_for.end ], [ %n, %f_for2.body ]
+  %m = zext i32 %l to i64
+  call void @g(i64 %m)
+  %n = add nuw nsw i32 %l, 1
+  br label %f_for2.body
+
+f.exit:                                ; preds = %f_entry
+  ret void
+}