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

[Alex MacLean via llvm-commits]

================
@@ -58,77 +64,55 @@ bool NVPTXLowerAlloca::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
 
-  bool Changed = false;
+  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),
+        "");
+    Cast->insertBefore(Alloca->getIterator());
+    NewAlloca->insertBefore(Cast->getIterator());
+    for (auto &U : llvm::make_early_inc_range(Alloca->uses())) {
+      auto *II = dyn_cast<IntrinsicInst>(U.getUser());
+      if (!II || (II->getIntrinsicID() != Intrinsic::lifetime_start &&
+                  II->getIntrinsicID() != Intrinsic::lifetime_end))
+        continue;
 
-        Instruction *AllocaInLocalAS = allocaInst;
-        auto ETy = allocaInst->getAllocatedType();
+      IRBuilder<> Builder(II);
+      Builder.CreateIntrinsic(II->getIntrinsicID(), {NewAlloca->getType()},
+                              {NewAlloca});
+      II->eraseFromParent();
+    }
 
-        // We need to make sure that LLVM has info that alloca needs to go to
-        // ADDRESS_SPACE_LOCAL for InferAddressSpace pass.
-        //
-        // For allocas in ADDRESS_SPACE_LOCAL, we add addrspacecast to
-        // ADDRESS_SPACE_LOCAL and back to ADDRESS_SPACE_GENERIC, so that
-        // the alloca's users still use a generic pointer to operate on.
-        //
-        // For allocas already in ADDRESS_SPACE_LOCAL, we just need
-        // addrspacecast to ADDRESS_SPACE_GENERIC.
-        if (AllocAddrSpace == ADDRESS_SPACE_GENERIC) {
-          auto ASCastToLocalAS = new AddrSpaceCastInst(
-              allocaInst,
-              PointerType::get(ETy->getContext(), ADDRESS_SPACE_LOCAL), "");
-          ASCastToLocalAS->insertAfter(allocaInst->getIterator());
-          AllocaInLocalAS = ASCastToLocalAS;
-        }
+    Alloca->replaceAllUsesWith(Cast);
+    Alloca->eraseFromParent();
+  }
+  return true;
+}
 
-        auto AllocaInGenericAS = new AddrSpaceCastInst(
-            AllocaInLocalAS,
-            PointerType::get(ETy->getContext(), ADDRESS_SPACE_GENERIC), "");
-        AllocaInGenericAS->insertAfter(AllocaInLocalAS->getIterator());
+bool NVPTXLowerAlloca::doInitialization(Module &M) {
+  const auto &DL = M.getDataLayout();
+  if (DL.getAllocaAddrSpace() == ADDRESS_SPACE_LOCAL)
+    return false;
+  auto DLStr = DL.getStringRepresentation();
 
-        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
-          // addrspacecast to NVPTXInferAddressSpaces. For other types
-          // of instructions this is unnecessary and may introduce redundant
-          // address cast.
-          auto LI = dyn_cast<LoadInst>(AllocaUse.getUser());
-          if (LI && LI->getPointerOperand() == allocaInst &&
-              !LI->isVolatile()) {
-            LI->setOperand(LI->getPointerOperandIndex(), AllocaInGenericAS);
-            continue;
-          }
-          auto SI = dyn_cast<StoreInst>(AllocaUse.getUser());
-          if (SI && SI->getPointerOperand() == allocaInst &&
-              !SI->isVolatile()) {
-            SI->setOperand(SI->getPointerOperandIndex(), AllocaInGenericAS);
-            continue;
-          }
-          auto GI = dyn_cast<GetElementPtrInst>(AllocaUse.getUser());
-          if (GI && GI->getPointerOperand() == allocaInst) {
-            GI->setOperand(GI->getPointerOperandIndex(), AllocaInGenericAS);
-            continue;
-          }
-          auto BI = dyn_cast<BitCastInst>(AllocaUse.getUser());
-          if (BI && BI->getOperand(0) == allocaInst) {
-            BI->setOperand(0, AllocaInGenericAS);
-            continue;
-          }
-        }
-      }
-    }
-  return Changed;
+  auto AddrSpaceStr = "A" + std::to_string(ADDRESS_SPACE_LOCAL);
+  assert(!StringRef(DLStr).contains("A") && "DataLayout should not contain A");
----------------
AlexMaclean wrote:

This seems like it should be a fatal error or diagnostic, not an assert.

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