[llvm] 6d05831 - Enable .ptr .global .align attributes for kernel attributes for CUDA (#114874)

[via llvm-commits]

Author: Lewis Crawford
Date: 2024-11-15T12:40:53Z
New Revision: 6d058317e60c25b71df8b8dc45b69e5202362678

URL: https://github.com/llvm/llvm-project/commit/6d058317e60c25b71df8b8dc45b69e5202362678
DIFF: https://github.com/llvm/llvm-project/commit/6d058317e60c25b71df8b8dc45b69e5202362678.diff

LOG: Enable .ptr .global .align attributes for kernel attributes for CUDA (#114874)

Emit .ptr, .address-space, and .align attributes for kernel
args in CUDA (previously handled only for OpenCL).

This allows for more vectorization opportunities if the PTX consumer
is able to know about the pointer alignments.

If no alignment is explicitly specified, .align 1 will be emitted
to match the LLVM IR semantics in this case.

PTX ISA doc -
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#kernel-parameter-attribute-ptr

This is a rework of the original patch proposed in #79646

---------

Co-authored-by: Vandana <vandanak at nvidia.com>

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

Modified: 
    llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
    llvm/test/CodeGen/NVPTX/i1-param.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
index f60091f38f246a..ac3e55df53c20a 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
@@ -1600,30 +1600,27 @@ void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
 
       if (isKernelFunc) {
         if (PTy) {
-          // Special handling for pointer arguments to kernel
-          O << "\t.param .u" << PTySizeInBits << " ";
-
-          if (static_cast<NVPTXTargetMachine &>(TM).getDrvInterface() !=
-              NVPTX::CUDA) {
-            int addrSpace = PTy->getAddressSpace();
-            switch (addrSpace) {
-            default:
-              O << ".ptr ";
-              break;
-            case ADDRESS_SPACE_CONST:
-              O << ".ptr .const ";
-              break;
-            case ADDRESS_SPACE_SHARED:
-              O << ".ptr .shared ";
-              break;
-            case ADDRESS_SPACE_GLOBAL:
-              O << ".ptr .global ";
-              break;
-            }
-            Align ParamAlign = I->getParamAlign().valueOrOne();
-            O << ".align " << ParamAlign.value() << " ";
+          O << "\t.param .u" << PTySizeInBits << " .ptr";
+
+          switch (PTy->getAddressSpace()) {
+          default:
+            break;
+          case ADDRESS_SPACE_GLOBAL:
+            O << " .global";
+            break;
+          case ADDRESS_SPACE_SHARED:
+            O << " .shared";
+            break;
+          case ADDRESS_SPACE_CONST:
+            O << " .const";
+            break;
+          case ADDRESS_SPACE_LOCAL:
+            O << " .local";
+            break;
           }
-          O << TLI->getParamName(F, paramIndex);
+
+          O << " .align " << I->getParamAlign().valueOrOne().value();
+          O << " " << TLI->getParamName(F, paramIndex);
           continue;
         }
 

diff  --git a/llvm/test/CodeGen/NVPTX/i1-param.ll b/llvm/test/CodeGen/NVPTX/i1-param.ll
index 375752b619a581..0878eb4fcd47b3 100644
--- a/llvm/test/CodeGen/NVPTX/i1-param.ll
+++ b/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 foo_param_1
+; CHECK:   .param .u64 .ptr .align 1 foo_param_1
 define void @foo(i1 %p, ptr %out) {
   %val = zext i1 %p to i32
   store i32 %val, ptr %out

diff  --git a/llvm/test/CodeGen/NVPTX/kernel-param-align.ll b/llvm/test/CodeGen/NVPTX/kernel-param-align.ll
new file mode 100644
index 00000000000000..f29b5823aa5fe6
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/kernel-param-align.ll
@@ -0,0 +1,42 @@
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_60 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_60 | %ptxas-verify %}
+
+%struct.Large = type { [16 x double] }
+
+; CHECK-LABEL: .entry func_align(
+; 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
+; 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:
+  ret void
+}
+
+; CHECK-LABEL: .entry func_noalign(
+; CHECK: .param .u64 .ptr         .align 1 func_noalign_param_0
+; CHECK: .param .u64 .ptr         .align 1 func_noalign_param_1
+; CHECK: .param .u64 .ptr .global .align 1 func_noalign_param_2
+; CHECK: .param .u64 .ptr .shared .align 1 func_noalign_param_3
+; CHECK: .param .u64 .ptr .const  .align 1 func_noalign_param_4
+; CHECK: .param .u64 .ptr .local  .align 1 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(5) %local) {
+entry:
+  ret void
+}
+
+!nvvm.annotations = !{!0, !1}
+!0 = !{ptr @func_align, !"kernel", i32 1}
+!1 = !{ptr @func_noalign, !"kernel", i32 1}