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

weiwei chen via llvm-commits llvm-commits at lists.llvm.org
Tue Sep 3 07:58:18 PDT 2024


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/6] 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/6] 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/6] 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 =

>From 1bdc7fa54436fc03a0f273a8389960a5c23d29d2 Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Mon, 26 Aug 2024 17:21:45 -0400
Subject: [PATCH 4/6] Address review comment.

---
 llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
index 69803ddfc6f920..ec3743f8f310cf 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
@@ -72,9 +72,8 @@ 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>(
+        PointerType *AllocInstPtrTy = cast<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

>From f1680bc0f13d1ca3eee9fe80d6ea651a71d9f93b Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Mon, 26 Aug 2024 17:36:23 -0400
Subject: [PATCH 5/6] Formatting again, ahhhhh.

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

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
index ec3743f8f310cf..b0893c0ad6e2b4 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
@@ -72,8 +72,8 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
         Changed = true;
         auto ETy = allocaInst->getAllocatedType();
         auto LocalAddrTy = PointerType::get(ETy, ADDRESS_SPACE_LOCAL);
-        PointerType *AllocInstPtrTy = cast<PointerType>(
-            allocaInst->getType()->getScalarType());
+        PointerType *AllocInstPtrTy =
+            cast<PointerType>(allocaInst->getType()->getScalarType());
         Instruction *NewASCToGeneric = allocaInst;
         if (AllocInstPtrTy->getAddressSpace() != ADDRESS_SPACE_LOCAL) {
           // Only insert a new AddrSpaceCastInst if

>From 0cee1cd529760c5d2a619666bea21d78662526cc Mon Sep 17 00:00:00 2001
From: Weiwei Chen <weiwei.chen at modular.com>
Date: Tue, 3 Sep 2024 10:57:49 -0400
Subject: [PATCH 6/6] Add early assert for alloc address space check.

---
 llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
index b0893c0ad6e2b4..7e8bb60af12881 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerAlloca.cpp
@@ -75,7 +75,12 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
         PointerType *AllocInstPtrTy =
             cast<PointerType>(allocaInst->getType()->getScalarType());
         Instruction *NewASCToGeneric = allocaInst;
-        if (AllocInstPtrTy->getAddressSpace() != ADDRESS_SPACE_LOCAL) {
+        unsigned AllocAddrSpace = AllocInstPtrTy->getAddressSpace();
+        assert((AllocAddrSpace == ADDRESS_SPACE_GENERIC ||
+                AllocAddrSpace == ADDRESS_SPACE_LOCAL) &&
+               "AllocaInst can only be in Generic or Local address space for "
+               "NVPTX.");
+        if (AllocAddrSpace != ADDRESS_SPACE_LOCAL) {
           // Only insert a new AddrSpaceCastInst if
           // allocaInst is not already in ADDRESS_SPACE_LOCAL.
           auto NewASCToLocal =



More information about the llvm-commits mailing list