[llvm] ce2d69b - [SROA][Mem2Reg] Do not crash on alloca + addrspacecast

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2020-07-24T15:15:38-05:00
New Revision: ce2d69b5577314ee97504876381273bbcddca3c5

URL: https://github.com/llvm/llvm-project/commit/ce2d69b5577314ee97504876381273bbcddca3c5
DIFF: https://github.com/llvm/llvm-project/commit/ce2d69b5577314ee97504876381273bbcddca3c5.diff

LOG: [SROA][Mem2Reg] Do not crash on alloca + addrspacecast

SROA knows that it can look through addrspacecast but
PromoteMemoryToRegister did not handle them. This caused an assertion
error for the test case, exposed while running
`Transforms/PhaseOrdering/inlining-alignment-assumptions.ll` with D83978
applied.

Reviewed By: arsenm

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

Added: 
    llvm/test/Transforms/Mem2Reg/alloca_addrspace.ll

Modified: 
    llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
    llvm/test/Transforms/SROA/alloca-address-space.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
index c7e9c919ec47..f079f81a6e8f 100644
--- a/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
+++ b/llvm/lib/Transforms/Utils/PromoteMemoryToRegister.cpp
@@ -95,6 +95,9 @@ bool llvm::isAllocaPromotable(const AllocaInst *AI) {
         return false;
       if (!onlyUsedByLifetimeMarkers(GEPI))
         return false;
+    } else if (const AddrSpaceCastInst *ASCI = dyn_cast<AddrSpaceCastInst>(U)) {
+      if (!onlyUsedByLifetimeMarkers(ASCI))
+        return false;
     } else {
       return false;
     }

diff  --git a/llvm/test/Transforms/Mem2Reg/alloca_addrspace.ll b/llvm/test/Transforms/Mem2Reg/alloca_addrspace.ll
new file mode 100644
index 000000000000..95db828fd117
--- /dev/null
+++ b/llvm/test/Transforms/Mem2Reg/alloca_addrspace.ll
@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -mem2reg -S | FileCheck %s
+; RUN: opt < %s -passes=mem2reg -S | FileCheck %s
+
+; Do not crash ;)
+
+target datalayout = "e-p:64:64-p5:32:32-A5"
+
+define amdgpu_kernel void @addressspace_alloca() {
+; CHECK-LABEL: @addressspace_alloca(
+; CHECK-NEXT:    ret void
+;
+  %alloca = alloca i8, align 8, addrspace(5)
+  %cast = addrspacecast i8 addrspace(5)* %alloca to i8*
+  call void @llvm.lifetime.start.p0i8(i64 2, i8* %cast)
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64 %size, i8* nocapture %ptr)

diff  --git a/llvm/test/Transforms/SROA/alloca-address-space.ll b/llvm/test/Transforms/SROA/alloca-address-space.ll
index c8e22f5b25d8..2ec59e66b9da 100644
--- a/llvm/test/Transforms/SROA/alloca-address-space.ll
+++ b/llvm/test/Transforms/SROA/alloca-address-space.ll
@@ -1,4 +1,5 @@
 ; RUN: opt < %s -sroa -S | FileCheck %s
+; RUN: opt < %s -passes=sroa -S | FileCheck %s
 target datalayout = "e-p:64:64:64-p1:16:16:16-p2:32:32-p3:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-n8:16:32:64-A2"
 
 declare void @llvm.memcpy.p2i8.p2i8.i32(i8 addrspace(2)* nocapture, i8 addrspace(2)* nocapture readonly, i32, i1)
@@ -127,3 +128,15 @@ define void @test_load_store_
diff _addr_space([2 x float] addrspace(1)* %complex1
   store i64 %v2, i64 addrspace(1)* %p2
   ret void
 }
+
+define void @addressspace_alloca_lifetime() {
+; CHECK-LABEL: @addressspace_alloca_lifetime(
+; CHECK-NEXT:    ret void
+;
+  %alloca = alloca i8, align 8, addrspace(2)
+  %cast = addrspacecast i8 addrspace(2)* %alloca to i8*
+  call void @llvm.lifetime.start.p0i8(i64 2, i8* %cast)
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64 %size, i8* nocapture %ptr)