[IR] Allow alignstack attribute on return values (#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.

Author: AlexMaclean

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

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]>;

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
+}