[SROA] Don't speculate phis with different load user types

Fixes an SROA crash.

Fallout from opaque pointers since with typed pointers we'd bail out at the bitcast.

Reviewed By: nikic

Differential Revision: https://reviews.llvm.org/D136119

(cherry picked from commit 6219ec07c6f8d1ead51beca7cf21fbf2323c51d7)

2 files changed, 53 insertions, 7 deletions

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 143a035749c7..644c5c82e58e 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -1210,8 +1210,7 @@ static bool isSafePHIToSpeculate(PHINode &PN) {
   BasicBlock *BB = PN.getParent();
   Align MaxAlign;
   uint64_t APWidth = DL.getIndexTypeSizeInBits(PN.getType());
-  APInt MaxSize(APWidth, 0);
-  bool HaveLoad = false;
+  Type *LoadType = nullptr;
   for (User *U : PN.users()) {
     LoadInst *LI = dyn_cast<LoadInst>(U);
     if (!LI || !LI->isSimple())
@@ -1223,21 +1222,27 @@ static bool isSafePHIToSpeculate(PHINode &PN) {
     if (LI->getParent() != BB)
       return false;
 
+    if (LoadType) {
+      if (LoadType != LI->getType())
+        return false;
+    } else {
+      LoadType = LI->getType();
+    }
+
     // Ensure that there are no instructions between the PHI and the load that
     // could store.
     for (BasicBlock::iterator BBI(PN); &*BBI != LI; ++BBI)
       if (BBI->mayWriteToMemory())
         return false;
 
-    uint64_t Size = DL.getTypeStoreSize(LI->getType()).getFixedSize();
     MaxAlign = std::max(MaxAlign, LI->getAlign());
-    MaxSize = MaxSize.ult(Size) ? APInt(APWidth, Size) : MaxSize;
-    HaveLoad = true;
   }
 
-  if (!HaveLoad)
+  if (!LoadType)
     return false;
 
+  APInt LoadSize = APInt(APWidth, DL.getTypeStoreSize(LoadType).getFixedSize());
+
   // We can only transform this if it is safe to push the loads into the
   // predecessor blocks. The only thing to watch out for is that we can't put
   // a possibly trapping load in the predecessor if it is a critical edge.
@@ -1259,7 +1264,7 @@ static bool isSafePHIToSpeculate(PHINode &PN) {
     // If this pointer is always safe to load, or if we can prove that there
     // is already a load in the block, then we can move the load to the pred
     // block.
-    if (isSafeToLoadUnconditionally(InVal, MaxAlign, MaxSize, DL, TI))
+    if (isSafeToLoadUnconditionally(InVal, MaxAlign, LoadSize, DL, TI))
       continue;
 
     return false;
diff --git a/llvm/test/Transforms/SROA/phi-speculate-different-load-types.ll b/llvm/test/Transforms/SROA/phi-speculate-different-load-types.ll
new file mode 100644
index 000000000000..7893da340736
--- /dev/null
+++ b/llvm/test/Transforms/SROA/phi-speculate-different-load-types.ll
@@ -0,0 +1,41 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=sroa < %s -S | FileCheck %s
+
+define void @f(i1 %i) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:    [[A1:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    [[A2:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    br i1 [[I:%.*]], label [[BB1:%.*]], label [[BB:%.*]]
+; CHECK:       bb:
+; CHECK-NEXT:    br label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    br label [[BB2]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[TMP3:%.*]] = phi ptr [ [[A1]], [[BB1]] ], [ [[A2]], [[BB]] ]
+; CHECK-NEXT:    [[TMP5:%.*]] = load i32, ptr [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP4:%.*]] = load i64, ptr [[TMP3]], align 4
+; CHECK-NEXT:    call void @use32(i32 [[TMP5]])
+; CHECK-NEXT:    call void @use64(i64 [[TMP4]])
+; CHECK-NEXT:    ret void
+;
+  %a1 = alloca i64
+  %a2 = alloca i64
+  br i1 %i, label %bb1, label %bb
+
+bb:
+  br label %bb2
+
+bb1:
+  br label %bb2
+
+bb2:
+  %tmp3 = phi ptr [ %a1, %bb1 ], [ %a2, %bb ]
+  %tmp5 = load i32, ptr %tmp3
+  %tmp4 = load i64, ptr %tmp3
+  call void @use32(i32 %tmp5)
+  call void @use64(i64 %tmp4)
+  ret void
+}
+
+declare void @use32(i32)
+declare void @use64(i64)