From eae4f56cb40c5db0ec2c4e967107cba2473246f1 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks@google.com>
Date: Thu, 7 Mar 2024 07:22:33 +0000
Subject: [PATCH] [SROA] Fix phi gep unfolding with an alloca not in entry
 block

Fixes a crash reported in #83494.
---
 llvm/lib/Transforms/Scalar/SROA.cpp  |  8 +++++--
 llvm/test/Transforms/SROA/phi-gep.ll | 35 ++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index a575dc74d8db..e11b984f13bb 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -4062,10 +4062,14 @@ private:
   bool unfoldGEPPhi(GetElementPtrInst &GEPI) {
     // To prevent infinitely expanding recursive phis, bail if the GEP pointer
     // operand (looking through the phi if it is the phi we want to unfold) is
-    // an instruction besides an alloca.
+    // an instruction besides a static alloca.
     PHINode *Phi = dyn_cast<PHINode>(GEPI.getPointerOperand());
     auto IsInvalidPointerOperand = [](Value *V) {
-      return isa<Instruction>(V) && !isa<AllocaInst>(V);
+      if (!isa<Instruction>(V))
+        return false;
+      if (auto *AI = dyn_cast<AllocaInst>(V))
+        return !AI->isStaticAlloca();
+      return true;
     };
     if (Phi) {
       if (any_of(Phi->operands(), IsInvalidPointerOperand))
diff --git a/llvm/test/Transforms/SROA/phi-gep.ll b/llvm/test/Transforms/SROA/phi-gep.ll
index 33f42af69608..c2dfa7578eb4 100644
--- a/llvm/test/Transforms/SROA/phi-gep.ll
+++ b/llvm/test/Transforms/SROA/phi-gep.ll
@@ -638,6 +638,41 @@ bb3:
   ret i1 %icmp
 }
 
+define i32 @test_phi_mem2reg_alloca_not_in_entry_block(i1 %arg) {
+; CHECK-LABEL: @test_phi_mem2reg_alloca_not_in_entry_block(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    store i64 123, ptr [[ALLOCA]], align 4
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[ALLOCA2:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    store i64 124, ptr [[ALLOCA]], align 4
+; CHECK-NEXT:    br i1 [[ARG:%.*]], label [[BB3:%.*]], label [[BB4:%.*]]
+; CHECK:       bb3:
+; CHECK-NEXT:    br label [[BB4]]
+; CHECK:       bb4:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr [ [[ALLOCA]], [[BB2]] ], [ [[ALLOCA2]], [[BB3]] ]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, ptr [[PHI]], i64 1
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT:    ret i32 [[LOAD]]
+;
+bb:
+  %alloca = alloca i64
+  store i64 123, ptr %alloca
+  br label %bb2
+bb2:
+  %alloca2 = alloca i64
+  store i64 124, ptr %alloca
+  br i1 %arg, label %bb3, label %bb4
+bb3:
+  br label %bb4
+bb4:
+  %phi = phi ptr [ %alloca, %bb2 ], [ %alloca2, %bb3 ]
+  %gep = getelementptr i32, ptr %phi, i64 1
+  %load = load i32, ptr %gep
+  ret i32 %load
+}
+
 define i64 @test_unfold_phi_duplicate_phi_entry(ptr %arg, i8 %arg1, i1 %arg2) {
 ; CHECK-LABEL: @test_unfold_phi_duplicate_phi_entry(
 ; CHECK-NEXT:  bb: