[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

1425

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

1560–1562

Here...

[pavelkopyl]

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

[pavelkopyl]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1425

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–1621

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–1621

Thank you. Done.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

for (I = F->arg_begin(), E = F->arg_end(); I != E; ++I, paramIndex++) {
     Type *ETy = PAL.getParamByValType(paramIndex);
     assert(ETy && "Param should have byval type");
 
     if (isABI || isKernelFunc) {
       // Just print .param .align <a> .b8 .param[size];
       // <a>  = optimal alignment for the element type; always multiple of
       //        PAL.getParamAlignment
       // size = typeallocsize of element type
-      Align OptimalAlign = getOptimalAlignForParam(ETy);
+      Align OptimalAlign =
+          isKernelFunc
+              ? getOptimalAlignForParam(ETy)
+              : TLI->getFunctionByValParamAlign(
+                    F, ETy, PAL.getParamAlignment(paramIndex).valueOrOne(), DL);
 
-      // Work around a bug in ptxas. When PTX code takes address of
-      // byval parameter with alignment < 4, ptxas generates code to
-      // spill argument into memory. Alas on sm_50+ ptxas generates
-      // SASS code that fails with misaligned access. To work around
-      // the problem, make sure that we align byval parameters by at
-      // least 4. Matching change must be made in LowerCall() where we
-      // prepare parameters for the call.
-      //
-      // TODO: this will need to be undone when we get to support multi-TU
-      // device-side compilation as it breaks ABI compatibility with nvcc.
-      // Hopefully ptxas bug is fixed by then.
-      if (!isKernelFunc && OptimalAlign < Align(4))
-        OptimalAlign = Align(4);
       unsigned sz = DL.getTypeAllocSize(ETy);
       O << "\t.param .align " << OptimalAlign.value() << " .b8 ";
       printParamName(I, paramIndex, O);
       O << "[" << sz << "]";
       continue;
     } else {
       // Split the ETy into constituent parts and
       // print .param .b<size> <name> for each part.
       // Further, if a part is vector, print the above for

llvm/lib/Target/NVPTX/NVPTXISelLowering.h

public:
   /// ensures that we can effectively vectorize their loads & stores. We can
   /// increase alignment only if the function has internal or has private
   /// linkage as for other linkage types callers may already rely on default
   /// alignment. To allow using 128-bit vectorized loads/stores, this function
   /// ensures that alignment is 16 or greater.
   Align getFunctionParamOptimizedAlign(const Function *F, Type *ArgTy,
                                        const DataLayout &DL) const;
 
+  /// Helper for computing alignment of a device function byval parameter.
+  Align getFunctionByValParamAlign(const Function *F, Type *ArgTy,
+                                   Align InitialAlign,
+                                   const DataLayout &DL) const;
+
   /// isLegalAddressingMode - Return true if the addressing mode represented
   /// by AM is legal for this target, for a load/store of the specified type
   /// Used to guide target specific optimizations, like loop strength
   /// reduction (LoopStrengthReduce.cpp) and memory optimization for
   /// address mode (CodeGenPrepare.cpp)
   bool isLegalAddressingMode(const DataLayout &DL, const AddrMode &AM, Type *Ty,
                              unsigned AS,
                              Instruction *I = nullptr) const override;

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

if (!Outs[OIdx].Flags.isByVal()) {
         sz = 32;
       else
         sz = Ty->getPrimitiveSizeInBits();
       O << ".param .b" << sz << " ";
       O << "_";
       continue;
     }
 
-    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);
+    Align InitialAlign = Outs[OIdx].Flags.getNonZeroByValAlign();
+    Align ParamByValAlign =
+        getFunctionByValParamAlign(F, ETy, InitialAlign, DL);
 
     O << ".param .align " << ParamByValAlign.value() << " .b8 ";
     O << "_";
     O << "[" << Outs[OIdx].Flags.getByValSize() << "]";
   }
 
   if (VAInfo)
     O << (first ? "" : ",") << " .param .align " << VAInfo->second
       << " .b8 _[]\n";
   O << ")";
   if (shouldEmitPTXNoReturn(&CB, *nvTM))
     O << " .noreturn";
   O << ";";
for (unsigned i = 0, e = Args.size(); i != e; ++i, ++OIdx) {
     Type *ETy = (IsByVal ? Args[i].IndirectType : Ty);
     ComputePTXValueVTs(*this, DL, ETy, VTs, &Offsets, IsByVal ? 0 : VAOffset);
 
     Align ArgAlign;
     if (IsByVal) {
       // The ByValAlign in the Outs[OIdx].Flags is always set at this point,
       // so we don't need to worry whether it's naturally aligned or not.
       // See TargetLowering::LowerCallTo().
-      ArgAlign = Outs[OIdx].Flags.getNonZeroByValAlign();
-
-      // 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));
+      Align InitialAlign = Outs[OIdx].Flags.getNonZeroByValAlign();
+      ArgAlign = getFunctionByValParamAlign(CB->getCalledFunction(), ETy,
+                                            InitialAlign, DL);
       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));
if (!F || !F->hasLocalLinkage() ||
                          /*IgnoreAssumeLikeCalls=*/true,
                          /*IgnoreLLVMUsed=*/true))
     return Align(ABITypeAlign);
 
   assert(!isKernelFunction(*F) && "Expect kernels to have non-local linkage");
   return Align(std::max(uint64_t(16), ABITypeAlign));
 }
 
+/// Helper for computing alignment of a device function byval parameter.
+Align NVPTXTargetLowering::getFunctionByValParamAlign(
+    const Function *F, Type *ArgTy, Align InitialAlign,
+    const DataLayout &DL) const {
+  Align ArgAlign = InitialAlign;
+  // Try to increase alignment to enhance vectorization options.
+  if (F)
+    ArgAlign = std::max(ArgAlign, getFunctionParamOptimizedAlign(F, ArgTy, DL));
+
+  // Work around a bug in ptxas. When PTX code takes address of
+  // byval parameter with alignment < 4, ptxas generates code to
+  // spill argument into memory. Alas on sm_50+ ptxas generates
+  // SASS code that fails with misaligned access. To work around
+  // the problem, make sure that we align byval parameters by at
+  // least 4.
+  // TODO: this will need to be undone when we get to support multi-TU
+  // device-side compilation as it breaks ABI compatibility with nvcc.
+  // Hopefully ptxas bug is fixed by then.
+  ArgAlign = std::max(ArgAlign, Align(4));
+
+  return ArgAlign;
+}
+
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.
 /// Used to guide target specific optimizations, like loop strength reduction
 /// (LoopStrengthReduce.cpp) and memory optimization for address mode
 /// (CodeGenPrepare.cpp)
 bool NVPTXTargetLowering::isLegalAddressingMode(const DataLayout &DL,
                                                 const AddrMode &AM, Type *Ty,
                                                 unsigned AS, Instruction *I) const {

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"))

llvm/test/CodeGen/NVPTX/param-align.ll

 ; CHECK: .func t5
 ; CHECK: .param .align 4 .b8 t5_param_0[1]
 ; CHECK: {
 ; CHECK: .param .align 4 .b8 param0[1];
 ; CHECK: call.uni
   call void @t4(ptr byval(i8) %x)
   ret void
 }
+
+;;; Make sure we adjust alignment for a function prototype
+;;; in case of an inderect call.
+
+declare ptr @getfp(i32 %n)
+%struct.half2 = type { half, half }
+define ptx_device void @t6() {
+; CHECK: .func t6
+  %fp = call ptr @getfp(i32 0)
+; CHECK: prototype_2 : .callprototype ()_ (.param .align 8 .b8 _[8]);
+  call void %fp(ptr byval(double) null);
+
+  %fp2 = call ptr @getfp(i32 1)
+; CHECK: prototype_4 : .callprototype ()_ (.param .align 4 .b8 _[4]);
+  call void %fp(ptr byval(%struct.half2) null);
+
+  %fp3 = call ptr @getfp(i32 2)
+; CHECK: prototype_6 : .callprototype ()_ (.param .align 4 .b8 _[1]);
+  call void %fp(ptr byval(i8) null);
+  ret void
+}