[llvm] [NVPTX] Don't use underlying alignment to align param (PR #96793)

[Hugh Delaney via llvm-commits]

https://github.com/hdelan updated https://github.com/llvm/llvm-project/pull/96793

>From f84714b57999d4480fe468efbe7686e48b6e2801 Mon Sep 17 00:00:00 2001
From: Hugh Delaney <hugh.delaney at codeplay.com>
Date: Wed, 26 Jun 2024 17:15:44 +0100
Subject: [PATCH] Don't use underlying alignment to align param

Previously, if a ptr had align N, then the NVPTX lowering was taking
this align N to refer to the alignment of the pointer type itself, as
opposed to the alignment of the memory that it points to.

As such, if a kernel of the form:

define void @foo(ptr align 4 %_arg_ptr)

Would take align 4 to be the alignment of the parameter, which would
result in breaking the ld.param into two separate loads.

	ld.param.u32 	%rd1, [foo_param_0+4];
	shl.b64 	%rd2, %rd1, 32;
	ld.param.u32 	%rd3, [foo_param_0];
	or.b64  	%rd4, %rd2, %rd3;

It isn't necessary as far as I can tell from the PTX ISA documents to
specify the alignment of the parameters themselves. So this patch
changes the codegen to the better:

	ld.param.u64 	%rd1, [foo_param_0];
---
 llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp |  4 +--
 llvm/test/CodeGen/NVPTX/param-align.ll      | 39 +++++++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 982c191875750..63cbdb0acfab6 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -3232,9 +3232,7 @@ SDValue NVPTXTargetLowering::LowerFormalArguments(
             if (NumElts != 1)
               return std::nullopt;
             Align PartAlign =
-                (Offsets[parti] == 0 && PAL.getParamAlignment(i))
-                    ? PAL.getParamAlignment(i).value()
-                    : DL.getABITypeAlign(EltVT.getTypeForEVT(F->getContext()));
+                DL.getABITypeAlign(EltVT.getTypeForEVT(F->getContext()));
             return commonAlignment(PartAlign, Offsets[parti]);
           }();
           SDValue P = DAG.getLoad(VecVT, dl, Root, VecAddr,
diff --git a/llvm/test/CodeGen/NVPTX/param-align.ll b/llvm/test/CodeGen/NVPTX/param-align.ll
index 5435ee238c88d..b8ef4a87bc752 100644
--- a/llvm/test/CodeGen/NVPTX/param-align.ll
+++ b/llvm/test/CodeGen/NVPTX/param-align.ll
@@ -69,3 +69,42 @@ define ptx_device void @t6() {
   call void %fp(ptr byval(i8) null);
   ret void
 }
+
+; CHECK-LABEL: .func check_ptr_align1(
+; CHECK: 	ld.param.u64 	%rd1
+; CHECK-NOT: 	ld.param.u8
+; CHECK: 	ret;
+define void @check_ptr_align1(ptr align 1 %_arg_ptr) {
+entry:
+  store i32 1, ptr %_arg_ptr, align 1
+  ret void
+}
+
+; CHECK-LABEL: .func check_ptr_align2(
+; CHECK: 	ld.param.u64 	%rd1
+; CHECK-NOT: 	ld.param.u16
+; CHECK: 	ret;
+define void @check_ptr_align2(ptr align 2 %_arg_ptr) {
+entry:
+  store i32 2, ptr %_arg_ptr, align 2
+  ret void
+}
+
+; CHECK-LABEL: .func check_ptr_align4(
+; CHECK: 	ld.param.u64 	%rd1
+; CHECK-NOT: 	ld.param.u32
+; CHECK: 	ret;
+define void @check_ptr_align4(ptr align 4 %_arg_ptr) {
+entry:
+  store i32 4, ptr %_arg_ptr, align 4
+  ret void
+}
+
+; CHECK-LABEL: .func check_ptr_align8(
+; CHECK: 	ld.param.u64 	%rd1
+; CHECK: 	ret;
+define void @check_ptr_align8(ptr align 8 %_arg_ptr) {
+entry:
+  store i32 8, ptr %_arg_ptr, align 8
+  ret void
+}