[llvm] [IR] Allow alignstack attribute on return values (PR #130439)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/130439

the PTX target allows an alignment to be specified on both return values and parameters to allow for more efficient vectorized stores. Currently we represent these parameter alignments via the "alignstack" attribute, but must fall back to metadata for the return value. This PR allows "alignstack" on return values as well.

>From 1ba3abb3d4ff7fdcfa86ed5d5fe4ad230c284d80 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Tue, 4 Mar 2025 17:46:09 +0000
Subject: [PATCH] [IR] Allow alignstack attribute on return values

---
 llvm/docs/LangRef.rst                      | 12 ++++++------
 llvm/include/llvm/IR/Attributes.td         |  2 +-
 llvm/test/CodeGen/NVPTX/param-overalign.ll | 11 ++++++++---
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 33c85c7ba9d29..8216b15512d9f 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -1616,12 +1616,12 @@ Currently, only the following parameter attributes are defined:
 
 ``alignstack(<n>)``
     This indicates the alignment that should be considered by the backend when
-    assigning this parameter to a stack slot during calling convention
-    lowering. The enforcement of the specified alignment is target-dependent,
-    as target-specific calling convention rules may override this value. This
-    attribute serves the purpose of carrying language specific alignment
-    information that is not mapped to base types in the backend (for example,
-    over-alignment specification through language attributes).
+    assigning this parameter or return value to a stack slot during calling
+    convention lowering. The enforcement of the specified alignment is
+    target-dependent, as target-specific calling convention rules may override
+    this value. This attribute serves the purpose of carrying language specific
+    alignment information that is not mapped to base types in the backend (for
+    example, over-alignment specification through language attributes).
 
 ``allocalign``
     The function parameter marked with this attribute is the alignment in bytes of the
diff --git a/llvm/include/llvm/IR/Attributes.td b/llvm/include/llvm/IR/Attributes.td
index 70b9a2c488d3e..fb94926043fc7 100644
--- a/llvm/include/llvm/IR/Attributes.td
+++ b/llvm/include/llvm/IR/Attributes.td
@@ -291,7 +291,7 @@ def SExt : EnumAttr<"signext", IntersectPreserve, [ParamAttr, RetAttr]>;
 
 /// Alignment of stack for function (3 bits)  stored as log2 of alignment with
 /// +1 bias 0 means unaligned (different from alignstack=(1)).
-def StackAlignment : IntAttr<"alignstack", IntersectPreserve, [FnAttr, ParamAttr]>;
+def StackAlignment : IntAttr<"alignstack", IntersectPreserve, [FnAttr, ParamAttr, RetAttr]>;
 
 /// Function can be speculated.
 def Speculatable : EnumAttr<"speculatable", IntersectAnd, [FnAttr]>;
diff --git a/llvm/test/CodeGen/NVPTX/param-overalign.ll b/llvm/test/CodeGen/NVPTX/param-overalign.ll
index b71b50a73f7cb..374475a29ffa1 100644
--- a/llvm/test/CodeGen/NVPTX/param-overalign.ll
+++ b/llvm/test/CodeGen/NVPTX/param-overalign.ll
@@ -45,7 +45,7 @@ define float @caller_md(float %a, float %b) {
   ret float %r
 }
 
-define float @callee_md(%struct.float2 %a) {
+define float @callee_md(%struct.float2 alignstack(8) %a) {
 ; CHECK-LABEL: .visible .func  (.param .b32 func_retval0) callee_md(
 ; CHECK-NEXT:         .param .align 8 .b8 callee_md_param_0[8]
 ; CHECK-NEXT: )
@@ -105,5 +105,10 @@ define float @callee(%struct.float2 alignstack(8) %a ) {
   ret float %2
 }
 
-!nvvm.annotations = !{!0}
-!0 = !{ptr @callee_md, !"align", i32 u0x00010008}
+define alignstack(8) %struct.float2 @aligned_return(%struct.float2 %a ) {
+; CHECK-LABEL: .visible .func  (.param .align 8 .b8 func_retval0[8]) aligned_return(
+; CHECK-NEXT:         .param .align 4 .b8 aligned_return_param_0[8]
+; CHECK-NEXT: )
+; CHECK-NEXT: {
+  ret %struct.float2 %a
+}