[NVPTX] Do not addrspacecast AS-specific kernel arguments.

[tra]

Summary

Fixes https://github.com/llvm/llvm-project/issues/46954

The assumption that generic pointers passed to a CUDA kernel is CUDA-specific and should not be applied to non-CUDA compilations.
Addrspacecasts to global AS and back should never be applied to AS-specific pointers.

In order to make tests actually do the testing for non-CUDA compilation, we need to get
TargetMachine from the TargetPassConfig, instead of passing it explicitly as a pass constructor argument.

[tra]

Cosmetic refactoring

[yaxunl]

LGTM. Thanks.

llvm/lib/Target/NVPTX/NVPTX.h

FunctionPass *createNVPTXISelDag(NVPTXTargetMachine &TM,
                                  llvm::CodeGenOpt::Level OptLevel);
 ModulePass *createNVPTXAssignValidGlobalNamesPass();
 ModulePass *createGenericToNVVMPass();
 FunctionPass *createNVVMIntrRangePass(unsigned int SmVersion);
 FunctionPass *createNVVMReflectPass(unsigned int SmVersion);
 MachineFunctionPass *createNVPTXPrologEpilogPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
-FunctionPass *createNVPTXLowerArgsPass(const NVPTXTargetMachine *TM);
+FunctionPass *createNVPTXLowerArgsPass();
 FunctionPass *createNVPTXLowerAllocaPass();
 MachineFunctionPass *createNVPTXPeephole();
 MachineFunctionPass *createNVPTXProxyRegErasurePass();
 
 struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
   NVVMIntrRangePass();
   NVVMIntrRangePass(unsigned SmVersion) : SmVersion(SmVersion) {}
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

 // cancel the addrspacecast pair this pass emits.
 //===----------------------------------------------------------------------===//
 
 #include "MCTargetDesc/NVPTXBaseInfo.h"
 #include "NVPTX.h"
 #include "NVPTXTargetMachine.h"
 #include "NVPTXUtilities.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
+#include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include <numeric>
 #include <queue>
 
 #define DEBUG_TYPE "nvptx-lower-args"
 
 using namespace llvm;
 
 namespace llvm {
 void initializeNVPTXLowerArgsPass(PassRegistry &);
 }
 
 namespace {
 class NVPTXLowerArgs : public FunctionPass {
   bool runOnFunction(Function &F) override;
 
-  bool runOnKernelFunction(Function &F);
-  bool runOnDeviceFunction(Function &F);
+  bool runOnKernelFunction(const NVPTXTargetMachine &TM, Function &F);
+  bool runOnDeviceFunction(const NVPTXTargetMachine &TM, Function &F);
 
   // handle byval parameters
-  void handleByValParam(Argument *Arg);
+  void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg);
   // Knowing Ptr must point to the global address space, this function
   // addrspacecasts Ptr to global and then back to generic. This allows
   // NVPTXInferAddressSpaces to fold the global-to-generic cast into
   // loads/stores that appear later.
   void markPointerAsGlobal(Value *Ptr);
 
 public:
   static char ID; // Pass identification, replacement for typeid
-  NVPTXLowerArgs(const NVPTXTargetMachine *TM = nullptr)
-      : FunctionPass(ID), TM(TM) {}
+  NVPTXLowerArgs() : FunctionPass(ID) {}
   StringRef getPassName() const override {
     return "Lower pointer arguments of CUDA kernels";
   }
-
-private:
-  const NVPTXTargetMachine *TM;
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetPassConfig>();
+  }
 };
 } // namespace
 
 char NVPTXLowerArgs::ID = 1;
 
-INITIALIZE_PASS(NVPTXLowerArgs, "nvptx-lower-args",
+INITIALIZE_PASS_BEGIN(NVPTXLowerArgs, "nvptx-lower-args",
+                      "Lower arguments (NVPTX)", false, false)
+INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
+INITIALIZE_PASS_END(NVPTXLowerArgs, "nvptx-lower-args",
                     "Lower arguments (NVPTX)", false, false)
 
 // =============================================================================
 // If the function had a byval struct ptr arg, say foo(%struct.x* byval %d),
 // and we can't guarantee that the only accesses are loads,
 // then add the following instructions to the first basic block:
 //
 // %temp = alloca %struct.x, align 8
 // %tempd = addrspacecast %struct.x* %d to %struct.x addrspace(101)*
static void adjustByValArgAlignment(Argument *Arg, Value *ArgInParamAS,
 
   for (Load &CurLoad : Loads) {
     Align NewLoadAlign(std::gcd(NewArgAlign, CurLoad.Offset));
     Align CurLoadAlign(CurLoad.Inst->getAlign());
     CurLoad.Inst->setAlignment(std::max(NewLoadAlign, CurLoadAlign));
   }
 }
 
-void NVPTXLowerArgs::handleByValParam(Argument *Arg) {
+void NVPTXLowerArgs::handleByValParam(const NVPTXTargetMachine &TM,
+                                      Argument *Arg) {
   Function *Func = Arg->getParent();
   Instruction *FirstInst = &(Func->getEntryBlock().front());
   Type *StructType = Arg->getParamByValType();
   assert(StructType && "Missing byval type");
 
   auto IsALoadChain = [&](Value *Start) {
     SmallVector<Value *, 16> ValuesToCheck = {Start};
     auto IsALoadChainInstr = [](Value *V) -> bool {
if (llvm::all_of(Arg->users(), IsALoadChain)) {
     SmallVector<User *, 16> UsersToUpdate(Arg->users());
     Value *ArgInParamAS = new AddrSpaceCastInst(
         Arg, PointerType::get(StructType, ADDRESS_SPACE_PARAM), Arg->getName(),
         FirstInst);
     for (Value *V : UsersToUpdate)
       convertToParamAS(V, ArgInParamAS);
     LLVM_DEBUG(dbgs() << "No need to copy " << *Arg << "\n");
 
-    // Further optimizations require target lowering info.
-    if (!TM)
-      return;
-
     const auto *TLI =
-        cast<NVPTXTargetLowering>(TM->getSubtargetImpl()->getTargetLowering());
+        cast<NVPTXTargetLowering>(TM.getSubtargetImpl()->getTargetLowering());
 
     adjustByValArgAlignment(Arg, ArgInParamAS, TLI);
 
     return;
   }
 
   // Otherwise we have to create a temporary copy.
   const DataLayout &DL = Func->getParent()->getDataLayout();
void NVPTXLowerArgs::handleByValParam(const NVPTXTargetMachine &TM,
   // definitely not volatile.
   LoadInst *LI =
       new LoadInst(StructType, ArgInParam, Arg->getName(),
                    /*isVolatile=*/false, AllocA->getAlign(), FirstInst);
   new StoreInst(LI, AllocA, FirstInst);
 }
 
 void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {
-  if (Ptr->getType()->getPointerAddressSpace() == ADDRESS_SPACE_GLOBAL)
+  if (Ptr->getType()->getPointerAddressSpace() != ADDRESS_SPACE_GENERIC)
     return;
 
   // Deciding where to emit the addrspacecast pair.
   BasicBlock::iterator InsertPt;
   if (Argument *Arg = dyn_cast<Argument>(Ptr)) {
     // Insert at the functon entry if Ptr is an argument.
     InsertPt = Arg->getParent()->getEntryBlock().begin();
   } else {
void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {
   // Replace with PtrInGeneric all uses of Ptr except PtrInGlobal.
   Ptr->replaceAllUsesWith(PtrInGeneric);
   PtrInGlobal->setOperand(0, Ptr);
 }
 
 // =============================================================================
 // Main function for this pass.
 // =============================================================================
-bool NVPTXLowerArgs::runOnKernelFunction(Function &F) {
-  if (TM && TM->getDrvInterface() == NVPTX::CUDA) {
+bool NVPTXLowerArgs::runOnKernelFunction(const NVPTXTargetMachine &TM,
+                                         Function &F) {
+  if (TM.getDrvInterface() == NVPTX::CUDA) {
     // Mark pointers in byval structs as global.
     for (auto &B : F) {
       for (auto &I : B) {
         if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
           if (LI->getType()->isPointerTy()) {
             Value *UO = getUnderlyingObject(LI->getPointerOperand());
             if (Argument *Arg = dyn_cast<Argument>(UO)) {
               if (Arg->hasByValAttr()) {
                 // LI is a load from a pointer within a byval kernel parameter.
                 markPointerAsGlobal(LI);
               }
             }
           }
         }
       }
     }
   }
 
   LLVM_DEBUG(dbgs() << "Lowering kernel args of " << F.getName() << "\n");
   for (Argument &Arg : F.args()) {
     if (Arg.getType()->isPointerTy()) {
       if (Arg.hasByValAttr())
-        handleByValParam(&Arg);
-      else if (TM && TM->getDrvInterface() == NVPTX::CUDA)
+        handleByValParam(TM, &Arg);
+      else if (TM.getDrvInterface() == NVPTX::CUDA)
         markPointerAsGlobal(&Arg);
     }
   }
   return true;
 }
 
 // Device functions only need to copy byval args into local memory.
-bool NVPTXLowerArgs::runOnDeviceFunction(Function &F) {
+bool NVPTXLowerArgs::runOnDeviceFunction(const NVPTXTargetMachine &TM,
+                                         Function &F) {
   LLVM_DEBUG(dbgs() << "Lowering function args of " << F.getName() << "\n");
   for (Argument &Arg : F.args())
     if (Arg.getType()->isPointerTy() && Arg.hasByValAttr())
-      handleByValParam(&Arg);
+      handleByValParam(TM, &Arg);
   return true;
 }
 
 bool NVPTXLowerArgs::runOnFunction(Function &F) {
-  return isKernelFunction(F) ? runOnKernelFunction(F) : runOnDeviceFunction(F);
-}
+  auto &TM = getAnalysis<TargetPassConfig>().getTM<NVPTXTargetMachine>();
 
-FunctionPass *
-llvm::createNVPTXLowerArgsPass(const NVPTXTargetMachine *TM) {
-  return new NVPTXLowerArgs(TM);
+  return isKernelFunction(F) ? runOnKernelFunction(TM, F)
+                             : runOnDeviceFunction(TM, F);
 }
+
+FunctionPass *llvm::createNVPTXLowerArgsPass() { return new NVPTXLowerArgs(); }

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

void NVPTXPassConfig::addIRPasses() {
 
   if (getOptLevel() != CodeGenOpt::None)
     addPass(createNVPTXImageOptimizerPass());
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
 
   // NVPTXLowerArgs is required for correctness and should be run right
   // before the address space inference passes.
-  addPass(createNVPTXLowerArgsPass(&getNVPTXTargetMachine()));
+  addPass(createNVPTXLowerArgsPass());
   if (getOptLevel() != CodeGenOpt::None) {
     addAddressSpaceInferencePasses();
     addStraightLineScalarOptimizationPasses();
   }
 
   addPass(createAtomicExpandPass());
 
   // === LSR and other generic IR passes ===

llvm/test/CodeGen/NVPTX/lower-args.ll

-; RUN: opt < %s -S -nvptx-lower-args | FileCheck %s --check-prefix IR
-; RUN: llc < %s -mcpu=sm_20 | FileCheck %s --check-prefix PTX
+; RUN: opt < %s -S -nvptx-lower-args --mtriple nvptx64-nvidia-cuda | FileCheck %s --check-prefixes COMMON,IR,IRC
+; RUN: opt < %s -S -nvptx-lower-args --mtriple nvptx64-nvidia-nvcl | FileCheck %s --check-prefixes COMMON,IR,IRO
+; RUN: llc < %s -mcpu=sm_20 --mtriple nvptx64-nvidia-cuda | FileCheck %s --check-prefixes COMMON,PTX,PTXC
+; RUN: llc < %s -mcpu=sm_20 --mtriple nvptx64-nvidia-nvcl| FileCheck %s --check-prefixes COMMON,PTX,PTXO
 ; RUN: %if ptxas %{ llc < %s -mcpu=sm_20 | %ptxas-verify %}
 
 target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
 target triple = "nvptx64-nvidia-cuda"
 
 %class.outer = type <{ %class.inner, i32, [4 x i8] }>
 %class.inner = type { ptr, ptr }
 
 ; Check that nvptx-lower-args preserves arg alignment
+; COMMON-LABEL: load_alignment
 define void @load_alignment(ptr nocapture readonly byval(%class.outer) align 8 %arg) {
 entry:
 ; IR: load %class.outer, ptr addrspace(101)
 ; IR-SAME: align 8
 ; PTX: ld.param.u64
 ; PTX-NOT: ld.param.u8
   %arg.idx.val = load ptr, ptr %arg, align 8
   %arg.idx1 = getelementptr %class.outer, ptr %arg, i64 0, i32 0, i32 1
   %arg.idx1.val = load ptr, ptr %arg.idx1, align 8
   %arg.idx2 = getelementptr %class.outer, ptr %arg, i64 0, i32 1
   %arg.idx2.val = load i32, ptr %arg.idx2, align 8
   %arg.idx.val.val = load i32, ptr %arg.idx.val, align 4
   %add.i = add nsw i32 %arg.idx.val.val, %arg.idx2.val
   store i32 %add.i, ptr %arg.idx1.val, align 4
 
   ; let the pointer escape so we still create a local copy this test uses to
   ; check the load alignment.
   %tmp = call ptr @escape(ptr nonnull %arg.idx2)
   ret void
 }
 
+
+; COMMON-LABEL: ptr_generic
+define void @ptr_generic(ptr %out, ptr %in) {
+; IRC:  %in3 = addrspacecast ptr %in to ptr addrspace(1)
+; IRC:  %in4 = addrspacecast ptr addrspace(1) %in3 to ptr
+; IRC:  %out1 = addrspacecast ptr %out to ptr addrspace(1)
+; IRC:  %out2 = addrspacecast ptr addrspace(1) %out1 to ptr
+; PTXC: cvta.to.global.u64
+; PTXC: cvta.to.global.u64
+; PTXC: ld.global.u32
+; PTXC: st.global.u32
+
+; OpenCL can't make assumptions about incoming pointer, so we should generate
+; generic pointers load/store.
+; IRO-NOT: addrspacecast
+; PTXO-NOT: cvta.to.global
+; PTXO: ld.u32
+; PTXO: st.u32
+  %v = load i32, ptr  %in, align 4
+  store i32 %v, ptr %out, align 4
+  ret void
+}
+
+; COMMON-LABEL: ptr_nongeneric
+define void @ptr_nongeneric(ptr addrspace(1) %out, ptr addrspace(4) %in) {
+; IR-NOT: addrspacecast
+; PTX-NOT: cvta.to.global
+; PTX:  ld.const.u32
+; PTX   st.global.u32
+  %v = load i32, ptr addrspace(4) %in, align 4
+  store i32 %v, ptr addrspace(1) %out, align 4
+  ret void
+}
+
+
 ; Function Attrs: convergent nounwind
 declare dso_local ptr @escape(ptr) local_unnamed_addr
+!nvvm.annotations = !{!0, !1}
+!0 = !{ptr @ptr_generic, !"kernel", i32 1}
+!1 = !{ptr @ptr_nongeneric, !"kernel", i32 1}