[NVPTX] cap param alignment at 128 (max supported by ptx)

[AlexMaclean]

Cap the alignment to 128 bytes as that is the maximum alignment supported by PTX. The restriction is mentioned in the parameter passing section (Note D) of the PTX Writer's Guide to Interoperability

D. The alignment must be 1, 2, 4, 8, 16, 32, 64, or 128 bytes.

[llvmbot]

@llvm/pr-subscribers-backend-nvptx

Author: Alex MacLean (AlexMaclean)

Changes

Cap the alignment to 128 bytes as that is the maximum alignment supported by PTX. The restriction is mentioned in the parameter passing section (Note D) of the PTX Writer's Guide to Interoperability

> D. The alignment must be 1, 2, 4, 8, 16, 32, 64, or 128 bytes.

---

Full diff: https://github.com/llvm/llvm-project/pull/96117.diff

2 Files Affected:

• (modified) llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp (+5-3)
• (added) llvm/test/CodeGen/NVPTX/max-align.ll (+14)

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index f4ef7c9914f13..982c191875750 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -5038,7 +5038,9 @@ bool NVPTXTargetLowering::getTgtMemIntrinsic(
 /// ensures that alignment is 16 or greater.
 Align NVPTXTargetLowering::getFunctionParamOptimizedAlign(
     const Function *F, Type *ArgTy, const DataLayout &DL) const {
-  const uint64_t ABITypeAlign = DL.getABITypeAlign(ArgTy).value();
+  // Capping the alignment to 128 bytes as that is the maximum alignment
+  // supported by PTX.
+  const Align ABITypeAlign = std::min(Align(128), DL.getABITypeAlign(ArgTy));
 
   // If a function has linkage different from internal or private, we
   // must use default ABI alignment as external users rely on it. Same
@@ -5048,10 +5050,10 @@ Align NVPTXTargetLowering::getFunctionParamOptimizedAlign(
                          /*IgnoreCallbackUses=*/false,
                          /*IgnoreAssumeLikeCalls=*/true,
                          /*IgnoreLLVMUsed=*/true))
-    return Align(ABITypeAlign);
+    return ABITypeAlign;
 
   assert(!isKernelFunction(*F) && "Expect kernels to have non-local linkage");
-  return Align(std::max(uint64_t(16), ABITypeAlign));
+  return std::max(Align(16), ABITypeAlign);
 }
 
 /// Helper for computing alignment of a device function byval parameter.
diff --git a/llvm/test/CodeGen/NVPTX/max-align.ll b/llvm/test/CodeGen/NVPTX/max-align.ll
new file mode 100644
index 0000000000000..c8b1cb12dee5f
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/max-align.ll
@@ -0,0 +1,14 @@
+; RUN: llc < %s -march=nvptx64 -O0 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -O0 | %ptxas-verify %}
+
+
+; CHECK: .visible .func  (.param .align 128 .b8 func_retval0[256]) repro()
+define <64 x i32> @repro() {
+
+  ; CHECK: .param .align 128 .b8 retval0[256];
+  %1 = tail call <64 x i32> @test(i32 0)
+  ret <64 x i32> %1
+}
+
+; Function Attrs: nounwind
+declare <64 x i32> @test(i32)

[AlexMaclean]

@jlebar ping

[jlebar]

Sorry I missed this one. LGTM!