From 58b9666dc1a09819377ae585b4718ef2ea951970 Mon Sep 17 00:00:00 2001
From: Zaara Syeda <syzaara@ca.ibm.com>
Date: Thu, 7 Jul 2022 14:29:06 -0400
Subject: [PATCH] [LSR] Fix bug - check if loop has preheader before calling
 isInductionPHI

Fix bug exposed by https://reviews.llvm.org/D125990
rewriteLoopExitValues calls InductionDescriptor::isInductionPHI which requires
the PHI node to have an incoming edge from the loop preheader. This adds checks
before calling InductionDescriptor::isInductionPHI to see that the loop has a
preheader. Also did some refactoring.

Differential Revision: https://reviews.llvm.org/D129297
---
 llvm/lib/Transforms/Utils/LoopUtils.cpp       | 25 ++++++++++++-------
 .../LoopStrengthReduce/remove_scev_indvars.ll | 21 ++++++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/LoopUtils.cpp b/llvm/lib/Transforms/Utils/LoopUtils.cpp
index afc7c61a5621..29d3cc712f98 100644
--- a/llvm/lib/Transforms/Utils/LoopUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopUtils.cpp
@@ -1243,6 +1243,20 @@ static bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet)
   return true;
 }
 
+/// Checks if it is safe to call InductionDescriptor::isInductionPHI for \p Phi,
+/// and returns true if this Phi is an induction phi in the loop. When
+/// isInductionPHI returns true, \p ID will be also be set by isInductionPHI.
+static bool checkIsIndPhi(PHINode *Phi, Loop *L, ScalarEvolution *SE,
+                          InductionDescriptor &ID) {
+  if (!Phi)
+    return false;
+  if (!L->getLoopPreheader())
+    return false;
+  if (Phi->getParent() != L->getHeader())
+    return false;
+  return InductionDescriptor::isInductionPHI(Phi, L, SE, ID);
+}
+
 int llvm::rewriteLoopExitValues(Loop *L, LoopInfo *LI, TargetLibraryInfo *TLI,
                                 ScalarEvolution *SE,
                                 const TargetTransformInfo *TTI,
@@ -1303,10 +1317,7 @@ int llvm::rewriteLoopExitValues(Loop *L, LoopInfo *LI, TargetLibraryInfo *TLI,
           InductionDescriptor ID;
           PHINode *IndPhi = dyn_cast<PHINode>(Inst);
           if (IndPhi) {
-            if (IndPhi->getParent() != L->getHeader())
-              continue;
-            // Do not consider non induction phis.
-            if (!InductionDescriptor::isInductionPHI(IndPhi, L, SE, ID))
+            if (!checkIsIndPhi(IndPhi, L, SE, ID))
               continue;
             // This is an induction PHI. Check that the only users are PHI
             // nodes, and induction variable update binary operators.
@@ -1327,12 +1338,8 @@ int llvm::rewriteLoopExitValues(Loop *L, LoopInfo *LI, TargetLibraryInfo *TLI,
               continue;
             if (llvm::any_of(Inst->users(), [&](User *U) {
                   PHINode *Phi = dyn_cast<PHINode>(U);
-                  if (!Phi)
+                  if (Phi != PN && !checkIsIndPhi(Phi, L, SE, ID))
                     return true;
-                  if (Phi->getParent() == L->getHeader()) {
-                    if (!InductionDescriptor::isInductionPHI(Phi, L, SE, ID))
-                      return true;
-                  }
                   return false;
                 }))
               continue;
diff --git a/llvm/test/Transforms/LoopStrengthReduce/remove_scev_indvars.ll b/llvm/test/Transforms/LoopStrengthReduce/remove_scev_indvars.ll
index 37c7d294ad05..f702732661f6 100644
--- a/llvm/test/Transforms/LoopStrengthReduce/remove_scev_indvars.ll
+++ b/llvm/test/Transforms/LoopStrengthReduce/remove_scev_indvars.ll
@@ -91,3 +91,24 @@ for.end6730:                                      ; preds = %for.cond6418
 dead:                                             ; No predecessors!
   br label %for.cond6403
 }
+
+
+; Check that this doesn't crash
+define void @kernfs_path_from_node() {
+entry:
+  callbr void asm sideeffect "", "i"(i8* blockaddress(@kernfs_path_from_node, %while.body))
+          to label %asm.fallthrough [label %while.body]
+
+asm.fallthrough:                                  ; preds = %entry
+  br label %while.body
+
+while.body:                                       ; preds = %while.body, %asm.fallthrough, %entry
+  %depth.04 = phi i32 [ %inc, %while.body ], [ 0, %asm.fallthrough ], [ 0, %entry ]
+  %inc = add i32 %depth.04, 1
+  br i1 false, label %while.end, label %while.body
+
+while.end:                                        ; preds = %while.body
+  %inc.lcssa = phi i32 [ %depth.04, %while.body ]
+  store i32 %inc.lcssa, i32* null, align 4
+  ret void
+}