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

[weiwei chen via llvm-commits]

================
@@ -72,12 +72,26 @@ 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 =
+            cast<PointerType>(allocaInst->getType()->getScalarType());
+        Instruction *NewASCToGeneric = allocaInst;
+        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 =
+              new AddrSpaceCastInst(allocaInst, LocalAddrTy, "");
+          auto GenericAddrTy = PointerType::get(ETy, ADDRESS_SPACE_GENERIC);
----------------
weiweichen wrote:

Perhaps it's a confusion from the [original test](https://github.com/llvm/llvm-project/blob/1324789a65665c27eda9e04bc93db81cc859924c/llvm/test/CodeGen/NVPTX/lower-alloca.ll#L8-L18) ⬇️ (as well as my explanation). Let me elaborate: 

```
define void @kernel() {
; LABEL: @lower_alloca
; PTX-LABEL: .visible .entry kernel(
  %A = alloca i32
; CHECK: addrspacecast ptr %A to ptr addrspace(5)
; CHECK: store i32 0, ptr addrspace(5) {{%.+}}
; PTX: st.local.u32 [{{%rd[0-9]+}}], {{%r[0-9]+}}
  store i32 0, ptr %A
  call void @callee(ptr %A)
  ret void
}
```

As shown above, FileCheck is checking that `store` has a `ptr addrsapce(5)`. This only happens if there is a chain of "asc(asc(alloca(AS(0)), to AS(5)), to AS(0)". Otherwise, the `store` here would have a `ptr` in generic space. 
Currently the test is [ran with](https://github.com/llvm/llvm-project/blob/1324789a65665c27eda9e04bc93db81cc859924c/llvm/test/CodeGen/NVPTX/lower-alloca.ll#L1C1-L1C76) `-infer-address-spaces`:
```
; RUN: opt < %s -S -nvptx-lower-alloca -infer-address-spaces | FileCheck %s
```

Here are the outputs with `main` with and without `-infer-address-space`

**without `-infer-address-space`**:
```
(autovenv) [M] weiwei.chen at Weiweis-MacBook-Pro ~/research/modularml/modular/third-party/llvm-project 🍔 opt < llvm/test/CodeGen/NVPTX/lower-alloca.ll -S -nvptx-lower-alloca 
; ModuleID = '<stdin>'
source_filename = "<stdin>"
target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
target triple = "nvptx64-unknown-unknown"

define void @kernel() {
  %A = alloca i32, align 4
  %1 = addrspacecast ptr %A to ptr addrspace(5)
  %2 = addrspacecast ptr addrspace(5) %1 to ptr
  store i32 0, ptr %2, align 4
  call void @callee(ptr %A)
  ret void
}

declare void @callee(ptr)

!nvvm.annotations = !{!0}

!0 = !{ptr @kernel, !"kernel", i32 1}
```

**with `-infer-address-space`** :
```
(autovenv) [M] weiwei.chen at Weiweis-MacBook-Pro ~/research/modularml/modular/third-party/llvm-project 🍔 opt < llvm/test/CodeGen/NVPTX/lower-alloca.ll -S -nvptx-lower-alloca -infer-address-spaces 
; ModuleID = '<stdin>'
source_filename = "<stdin>"
target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
target triple = "nvptx64-unknown-unknown"

define void @kernel() {
  %A = alloca i32, align 4
  %1 = addrspacecast ptr %A to ptr addrspace(5)
  store i32 0, ptr addrspace(5) %1, align 4
  call void @callee(ptr %A)
  ret void
}

declare void @callee(ptr)

!nvvm.annotations = !{!0}

!0 = !{ptr @kernel, !"kernel", i32 1}
```


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