[NVPTX] Enforce minumum alignment of 4 for byval parametrs in a function prototype

[pavelkopyl]

Summary

As a result, we have identical alignment calculation of byval parameters for

• LowerCall() - getting alignment of an argument (.param)
• emitFunctionParamList() - getting alignment of a parameter (.param) in a function declaration
• getPrototype() - getting alignment of a parameter (.param) in a function prototypes that is used for indirect calls

This change is required to avoid ptxas error:
'Alignment of argument does not match formal parameter'. This
error happens even in cases where it logically shouldn't.
For instance:

.param .align 4 .b8 param0[4];
...
callprototype ()_ (.param .align 2 .b8 _[4]);
...

Here we allocate 'param0' with alignment of 4 and it should be
fine to pass it to a function that requires minimum alignment of 2.
At least ptxas v12.0 rejects this code.

[ldrumm]

In your commit message:

• s/parametrs/parameters

I'd prefer to split off a separate test that only checks the alignment of various calls, rather than reuse the byval bitcast test. Otherwise LGTM, but I'd wait on someone more familiar with ptx, e.g. ~tra

[pavelkopyl]

• Rebased
• Added fixes according to review notes

[pavelkopyl]

In D140581#4015109, @ldrumm wrote:

In your commit message:

• s/parametrs/parameters

I'd prefer to split off a separate test that only checks the alignment of various calls, rather than reuse the byval bitcast test. Otherwise LGTM, but I'd wait on someone more familiar with ptx, e.g. ~tra

Thank you for review. I added new case for checking .param alignment in function prototype to llvm/test/CodeGen/NVPTX/param-align.ll

[tra]

Looks good overall.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1428

This seems to be the same adjustment we do in LowerCall below. Perhaps we can consolidate alignment calculation into a single helper function.

1575

Here...

[pavelkopyl]

• Rebased
• Added hepler function for computing byval param alignments of device functions

[pavelkopyl]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1428

Thank you for suggestion. Actually I introduced a helper to use in all three places: emitFunctionParamList (in AsmPrinter), getPrototype and LowerCall. Is that OK?

[tra]

LGTM. Thank you!

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1615–1623 ↗

(On Diff #487011)

Nit. This all could be folded into

Align OptimalAlign = isKernelFunc 
        ? getOptimalAlignForParam(ETy);
        : TLI->getFunctionByValParamAlign(F, ETy, PAL.getParamAlignment(paramIndex).valueOrOne(), DL);

[pavelkopyl]

• Fixes according to review notes

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1615–1623 ↗

(On Diff #487011)

Thank you. Done.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

for (unsigned i = 0, OIdx = 0; i != NumArgs; ++i, ++OIdx) {
     Align ParamByValAlign = Outs[OIdx].Flags.getNonZeroByValAlign();
 
     // Try to increase alignment. This code matches logic in LowerCall when
     // alignment increase is performed to increase vectorization options.
     Type *ETy = Args[i].IndirectType;
     Align AlignCandidate = getFunctionParamOptimizedAlign(F, ETy, DL);
     ParamByValAlign = std::max(ParamByValAlign, AlignCandidate);
 
+    // Enforce minumum alignment of 4 to work around ptxas miscompile
+    // for sm_50+. See corresponding alignment adjustment in
+    // emitFunctionParamList() for details.
+    ParamByValAlign = std::max(ParamByValAlign, Align(4));
+
     O << ".param .align " << ParamByValAlign.value() << " .b8 ";
     O << "_";
     O << "[" << Outs[OIdx].Flags.getByValSize() << "]";
   }
 
   if (VAInfo)
     O << (first ? "" : ",") << " .param .align " << VAInfo->second
       << " .b8 _[]\n";
if (IsByVal) {
       // Try to increase alignment to enhance vectorization options.
       if (const Function *DirectCallee = CB->getCalledFunction())
         ArgAlign = std::max(
             ArgAlign, getFunctionParamOptimizedAlign(DirectCallee, ETy, DL));
 
       // Enforce minumum alignment of 4 to work around ptxas miscompile
       // for sm_50+. See corresponding alignment adjustment in
       // emitFunctionParamList() for details.
       ArgAlign = std::max(ArgAlign, Align(4));
       if (IsVAArg)
         VAOffset = alignTo(VAOffset, ArgAlign);
     } else {
       ArgAlign = getArgumentAlignment(Callee, CB, Ty, ParamCount + 1, DL);
     }
 
     unsigned TypeSize =
         (IsByVal ? Outs[OIdx].Flags.getByValSize() : DL.getTypeAllocSize(Ty));

llvm/test/CodeGen/NVPTX/call_bitcast_byval.ll

 ; CHECK-NEXT: .param .align 4 .b8 callee_param_0[4]
 define internal void @callee(ptr byval(%"class.complex") %byval_arg) {
   ret void
 }
 define void @boom() {
   %fp = call ptr @usefp(ptr @callee)
   ; CHECK: .param .align 4 .b8 param0[4];
   ; CHECK: st.param.v2.b16 [param0+0]
-  ; CHECK: .callprototype ()_ (.param .align 2 .b8 _[4]);
+  ; CHECK: .callprototype ()_ (.param .align 4 .b8 _[4]);
   call void %fp(ptr byval(%"class.complex") null)
   ret void
 }
 
-declare %complex_half @_Z20__spirv_GroupCMulKHRjjN5__spv12complex_halfE()
+declare %complex_half @_Z20__spirv_GroupCMulKHRjjN5__spv12complex_halfE(i32, i32, ptr byval(%"class.complex"))