[llvm] [NVPTX] Rewrite kernel signatures in param AS (PR #204192)

[Akshay Deodhar via llvm-commits]

================
@@ -395,108 +268,165 @@ void copyByValParam(Function &F, Argument &Arg) {
       Arg.getParamAlign().value_or(DL.getPrefTypeAlign(ByValType)));
   Arg.replaceAllUsesWith(AllocA);
 
-  Value *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, Arg);
-
-  // Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
-  // addrspacecast preserves alignment.  Since params are constant, this load
+  // Be sure to propagate alignment to this copy; LLVM doesn't know that NVPTX
+  // addrspacecast preserves alignment.  Since params are constant, this copy
   // is definitely not volatile.
   const auto ArgSize = *AllocA->getAllocationSize(DL);
-  IRB.CreateMemCpy(AllocA, AllocA->getAlign(), ArgInParamAS, AllocA->getAlign(),
+  IRB.CreateMemCpy(AllocA, AllocA->getAlign(), &ParamPtr, AllocA->getAlign(),
                    ArgSize);
 }
 } // namespace
 
-static bool argIsProcessed(Argument *Arg) {
-  if (Arg->use_empty())
-    return true;
-
-  // If the argument is already wrapped, it was processed by this pass before.
-  if (Arg->hasOneUse())
-    if (const auto *II = dyn_cast<IntrinsicInst>(*Arg->user_begin()))
-      if (II->getIntrinsicID() == Intrinsic::nvvm_internal_addrspace_wrap)
-        return true;
-
-  return false;
+// Returns true if F has a byval argument not yet in the param address space.
+// Such arguments are lowered exactly once, so one already in param space means
+// the kernel has already been processed.
+static bool kernelNeedsByValLowering(const Function &F) {
+  return any_of(F.args(), [](const Argument &A) {
+    return A.hasByValAttr() &&
+           A.getType()->getPointerAddressSpace() != ADDRESS_SPACE_ENTRY_PARAM;
+  });
 }
 
-static void lowerKernelByValParam(Argument *Arg, Function &F,
-                                  const bool HasCvtaParam) {
+// Lower the uses of a single kernel byval argument. \p OldArg is the original
+// (generic) argument whose uses are being rewritten; \p NewParamArg is its
+// replacement, natively in the param address space.
+static void lowerKernelByValParam(Argument &OldArg, Argument &NewParamArg,
+                                  Function &F, const bool HasCvtaParam) {
   assert(isKernelFunction(F));
 
   const DataLayout &DL = F.getDataLayout();
   IRBuilder<> IRB(&F.getEntryBlock().front());
 
-  if (argIsProcessed(Arg))
+  if (OldArg.use_empty())
     return;
 
   // (1) First check the easy case, if were able to trace through all the uses
   // and we can convert them all to param AS, then we'll do this.
   ArgUseChecker AUC(DL);
-  ArgUseChecker::PtrInfo PI = AUC.visitArgPtr(*Arg);
+  ArgUseChecker::PtrInfo PI = AUC.visitArgPtr(OldArg);
   const bool ArgUseIsReadOnly = !(PI.isEscaped() || PI.isAborted());
   if (ArgUseIsReadOnly && AUC.Conditionals.empty()) {
     // Convert all loads and intermediate operations to use parameter AS and
     // skip creation of a local copy of the argument.
-    SmallVector<Use *, 16> UsesToUpdate(llvm::make_pointer_range(Arg->uses()));
-    Value *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, *Arg);
+    SmallVector<Use *, 16> UsesToUpdate(make_pointer_range(OldArg.uses()));
     for (Use *U : UsesToUpdate)
-      convertToParamAS(U, ArgInParamAS);
+      convertToParamAS(U, &NewParamArg);
+    // This path does not replaceAllUsesWith the old argument, so any debug-info
+    // uses would be left dangling and reset to poison when the old function is
+    // erased. Point them at the new param-space argument instead.
+    if (OldArg.isUsedByMetadata()) {
+      SmallVector<DbgVariableRecord *, 4> DbgUsers;
+      findDbgUsers(&OldArg, DbgUsers);
+      for (DbgVariableRecord *DVR : DbgUsers)
+        DVR->replaceVariableLocationOp(&OldArg, &NewParamArg);
+    }
     return;
   }
 
   // (2) If the argument is grid constant, we get to use the pointer directly.
-  if (HasCvtaParam && (ArgUseIsReadOnly || isParamGridConstant(*Arg))) {
-    LLVM_DEBUG(dbgs() << "Using non-copy pointer to " << *Arg << "\n");
-
-    // Cast argument to param address space. Because the backend will emit the
-    // argument already in the param address space, we need to use the noop
-    // intrinsic, this had the added benefit of preventing other optimizations
-    // from folding away this pair of addrspacecasts.
-    Instruction *ArgInParamAS = createNVVMInternalAddrspaceWrap(IRB, *Arg);
+  if (HasCvtaParam && (ArgUseIsReadOnly || isParamGridConstant(OldArg))) {
+    LLVM_DEBUG(dbgs() << "Using non-copy pointer to " << OldArg << "\n");
 
-    // Cast param address to generic address space.
+    // Cast the param-space argument to the generic address space. Because the
+    // argument is natively in param space, this cast only ever goes
+    // param -> generic and lowers to cvta.param; there is no inverse cast for
+    // InferAddressSpaces to fold it away with.
     Value *GenericArg = IRB.CreateAddrSpaceCast(
-        ArgInParamAS, IRB.getPtrTy(ADDRESS_SPACE_GENERIC),
-        Arg->getName() + ".gen");
-
-    Arg->replaceAllUsesWith(GenericArg);
+        &NewParamArg, IRB.getPtrTy(ADDRESS_SPACE_GENERIC),
+        OldArg.getName() + ".gen");
 
-    // Do not replace Arg in the cast to param space
-    ArgInParamAS->setOperand(0, Arg);
+    OldArg.replaceAllUsesWith(GenericArg);
     return;
   }
 
   // (3) Otherwise we have to create a copy of the argument in local memory.
-  copyByValParam(F, *Arg);
+  copyByValParam(F, OldArg, NewParamArg);
+}
+
+// Rewrite a kernel's signature so that each byval argument is declared directly
+// as a pointer in the param address space, then lower the body to match. This
+// creates a new function, moves the body across, and erases \p F.
+static void rewriteKernelByValSignature(Function &F, const bool HasCvtaParam) {
+  LLVMContext &Ctx = F.getContext();
+  FunctionType *FTy = F.getFunctionType();
+
+  // Build the new signature: byval pointer arguments move to the param address
+  // space; all other arguments are unchanged.
+  SmallVector<Type *> Params(FTy->params());
+  for (const Argument &Arg : F.args())
+    if (Arg.hasByValAttr())
+      Params[Arg.getArgNo()] = PointerType::get(Ctx, ADDRESS_SPACE_ENTRY_PARAM);
+
+  Function *NF = Function::Create(
+      FunctionType::get(FTy->getReturnType(), Params, FTy->isVarArg()),
+      F.getLinkage(), F.getAddressSpace());
+  NF->copyAttributesFrom(&F);
+  NF->setComdat(F.getComdat());
+  F.getParent()->getFunctionList().insert(F.getIterator(), NF);
+
+  // ISel reads the param symbol directly for kernel byval arguments; this is
+  // valid because the signature rewrite above puts them in the param address
+  // space. Mark them readonly: any mutation is redirected to a local copy
+  // below, so the param itself is never written.
+  for (Argument &NewArg : NF->args())
+    if (NewArg.hasByValAttr())
+      NewArg.addAttr(Attribute::ReadOnly);
+
+  // Take over F's name and uses (e.g. @llvm.used, nvvm.annotations metadata),
+  // then move the body across.
+  F.replaceAllUsesWith(NF);
+  NF->takeName(&F);
+  NF->splice(NF->begin(), &F);
+
+  // Remap arguments. Non-byval arguments keep their type and are replaced
+  // directly; byval arguments change address space, so their uses are lowered
+  // to operate on the new param-space argument.
+  for (auto [OldArg, NewArg] : zip_equal(F.args(), NF->args())) {
+    if (OldArg.hasByValAttr())
+      lowerKernelByValParam(OldArg, NewArg, *NF, HasCvtaParam);
----------------
akshayrdeodhar wrote:

Would it be possible to decouple the function rewrite and the lowering? Or would the "rewritten" function be in an inconsistent state? 
Original -> rewritten -> lower (use param space ptr / cast / create copy)

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