[NVPTX] Infer AS of pointers passed to kernels as integers.

[tra]

Summary

When pointers are passed within aggregates, we sometimes end up with IR that
loads them as integers and the converts them back to pointers. Typically it's
due to a memcpy or SROA. E.g. https://godbolt.org/z/xM3n5daaa

Normally we treat all pointers passed to a CUDA kernel as global pointers and
the same treatment should be applied to the pointers we load/store as integers.

[tra]

Updated tests.

[jlebar]

Maybe worth a comment in the code that says basically what you say in the commit message? Otherwise LGTM.

[tra]

Added a comment.

llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp

void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {
   PtrInGlobal->setOperand(0, Ptr);
 }
 
 // =============================================================================
 // Main function for this pass.
 // =============================================================================
 bool NVPTXLowerArgs::runOnKernelFunction(const NVPTXTargetMachine &TM,
                                          Function &F) {
+  // Copying of byval aggregates + SROA may result in pointers being loaded as
+  // integers, followed by intotoptr. We may want to mark those as global, too,
+  // but only if the loaded integer is used exclusively for conversion to a
+  // pointer with inttoptr.
+  auto HandleIntToPtr = [this](Value &V) {
+    if (llvm::all_of(V.users(), [](User *U) { return isa<IntToPtrInst>(U); })) {
+      SmallVector<User *, 16> UsersToUpdate(V.users());
+      llvm::for_each(UsersToUpdate, [&](User *U) { markPointerAsGlobal(U); });
+    }
+  };
   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()) {
+          if (LI->getType()->isPointerTy() || LI->getType()->isIntegerTy()) {
             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.
+                if (LI->getType()->isPointerTy())
                   markPointerAsGlobal(LI);
+                else
+                  HandleIntToPtr(*LI);
               }
             }
           }
         }
       }
     }
   }
 
   LLVM_DEBUG(dbgs() << "Lowering kernel args of " << F.getName() << "\n");
   for (Argument &Arg : F.args()) {
     if (Arg.getType()->isPointerTy()) {
       if (Arg.hasByValAttr())
         handleByValParam(TM, &Arg);
       else if (TM.getDrvInterface() == NVPTX::CUDA)
         markPointerAsGlobal(&Arg);
+    } else if (Arg.getType()->isIntegerTy() &&
+               TM.getDrvInterface() == NVPTX::CUDA) {
+      HandleIntToPtr(Arg);
     }
   }
   return true;
 }
 
 // Device functions only need to copy byval args into local memory.
 bool NVPTXLowerArgs::runOnDeviceFunction(const NVPTXTargetMachine &TM,
                                          Function &F) {

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

 ; 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
 }
 
+; COMMON-LABEL: ptr_as_int
+ define void @ptr_as_int(i64 noundef %i, i32 noundef %v) {
+; IR:   [[P:%.*]] = inttoptr i64 %i to ptr
+; IRC:  [[P1:%.*]] = addrspacecast ptr [[P]] to ptr addrspace(1)
+; IRC:  addrspacecast ptr addrspace(1) [[P1]] to ptr
+; IRO-NOT: addrspacecast
+
+; PTXC-DAG:  ld.param.u64    [[I:%rd.*]], [ptr_as_int_param_0];
+; PTXC-DAG:  ld.param.u32    [[V:%r.*]], [ptr_as_int_param_1];
+; PTXC:      cvta.to.global.u64 %[[P:rd.*]], [[I]];
+; PTXC:      st.global.u32    [%[[P]]], [[V]];
+
+; PTXO-DAG:  ld.param.u64    %[[P:rd.*]], [ptr_as_int_param_0];
+; PTXO-DAG:  ld.param.u32    [[V:%r.*]], [ptr_as_int_param_1];
+; PTXO:      st.u32   [%[[P]]], [[V]];
+
+  %p = inttoptr i64 %i to ptr
+  store i32 %v, ptr %p, align 4
+  ret void
+}
+
+%struct.S = type { i64 }
+
+; COMMON-LABEL: ptr_as_int_aggr
+define void @ptr_as_int_aggr(ptr nocapture noundef readonly byval(%struct.S) align 8 %s, i32 noundef %v) {
+; IR:   [[S:%.*]] = addrspacecast ptr %s to ptr addrspace(101)
+; IR:   [[I:%.*]] = load i64, ptr addrspace(101) [[S]], align 8
+; IR:   [[P0:%.*]] = inttoptr i64 [[I]] to ptr
+; IRC:  [[P1:%.*]] = addrspacecast ptr [[P]] to ptr addrspace(1)
+; IRC:  [[P:%.*]] = addrspacecast ptr addrspace(1) [[P1]] to ptr
+; IRO-NOT: addrspacecast
+
+; PTXC-DAG:  ld.param.u64    [[I:%rd.*]], [ptr_as_int_aggr_param_0];
+; PTXC-DAG:  ld.param.u32    [[V:%r.*]], [ptr_as_int_aggr_param_1];
+; PTXC:      cvta.to.global.u64 %[[P:rd.*]], [[I]];
+; PTXC:      st.global.u32    [%[[P]]], [[V]];
+
+; PTXO-DAG:  ld.param.u64    %[[P:rd.*]], [ptr_as_int_aggr_param_0];
+; PTXO-DAG:  ld.param.u32    [[V:%r.*]], [ptr_as_int_aggr_param_1];
+; PTXO:      st.u32   [%[[P]]], [[V]];
+  %i = load i64, ptr %s, align 8
+  %p = inttoptr i64 %i to ptr
+  store i32 %v, ptr %p, align 4
+  ret void
+}
+
 
 ; Function Attrs: convergent nounwind
 declare dso_local ptr @escape(ptr) local_unnamed_addr
-!nvvm.annotations = !{!0, !1}
+!nvvm.annotations = !{!0, !1, !2, !3}
 !0 = !{ptr @ptr_generic, !"kernel", i32 1}
 !1 = !{ptr @ptr_nongeneric, !"kernel", i32 1}
+!2 = !{ptr @ptr_as_int, !"kernel", i32 1}
+!3 = !{ptr @ptr_as_int_aggr, !"kernel", i32 1}