[llvm] [NVPTX] Check Before inserting AddrSpaceCastInst in NVPTXLoweringAlloca (PR #106127)

[weiwei chen via llvm-commits]

https://github.com/weiweichen updated https://github.com/llvm/llvm-project/pull/106127

>From 2506e94c21e3c61a584332b34a9f0aea4fb17a48 Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Mon, 26 Aug 2024 15:28:38 -0400
Subject: [PATCH 1/3] Only insert AddrSpaceCastInst when necessary.

---
 llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
index 369238436083c7..62b789c008bb43 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
@@ -72,12 +72,21 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
         Changed = true;
         auto ETy = allocaInst->getAllocatedType();
         auto LocalAddrTy = PointerType::get(ETy, ADDRESS_SPACE_LOCAL);
-        auto NewASCToLocal = new AddrSpaceCastInst(allocaInst, LocalAddrTy, "");
-        auto GenericAddrTy = PointerType::get(ETy, ADDRESS_SPACE_GENERIC);
-        auto NewASCToGeneric =
-            new AddrSpaceCastInst(NewASCToLocal, GenericAddrTy, "");
-        NewASCToLocal->insertAfter(allocaInst);
-        NewASCToGeneric->insertAfter(NewASCToLocal);
+        PointerType *allocInstPtrTy = dyn_cast_or_null<PointerType>(allocaInst->getType()->getScalarType());
+        assert(allocInstPtrTy && "AllocInst scalar type is not a PointerType.");
+        Instruction* NewASCToGeneric = allocaInst;
+        if(allocInstPtrTy->getAddressSpace() != ADDRESS_SPACE_LOCAL) {
+          // Only insert a new AddrSpaceCastInst if
+          // allocaInst is not already in ADDRESS_SPACE_LOCAL.
+          auto NewASCToLocal =
+              new AddrSpaceCastInst(allocaInst, LocalAddrTy, "");
+          auto GenericAddrTy = PointerType::get(ETy, ADDRESS_SPACE_GENERIC);
+          NewASCToGeneric =
+              new AddrSpaceCastInst(NewASCToLocal, GenericAddrTy, "");
+          NewASCToLocal->insertAfter(allocaInst);
+          NewASCToGeneric->insertAfter(NewASCToLocal);
+        }
+
         for (Use &AllocaUse : llvm::make_early_inc_range(allocaInst->uses())) {
           // Check Load, Store, GEP, and BitCast Uses on alloca and make them
           // use the converted generic address, in order to expose non-generic

>From 0c9a46d59bf1df9798e2e868344db3a93ae83d49 Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Mon, 26 Aug 2024 16:06:47 -0400
Subject: [PATCH 2/3] Add unittest.

---
 llvm/test/CodeGen/NVPTX/lower-alloca.ll | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/llvm/test/CodeGen/NVPTX/lower-alloca.ll b/llvm/test/CodeGen/NVPTX/lower-alloca.ll
index b1c34c8b5ecd78..68793dffb01e8a 100644
--- a/llvm/test/CodeGen/NVPTX/lower-alloca.ll
+++ b/llvm/test/CodeGen/NVPTX/lower-alloca.ll
@@ -17,7 +17,21 @@ define void @kernel() {
   ret void
 }
 
+define void @alloc_already_in_addrspace5() {
+; LABEL: @lower_alloca_addrspace5
+; PTX-LABEL: .visible .func alloc_already_in_addrspace5(
+  %A = alloca i32, addrspace(5)
+; CHECK-NOT: addrspacecast ptr %A to ptr addrspace(5)
+; CHECK: store i32 0, ptr addrspace(5) {{%.+}}
+; PTX: st.local.u32 [%SP+0], {{%r[0-9]+}}
+  store i32 0, ptr addrspace(5) %A
+  call void @callee(ptr addrspace(5) %A)
+  ret void
+}
+
 declare void @callee(ptr)
+declare void @callee_addrspace5(ptr addrspace(5))
 
 !nvvm.annotations = !{!0}
 !0 = !{ptr @kernel, !"kernel", i32 1}
+!1 = !{ptr @alloc_already_in_addrspace5, !"alloc_already_in_addrspace5", i32 1}

>From 6ec68a5bc9e8218d5ae6438408618d13fcf23328 Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Mon, 26 Aug 2024 16:18:51 -0400
Subject: [PATCH 3/3] Formattttttttt.

---
 llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
index 62b789c008bb43..69803ddfc6f920 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
@@ -72,10 +72,11 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
         Changed = true;
         auto ETy = allocaInst->getAllocatedType();
         auto LocalAddrTy = PointerType::get(ETy, ADDRESS_SPACE_LOCAL);
-        PointerType *allocInstPtrTy = dyn_cast_or_null<PointerType>(allocaInst->getType()->getScalarType());
-        assert(allocInstPtrTy && "AllocInst scalar type is not a PointerType.");
-        Instruction* NewASCToGeneric = allocaInst;
-        if(allocInstPtrTy->getAddressSpace() != ADDRESS_SPACE_LOCAL) {
+        PointerType *AllocInstPtrTy = dyn_cast_or_null<PointerType>(
+            allocaInst->getType()->getScalarType());
+        assert(AllocInstPtrTy && "AllocInst scalar type is not a PointerType.");
+        Instruction *NewASCToGeneric = allocaInst;
+        if (AllocInstPtrTy->getAddressSpace() != ADDRESS_SPACE_LOCAL) {
           // Only insert a new AddrSpaceCastInst if
           // allocaInst is not already in ADDRESS_SPACE_LOCAL.
           auto NewASCToLocal =