Improve kernel arg attrs accuracy

Allow explicitly defined kernel .ptr arg alignments of 1 to be
specified.

Avoid forcing .global annotations for generic pointers on CUDA.

Allow .local pointers, since the PTX specification says they are legal
here.

Avoid outputting unnecessary extra spaces in PTX between .ptr and the
memory space.

Improve test-cases by removing unnecessary function bodies, and adding
more varied alignments + address spaces.

Author: LewisCrawford

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

@@ -1600,37 +1600,33 @@ void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
 
       if (isKernelFunc) {
         if (PTy) {
-          O << "\t.param .u" << PTySizeInBits << " ";
+          O << "\t.param .u" << PTySizeInBits << " .ptr ";
 
-          int addrSpace = PTy->getAddressSpace();
-          const bool IsCUDA =
-              static_cast<NVPTXTargetMachine &>(TM).getDrvInterface() ==
-              NVPTX::CUDA;
-
-          O << ".ptr ";
-          switch (addrSpace) {
+          switch (PTy->getAddressSpace()) {
           default:
-            // Special handling for pointer arguments to kernel
-            // CUDA kernels assume that pointers are in global address space
-            // See:
-            // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parameter-state-space
-            if (IsCUDA)
-              O << " .global ";
             break;
-          case ADDRESS_SPACE_CONST:
-            O << " .const ";
+          case ADDRESS_SPACE_GLOBAL:
+            O << ".global ";
             break;
           case ADDRESS_SPACE_SHARED:
-            O << " .shared ";
+            O << ".shared ";
             break;
-          case ADDRESS_SPACE_GLOBAL:
-            O << " .global ";
+          case ADDRESS_SPACE_CONST:
+            O << ".const ";
+            break;
+          case ADDRESS_SPACE_LOCAL:
+            O << ".local ";
             break;
           }
 
-          Align ParamAlign = I->getParamAlign().valueOrOne();
-          if (ParamAlign != 1 || !IsCUDA)
-            O << ".align " << ParamAlign.value() << " ";
+          const bool IsCUDA =
+              static_cast<NVPTXTargetMachine &>(TM).getDrvInterface() ==
+              NVPTX::CUDA;
+
+          MaybeAlign ParamAlign = I->getParamAlign();
+          if (ParamAlign.has_value() || !IsCUDA)
+            O << ".align " << ParamAlign.valueOrOne().value() << " ";
+
           O << TLI->getParamName(F, paramIndex);
           continue;
         }

llvm/test/CodeGen/NVPTX/i1-param.ll

@@ -8,7 +8,7 @@ target triple = "nvptx-nvidia-cuda"
 
 ; CHECK: .entry foo
 ; CHECK:   .param .u8 foo_param_0
-; CHECK:   .param .u64 .ptr .global foo_param_1
+; CHECK:   .param .u64 .ptr foo_param_1
 define void @foo(i1 %p, ptr %out) {
   %val = zext i1 %p to i32
   store i32 %val, ptr %out

llvm/test/CodeGen/NVPTX/kernel-param-align.ll

@@ -4,42 +4,36 @@
 %struct.Large = type { [16 x double] }
 
 ; CHECK-LABEL: .entry func_align(
-; CHECK: .param .u64 .ptr .global .align 16 func_align_param_0
-; CHECK: .param .u64 .ptr .global .align 16 func_align_param_1
-; CHECK: .param .u64 .ptr .global .align 16 func_align_param_2
-; CHECK: .param .u64 .ptr .shared .align 16 func_align_param_3
+; CHECK: .param .u64 .ptr         .align 1  func_align_param_0
+; CHECK: .param .u64 .ptr         .align 2  func_align_param_1
+; CHECK: .param .u64 .ptr .global .align 4  func_align_param_2
+; CHECK: .param .u64 .ptr .shared .align 8  func_align_param_3
 ; CHECK: .param .u64 .ptr .const  .align 16 func_align_param_4
-define void @func_align(ptr nocapture readonly align 16 %input,
-                        ptr nocapture align 16 %out,
-                        ptr addrspace(1) align 16 %global,
-                        ptr addrspace(3) align 16 %shared,
-                        ptr addrspace(4) align 16 %const) {
+; CHECK: .param .u64 .ptr .local  .align 32 func_align_param_5
+define void @func_align(ptr nocapture readonly align 1 %input,
+                        ptr nocapture align 2 %out,
+                        ptr addrspace(1) align 4 %global,
+                        ptr addrspace(3) align 8 %shared,
+                        ptr addrspace(4) align 16 %const,
+                        ptr addrspace(5) align 32 %local) {
 entry:
-  %0 = addrspacecast ptr %out to ptr addrspace(1)
-  %1 = addrspacecast ptr %input to ptr addrspace(1)
-  %getElem = getelementptr inbounds %struct.Large, ptr addrspace(1) %1, i64 0, i32 0, i64 5
-  %tmp2 = load i32, ptr addrspace(1) %getElem, align 8
-  store i32 %tmp2, ptr addrspace(1) %0, align 4
   ret void
 }
 
 ; CHECK-LABEL: .entry func_noalign(
-; CHECK: .param .u64 .ptr .global func_noalign_param_0
-; CHECK: .param .u64 .ptr .global func_noalign_param_1
+; CHECK: .param .u64 .ptr         func_noalign_param_0
+; CHECK: .param .u64 .ptr         func_noalign_param_1
 ; CHECK: .param .u64 .ptr .global func_noalign_param_2
 ; CHECK: .param .u64 .ptr .shared func_noalign_param_3
-; CHECK: .param .u64 .ptr .const func_noalign_param_4
+; CHECK: .param .u64 .ptr .const  func_noalign_param_4
+; CHECK: .param .u64 .ptr .local  func_noalign_param_5
 define void @func_noalign(ptr nocapture readonly %input,
                           ptr nocapture %out,
                           ptr addrspace(1) %global,
                           ptr addrspace(3) %shared,
-                          ptr addrspace(4) %const) {
+                          ptr addrspace(4) %const,
+                          ptr addrspace(5) %local) {
 entry:
-  %0 = addrspacecast ptr %out to ptr addrspace(1)
-  %1 = addrspacecast ptr %input to ptr addrspace(1)
-  %getElem = getelementptr inbounds %struct.Large, ptr addrspace(1) %1, i64 0, i32 0, i64 5
-  %tmp2 = load i32, ptr addrspace(1) %getElem, align 8
-  store i32 %tmp2, ptr addrspace(1) %0, align 4
   ret void
 }