[clang] [llvm] [NVPTX] Change the alloca address space in NVPTXLowerAlloca (PR #154814)

[Alex MacLean via llvm-commits]

================
@@ -51,86 +59,66 @@ char NVPTXLowerAlloca::ID = 1;
 INITIALIZE_PASS(NVPTXLowerAlloca, "nvptx-lower-alloca", "Lower Alloca", false,
                 false)
 
-// =============================================================================
-// Main function for this pass.
-// =============================================================================
-bool NVPTXLowerAlloca::runOnFunction(Function &F) {
-  if (skipFunction(F))
-    return false;
+bool NVPTXLowerAlloca::runOnModule(Module &M) {
+  bool Changed = changeDataLayout(M);
+  for (auto &F : M)
+    Changed |= lowerFunctionAllocas(F);
+  return Changed;
+}
 
-  bool Changed = false;
+bool NVPTXLowerAlloca::lowerFunctionAllocas(Function &F) {
+  SmallVector<AllocaInst *, 16> Allocas;
   for (auto &BB : F)
-    for (auto &I : BB) {
-      if (auto allocaInst = dyn_cast<AllocaInst>(&I)) {
-        Changed = true;
+    for (auto &I : BB)
+      if (auto *Alloca = dyn_cast<AllocaInst>(&I))
+        if (Alloca->getAddressSpace() != ADDRESS_SPACE_LOCAL)
+          Allocas.push_back(Alloca);
+
+  if (Allocas.empty())
+    return false;
 
-        PointerType *AllocInstPtrTy =
-            cast<PointerType>(allocaInst->getType()->getScalarType());
-        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.");
+  for (AllocaInst *Alloca : Allocas) {
+    auto *NewAlloca = new AllocaInst(
+        Alloca->getAllocatedType(), ADDRESS_SPACE_LOCAL, Alloca->getArraySize(),
+        Alloca->getAlign(), Alloca->getName());
+    auto *Cast = new AddrSpaceCastInst(
+        NewAlloca,
+        PointerType::get(Alloca->getAllocatedType()->getContext(),
+                         ADDRESS_SPACE_GENERIC),
+        "");
----------------
AlexMaclean wrote:

Nit: can we use an IRBuilder here?

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