[llvm] [InstCombine] Remove Store which has one-use AddrSpaceCastInst as Ptr (#68120) (PR #79565)

[via llvm-commits]

https://github.com/ParkHanbum updated https://github.com/llvm/llvm-project/pull/79565

>From 62e63cd0e5d0f60bd4fd5d614cb87b1f16a2c5a7 Mon Sep 17 00:00:00 2001
From: Hanbum Park <kese111 at gmail.com>
Date: Fri, 26 Jan 2024 18:30:47 +0900
Subject: [PATCH] [InstCombine] Remove Store which has one-use
 AddrSpaceCastInst as Ptr (#68120)

If the `store` target is OneUse and the user is `AddrSpaceCastInst`,
proper processing was not performed. Therefore, when processing
`Alloca` in the second process, this storage space was deleted and
a problem occurred.

This patch removes the `store` when the store target is OneUse and
the user is `AddrSpaceCastInst`.
---
 .../InstCombine/InstCombineLoadStoreAlloca.cpp   |  8 ++++----
 llvm/test/Transforms/InstCombine/store.ll        | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
index bb2a77daa60a764..a7c60852bc49e1b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp
@@ -1385,11 +1385,11 @@ Instruction *InstCombinerImpl::visitStoreInst(StoreInst &SI) {
   if (Ptr->hasOneUse()) {
     if (isa<AllocaInst>(Ptr))
       return eraseInstFromFunction(SI);
-    if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Ptr)) {
-      if (isa<AllocaInst>(GEP->getOperand(0))) {
-        if (GEP->getOperand(0)->hasOneUse())
+    if (isa<GetElementPtrInst>(Ptr) || isa<AddrSpaceCastInst>(Ptr)) {
+      Instruction *PtrI = cast<Instruction>(Ptr);
+      if (isa<AllocaInst>(PtrI->getOperand(0)))
+        if (PtrI->getOperand(0)->hasOneUse())
           return eraseInstFromFunction(SI);
-      }
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/store.ll b/llvm/test/Transforms/InstCombine/store.ll
index 673395464c85aaf..969d766a2d1e2bf 100644
--- a/llvm/test/Transforms/InstCombine/store.ll
+++ b/llvm/test/Transforms/InstCombine/store.ll
@@ -345,6 +345,22 @@ define void @store_to_readonly_noalias(ptr readonly noalias %0) {
   ret void
 }
 
+; if ptr which store to is addrspacecast, and its target is alloca. remove it.
+define void @src_store_oneuse_addrspaceast_alloca(ptr align 8 %arg) {
+; CHECK-LABEL: @src_store_oneuse_addrspaceast_alloca(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    store ptr poison, ptr null, align 8
+; CHECK-NEXT:    ret void
+;
+bb:
+  %i = alloca ptr, align 8, addrspace(5)
+  %i1 = addrspacecast ptr addrspace(5) %i to ptr
+  store ptr %arg, ptr %i1, align 8
+  %i2 = load ptr, ptr %i1, align 8
+  store ptr %i2, ptr null, align 8
+  ret void
+}
+
 !0 = !{!4, !4, i64 0}
 !1 = !{!"omnipotent char", !2}
 !2 = !{!"Simple C/C++ TBAA"}