[NVPTX] kernel pointer arguments point to the global address space

[jingyue]

Summary

With this patch, NVPTXLowerKernelArgs converts a kernel pointer argument to a
pointer in the global address space. This change, along with
NVPTXFavorNonGenericAddrSpaces, allows the NVPTX backend to emit ld.global.*
and st.global.* for accessing kernel pointer arguments.

Minor changes:

1. refactor: extract function convertToPointerInAddrSpace
2. fix a bug in the test case in bug21465.ll

[jholewinski]

Comments inlined.

Are you planning on enabling this pass by default?

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

37

The address space conversion intrinsics are deprecated in favor of the new addrspacecast instruction (I need to document that in IntrinsicsNVVM.td).

40

Isn't it possible that an optimization would remove all of these casts before NVVMFavorNonGenericAddrSpaces runs? I know we control the pass pipeline in the backend, but I worry about these pass ordering constraints.

189

Strictly speaking, this is only valid for CUDA (DrvInterface::CUDA).

[jingyue]

Yes. I'll add that to NVPTXPassConfig::addIRPass.

[jingyue]

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

37

Thank you for pointing this out. So, do you think we should emit addrspacecast here too instead of gen.to.global and gen.to.param?

40

Not sure. Removing bitcast is fine for NVPTXFavorNonGenericAddrSpace. Why would any pass remove addrspacecast? addrspacecast has a richer semantics -- it changes the address space of a pointer and accessing pointers in different address spaces are not necessarily as fast. If any pass would want to remove an addrspacecast, it should reason about the performance effects of that too.

189

ACK

[jingyue]

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

40

Your concern makes more sense to me if we replace llvm.nvvm.ptr.gen.to.* with addrspacecast. In that case, this pass would addrspacecasts value back and forth, e.g.,

t1 = cast t0 to global
t2 = cast t1 to generic

Then, it's reasonable for some optimizations to cancel such addrspacecast pairs.

Maybe not in this patch, we can merge this pass and NVPTXFavorNonGeneric into something like NVPTXTypeInference? That should prevent other passes from messing up the intermediate code (besides the name becomes fancier :) How does that sound to you?

[jingyue]

Enable NVPTXLowerKernelArgs by default

[jingyue]

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

189

Done.

[jingyue]

use addrspacecast instead of llvm.nvvm.ptr.gen.to.* intrinsics

[jingyue]

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

37

The updated version emits addrspacecast instead of llvm.nvvm.ptr.gen.to.*.

[jingyue]

more comments

[jholewinski]

This looks good to me now. I agree that we should merge these two passes to prevent any ordering issues.

[jingyue]

Thank you. FYI, we are working on merging LowerKernelArgs and FavorNonGeneric, and handling the local address space as well.

lib/Target/NVPTX/CMakeLists.txt

set(NVPTXCodeGen_sources
   NVPTXFavorNonGenericAddrSpaces.cpp
   NVPTXFrameLowering.cpp
   NVPTXGenericToNVVM.cpp
   NVPTXISelDAGToDAG.cpp
   NVPTXISelLowering.cpp
   NVPTXImageOptimizer.cpp
   NVPTXInstrInfo.cpp
   NVPTXLowerAggrCopies.cpp
-  NVPTXLowerStructArgs.cpp
+  NVPTXLowerKernelArgs.cpp
   NVPTXMCExpr.cpp
   NVPTXPrologEpilogPass.cpp
   NVPTXRegisterInfo.cpp
   NVPTXReplaceImageHandles.cpp
   NVPTXSubtarget.cpp
   NVPTXTargetMachine.cpp
   NVPTXTargetTransformInfo.cpp
   NVPTXUtilities.cpp
   NVVMReflect.cpp
   )
 
 add_llvm_target(NVPTXCodeGen ${NVPTXCodeGen_sources})
 
 add_subdirectory(TargetInfo)
 add_subdirectory(InstPrinter)
 add_subdirectory(MCTargetDesc)

lib/Target/NVPTX/NVPTX.h

 ModulePass *createNVPTXAssignValidGlobalNamesPass();
 ModulePass *createGenericToNVVMPass();
 FunctionPass *createNVPTXFavorNonGenericAddrSpacesPass();
 ModulePass *createNVVMReflectPass();
 ModulePass *createNVVMReflectPass(const StringMap<int>& Mapping);
 MachineFunctionPass *createNVPTXPrologEpilogPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
-FunctionPass *createNVPTXLowerStructArgsPass();
+FunctionPass *createNVPTXLowerKernelArgsPass();
 
 bool isImageOrSamplerVal(const Value *, const Module *);
 
 extern Target TheNVPTXTarget32;
 extern Target TheNVPTXTarget64;
 
 namespace NVPTX {
 enum DrvInterface {

lib/Target/NVPTX/NVPTXLowerKernelArgs.cpp

+//===-- NVPTXLowerKernelArgs.cpp - Lower kernel arguments -----------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Pointer arguments to kernel functions need to be lowered specially.
+//
+// 1. Copy byval struct args to local memory. This is a preparation for handling
+//    cases like
+//
+//    kernel void foo(struct A arg, ...)
+//    {
+//      struct A *p = &arg;
+//      ...
+//      ... = p->filed1 ...  (this is no generic address for .param)
+//      p->filed2 = ...      (this is no write access to .param)
+//    }
+//
+// 2. Convert non-byval pointer arguments of CUDA kernels to pointers in the
+//    global address space. This allows later optimizations to emit
+//    ld.global.*/st.global.* for accessing these pointer arguments. For
+//    example,
+//
+//    define void @foo(float* %input) {
+//      %v = load float, float* %input, align 4
+//      ...
+//    }
+//
+//    becomes
+//
+//    define void @foo(float* %input) {
+//      %input2 = bitcast float* %input to i8*
+//      %cvt = call i8 addrspace(1)* @llvm.nvvm.ptr.gen.to.global.p1i8.p0i8(i8* %input2)
+//      %input3 = bitcast i8 addrspace(1)* %cvt to float addrspace(1)*
+//      %input4 = addrspacecast float addrspace(1)* %input3 to float*
+//      %v = load float, float* %input4, align 4
+//      ...
+//    }
+//
+//    Later, NVPTXFavorNonGenericAddrSpaces will optimize
+//
+//    %input4 = addrspacecast float addrspace(1)* %input3 to float*
+//    %v = load float, float* %input4, align 4
+//
+//    to
+//
+//    %v = load float, float addrspace(1)* %input3, align 4
+//===----------------------------------------------------------------------===//
+
+#include "NVPTX.h"
+#include "NVPTXUtilities.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Pass.h"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeNVPTXLowerKernelArgsPass(PassRegistry &);
+}
+
+namespace {
+class NVPTXLowerKernelArgs : public FunctionPass {
+  bool runOnFunction(Function &F) override;
+
+  // handle byval parameters
+  void handleByValParam(Argument *);
+  // handle non-byval pointer parameters
+  void handlePointerParam(Argument *);
+  // a helper function to convert a pointer argument in generic to a pointer in
+  // the target address space.
+  Value *convertToPointerInAddrSpace(Value *Arg, enum AddressSpace TargetAS,
+                                     Instruction *InsertPt);
+
+public:
+  static char ID; // Pass identification, replacement for typeid
+  NVPTXLowerKernelArgs() : FunctionPass(ID) {}
+  const char *getPassName() const override {
+    return "Lower pointer arguments of CUDA kernels";
+  }
+};
+} // namespace
+
+char NVPTXLowerKernelArgs::ID = 1;
+
+INITIALIZE_PASS(NVPTXLowerKernelArgs, "nvptx-lower-kernel-args",
+                "Lower kernel arguments (NVPTX)", false, false)
+
+// =============================================================================
+// If the function had a byval struct ptr arg, say foo(%struct.x *byval %d),
+// then add the following instructions to the first basic block:
+//
+// %temp = alloca %struct.x, align 8
+// %tt1 = bitcast %struct.x * %d to i8 *
+// %tt2 = llvm.nvvm.cvt.gen.to.param %tt2
+// %tempd = bitcast i8 addrspace(101) * to %struct.x addrspace(101) *
+// %tv = load %struct.x addrspace(101) * %tempd
+// store %struct.x %tv, %struct.x * %temp, align 8
+//
+// The above code allocates some space in the stack and copies the incoming
+// struct from param space to local space.
+// Then replace all occurences of %d by %temp.
+// =============================================================================
+void NVPTXLowerKernelArgs::handleByValParam(Argument *Arg) {
+  Function *Func = Arg->getParent();
+  Instruction *FirstInst = &(Func->getEntryBlock().front());
+  PointerType *PType = dyn_cast<PointerType>(Arg->getType());
+
+  assert(PType && "Expecting pointer type in handleByValParam");
+
+  Type *StructType = PType->getElementType();
+  AllocaInst *AllocA = new AllocaInst(StructType, Arg->getName(), FirstInst);
+  // Set the alignment to alignment of the byval parameter. This is because,
+  // later load/stores assume that alignment, and we are going to replace
+  // the use of the byval parameter with this alloca instruction.
+  AllocA->setAlignment(Func->getParamAlignment(Arg->getArgNo() + 1));
+  Arg->replaceAllUsesWith(AllocA);
+
+  Value *ArgInParam =
+      convertToPointerInAddrSpace(Arg, ADDRESS_SPACE_PARAM, FirstInst);
+  LoadInst *LI = new LoadInst(ArgInParam, Arg->getName(), FirstInst);
+  new StoreInst(LI, AllocA, FirstInst);
+}
+
+Value *NVPTXLowerKernelArgs::convertToPointerInAddrSpace(
+    Value *Arg, enum AddressSpace TargetAS, Instruction *InsertPt) {
+  // Get the llvm.nvvm.ptr.gen.to.<TargetAS> intrinsic.
+  Type *Int8PtrInTargetAS = Type::getInt8PtrTy(Arg->getContext(), TargetAS);
+  Type *Int8PtrInGenericAS =
+      Type::getInt8PtrTy(Arg->getContext(), ADDRESS_SPACE_GENERIC);
+  Type *CvtTypes[] = {Int8PtrInTargetAS, Int8PtrInGenericAS};
+  Function *CvtFunc = Intrinsic::getDeclaration(
+      InsertPt->getModule(), Intrinsic::nvvm_ptr_gen_to_param, CvtTypes);
+
+  // llvm.nvvm.ptr.gen.to.* converts i8* only. Therefore, we bitcast the
+  // argument to i8*, call the conversion function, and bitcast the result back
+  // to the original type in the target address space.
+  Value *ArgInI8Star =
+      new BitCastInst(Arg, Int8PtrInGenericAS, Arg->getName(), InsertPt);
+  CallInst *CallCVT = CallInst::Create(CvtFunc, ArgInI8Star, "cvt", InsertPt);
+  return new BitCastInst(
+      CallCVT,
+      PointerType::get(Arg->getType()->getPointerElementType(), TargetAS),
+      Arg->getName(), InsertPt);
+}
+
+void NVPTXLowerKernelArgs::handlePointerParam(Argument *Arg) {
+  assert(!Arg->hasByValAttr() &&
+         "byval params should be handled by handleByValParam");
+
+  Instruction *FirstInst = Arg->getParent()->getEntryBlock().begin();
+  // Collect the original uses of Arg.
+  SmallVector<Use *, 4> UsesOfArg;
+  for (Use &U : Arg->uses())
+    UsesOfArg.push_back(&U);
+
+  Value *ArgInGlobal =
+      convertToPointerInAddrSpace(Arg, ADDRESS_SPACE_GLOBAL, FirstInst);
+  Value *ArgInGeneric = new AddrSpaceCastInst(ArgInGlobal, Arg->getType(),
+                                              Arg->getName(), FirstInst);
+
+  // Replace all the original uses, avoiding modifying Arg's user added in
+  // convertToPointerInAddrSpace.
+  for (Use *U : UsesOfArg)
+    U->set(ArgInGeneric);
+}
+
+
+// =============================================================================
+// Main function for this pass.
+// =============================================================================
+bool NVPTXLowerKernelArgs::runOnFunction(Function &F) {
+  // Skip non-kernels. See the comments at the top of this file.
+  if (!isKernelFunction(F))
+    return false;
+
+  for (Argument &Arg : F.args()) {
+    if (Arg.getType()->isPointerTy()) {
+      if (Arg.hasByValAttr())
+        handleByValParam(&Arg);
+      else
+        handlePointerParam(&Arg);
+    }
+  }
+  return true;
+}
+
+FunctionPass *llvm::createNVPTXLowerKernelArgsPass() {
+  return new NVPTXLowerKernelArgs();
+}

lib/Target/NVPTX/NVPTXLowerStructArgs.cpp

-//===-- NVPTXLowerStructArgs.cpp - Copy struct args to local memory =====--===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-//
-// Copy struct args to local memory. This is needed for kernel functions only.
-// This is a preparation for handling cases like
-//
-// kernel void foo(struct A arg, ...)
-// {
-//     struct A *p = &arg;
-//     ...
-//     ... = p->filed1 ...  (this is no generic address for .param)
-//     p->filed2 = ...      (this is no write access to .param)
-// }
-//
-//===----------------------------------------------------------------------===//
-
-#include "NVPTX.h"
-#include "NVPTXUtilities.h"
-#include "llvm/IR/Function.h"
-#include "llvm/IR/Instructions.h"
-#include "llvm/IR/IntrinsicInst.h"
-#include "llvm/IR/Module.h"
-#include "llvm/IR/Type.h"
-#include "llvm/Pass.h"
-
-using namespace llvm;
-
-namespace llvm {
-void initializeNVPTXLowerStructArgsPass(PassRegistry &);
-}
-
-namespace {
-class NVPTXLowerStructArgs : public FunctionPass {
-  bool runOnFunction(Function &F) override;
-
-  void handleStructPtrArgs(Function &);
-  void handleParam(Argument *);
-
-public:
-  static char ID; // Pass identification, replacement for typeid
-  NVPTXLowerStructArgs() : FunctionPass(ID) {}
-  const char *getPassName() const override {
-    return "Copy structure (byval *) arguments to stack";
-  }
-};
-} // namespace
-
-char NVPTXLowerStructArgs::ID = 1;
-
-INITIALIZE_PASS(NVPTXLowerStructArgs, "nvptx-lower-struct-args",
-                "Lower structure arguments (NVPTX)", false, false)
-
-void NVPTXLowerStructArgs::handleParam(Argument *Arg) {
-  Function *Func = Arg->getParent();
-  Instruction *FirstInst = &(Func->getEntryBlock().front());
-  PointerType *PType = dyn_cast<PointerType>(Arg->getType());
-
-  assert(PType && "Expecting pointer type in handleParam");
-
-  Type *StructType = PType->getElementType();
-  AllocaInst *AllocA = new AllocaInst(StructType, Arg->getName(), FirstInst);
-
-  /* Set the alignment to alignment of the byval parameter. This is because,
-   * later load/stores assume that alignment, and we are going to replace
-   * the use of the byval parameter with this alloca instruction.
-   */
-  AllocA->setAlignment(Func->getParamAlignment(Arg->getArgNo() + 1));
-
-  Arg->replaceAllUsesWith(AllocA);
-
-  // Get the cvt.gen.to.param intrinsic
-  Type *CvtTypes[] = {
-      Type::getInt8PtrTy(Func->getParent()->getContext(), ADDRESS_SPACE_PARAM),
-      Type::getInt8PtrTy(Func->getParent()->getContext(),
-                         ADDRESS_SPACE_GENERIC)};
-  Function *CvtFunc = Intrinsic::getDeclaration(
-      Func->getParent(), Intrinsic::nvvm_ptr_gen_to_param, CvtTypes);
-
-  Value *BitcastArgs[] = {
-      new BitCastInst(Arg, Type::getInt8PtrTy(Func->getParent()->getContext(),
-                                              ADDRESS_SPACE_GENERIC),
-                      Arg->getName(), FirstInst)};
-  CallInst *CallCVT =
-      CallInst::Create(CvtFunc, BitcastArgs, "cvt_to_param", FirstInst);
-
-  BitCastInst *BitCast = new BitCastInst(
-      CallCVT, PointerType::get(StructType, ADDRESS_SPACE_PARAM),
-      Arg->getName(), FirstInst);
-  LoadInst *LI = new LoadInst(BitCast, Arg->getName(), FirstInst);
-  new StoreInst(LI, AllocA, FirstInst);
-}
-
-// =============================================================================
-// If the function had a struct ptr arg, say foo(%struct.x *byval %d), then
-// add the following instructions to the first basic block :
-//
-// %temp = alloca %struct.x, align 8
-// %tt1 = bitcast %struct.x * %d to i8 *
-// %tt2 = llvm.nvvm.cvt.gen.to.param %tt2
-// %tempd = bitcast i8 addrspace(101) * to %struct.x addrspace(101) *
-// %tv = load %struct.x addrspace(101) * %tempd
-// store %struct.x %tv, %struct.x * %temp, align 8
-//
-// The above code allocates some space in the stack and copies the incoming
-// struct from param space to local space.
-// Then replace all occurences of %d by %temp.
-// =============================================================================
-void NVPTXLowerStructArgs::handleStructPtrArgs(Function &F) {
-  for (Argument &Arg : F.args()) {
-    if (Arg.getType()->isPointerTy() && Arg.hasByValAttr()) {
-      handleParam(&Arg);
-    }
-  }
-}
-
-// =============================================================================
-// Main function for this pass.
-// =============================================================================
-bool NVPTXLowerStructArgs::runOnFunction(Function &F) {
-  // Skip non-kernels. See the comments at the top of this file.
-  if (!isKernelFunction(F))
-    return false;
-
-  handleStructPtrArgs(F);
-  return true;
-}
-
-FunctionPass *llvm::createNVPTXLowerStructArgsPass() {
-  return new NVPTXLowerStructArgs();
-}

lib/Target/NVPTX/NVPTXTargetMachine.cpp

 using namespace llvm;
 
 namespace llvm {
 void initializeNVVMReflectPass(PassRegistry&);
 void initializeGenericToNVVMPass(PassRegistry&);
 void initializeNVPTXAllocaHoistingPass(PassRegistry &);
 void initializeNVPTXAssignValidGlobalNamesPass(PassRegistry&);
 void initializeNVPTXFavorNonGenericAddrSpacesPass(PassRegistry &);
-void initializeNVPTXLowerStructArgsPass(PassRegistry &);
+void initializeNVPTXLowerKernelArgsPass(PassRegistry &);
 }
 
 extern "C" void LLVMInitializeNVPTXTarget() {
   // Register the target.
   RegisterTargetMachine<NVPTXTargetMachine32> X(TheNVPTXTarget32);
   RegisterTargetMachine<NVPTXTargetMachine64> Y(TheNVPTXTarget64);
 
   // FIXME: This pass is really intended to be invoked during IR optimization,
   // but it's very NVPTX-specific.
   initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
   initializeGenericToNVVMPass(*PassRegistry::getPassRegistry());
   initializeNVPTXAllocaHoistingPass(*PassRegistry::getPassRegistry());
   initializeNVPTXAssignValidGlobalNamesPass(*PassRegistry::getPassRegistry());
   initializeNVPTXFavorNonGenericAddrSpacesPass(
     *PassRegistry::getPassRegistry());
-  initializeNVPTXLowerStructArgsPass(*PassRegistry::getPassRegistry());
+  initializeNVPTXLowerKernelArgsPass(*PassRegistry::getPassRegistry());
 }
 
 static std::string computeDataLayout(bool is64Bit) {
   std::string Ret = "e";
 
   if (!is64Bit)
     Ret += "-p:32:32";
 

test/CodeGen/NVPTX/bug21465.ll

-; RUN: opt < %s -nvptx-lower-struct-args -S | FileCheck %s
+; RUN: opt < %s -nvptx-lower-kernel-args -S | FileCheck %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
 target triple = "nvptx64-unknown-unknown"
 
 %struct.S = type { i32, i32 }
 
 ; Function Attrs: nounwind
 define void @_Z11TakesStruct1SPi(%struct.S* byval nocapture readonly %input, i32* nocapture %output) #0 {
 entry:
-; CHECK-LABEL @_Z22TakesStruct1SPi
+; CHECK-LABEL: @_Z11TakesStruct1SPi
 ; CHECK:   bitcast %struct.S* %input to i8*
 ; CHECK:   call i8 addrspace(101)* @llvm.nvvm.ptr.gen.to.param.p101i8.p0i8
   %b = getelementptr inbounds %struct.S, %struct.S* %input, i64 0, i32 1
   %0 = load i32, i32* %b, align 4
   store i32 %0, i32* %output, align 4
   ret void
 }
 
 attributes #0 = { nounwind "less-precise-fpmad"="false" "no-frame-pointer-elim"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "stack-protector-buffer-size"="8" "unsafe-fp-math"="false" "use-soft-float"="false" }
 
 !nvvm.annotations = !{!0}
 
 !0 = !{void (%struct.S*, i32*)* @_Z11TakesStruct1SPi, !"kernel", i32 1}

test/CodeGen/NVPTX/lower-kernel-ptr-arg.ll

+; RUN: opt < %s -nvptx-lower-kernel-args -S | FileCheck %s
+
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+target triple = "nvptx64-unknown-unknown"
+
+; Verify that both %input and %output are converted to global pointers and then
+; addrspacecast'ed back to the original type.
+define void @kernel(float* %input, float* %output) {
+; CHECK-LABEL: @kernel(
+; CHECK-NEXT: bitcast float* %output to i8*
+; CHECK-NEXT: call i8 addrspace(1)* @llvm.nvvm.ptr.gen.to.param.p1i8.p0i8
+; CHECK-NEXT: bitcast i8 addrspace(1)* {{[^ ]+}} to float addrspace(1)*
+; CHECK-NEXT: addrspacecast float addrspace(1)* {{[^ ]+}} to float*
+; CHECK-NEXT: bitcast float* %input to i8*
+; CHECK-NEXT: call i8 addrspace(1)* @llvm.nvvm.ptr.gen.to.param.p1i8.p0i8
+; CHECK-NEXT: bitcast i8 addrspace(1)* {{[^ ]+}} to float addrspace(1)*
+; CHECK-NEXT: addrspacecast float addrspace(1)* {{[^ ]+}} to float*
+  %1 = load float, float* %input, align 4
+  store float %1, float* %output, align 4
+  ret void
+}
+
+!nvvm.annotations = !{!0}
+!0 = !{void (float*, float*)* @kernel, !"kernel", i32 1}