[NVPTX] Force minimum alignment of 4 for byval arguments of device-side functions.

[tra]

Summary

Taking address of parameter is legal in PTX and we do generate code that does it
Alas such code currently runs into miscompilation by ptxas on sm_50+ (NVIDIA issue 1789042).
Work around the issue by enforcing minimum alignment on byval arguments of device functions.

The change is effectively a no-op on SASS level for sm_3x where ptxas already aligns local copy by at least 4.

[jlebar]

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1601

I think the condition is when we support cross-TU linking with nvcc, because this breaks ABI compatibility with nvcc? We should be able to do multi-TU compilations just fine, if all of the object files come from clang.

lib/Target/NVPTX/NVPTXISelLowering.cpp

1344

Maybe s/change/condition/?

[tra]

updated comments.

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1601

That's correct. I've updated the comment.

lib/Target/NVPTX/NVPTXISelLowering.cpp

1344

'alignment adjustment' ?

lib/Target/NVPTX/NVPTXAsmPrinter.cpp

for (I = F->arg_begin(), E = F->arg_end(); I != E; ++I, paramIndex++) {
 
     if (isABI || isKernelFunc) {
       // Just print .param .align <a> .b8 .param[size];
       // <a> = PAL.getparamalignment
       // size = typeallocsize of element type
       unsigned align = PAL.getParamAlignment(paramIndex + 1);
       if (align == 0)
         align = DL.getABITypeAlignment(ETy);
-
+      // 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. Hopefully ptxas bug is fixed by then.
+      if (!isKernelFunc && align < 4)
+        align = 4;
       unsigned sz = DL.getTypeAllocSize(ETy);
       O << "\t.param .align " << align << " .b8 ";
       printParamName(I, paramIndex, O);
       O << "[" << sz << "]";
       continue;
     } else {
       // Split the ETy into constituent parts and
       // print .param .b<size> <name> for each part.

lib/Target/NVPTX/NVPTXISelLowering.cpp

SDValue NVPTXTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
 
   bool isABI = (STI.getSmVersion() >= 20);
   assert(isABI && "Non-ABI compilation is not supported");
   if (!isABI)
     return Chain;
   MachineFunction &MF = DAG.getMachineFunction();
   const Function *F = MF.getFunction();
   auto &DL = MF.getDataLayout();
+  bool isKernel = llvm::isKernelFunction(*F);
 
   SDValue tempChain = Chain;
   Chain = DAG.getCALLSEQ_START(Chain,
                                DAG.getIntPtrConstant(uniqueCallSite, dl, true),
                                dl);
   SDValue InFlag = Chain.getValue(1);
 
   unsigned paramCount = 0;
for (unsigned i = 0, e = Args.size(); i != e; ++i, ++OIdx) {
 
     // declare .param .align <align> .b8 .param<n>[<size>];
     unsigned sz = Outs[OIdx].Flags.getByValSize();
     SDVTList DeclareParamVTs = DAG.getVTList(MVT::Other, MVT::Glue);
     unsigned ArgAlign = Outs[OIdx].Flags.getByValAlign();
     // The ByValAlign in the Outs[OIdx].Flags is alway set at this point,
     // so we don't need to worry about natural alignment or not.
     // See TargetLowering::LowerCallTo().
-    SDValue DeclareParamOps[] = {
-      Chain, DAG.getConstant(Outs[OIdx].Flags.getByValAlign(), dl, MVT::i32),
+
+    // Enforce minumum alignment of 4 to work around ptxas miscompile
+    // for sm_50+. See corresponding change in emitFunctionParamList()
+    // for details.
+    if (!isKernel && ArgAlign < 4)
+      ArgAlign = 4;
+    SDValue DeclareParamOps[] = {Chain, DAG.getConstant(ArgAlign, dl, MVT::i32),
                                  DAG.getConstant(paramCount, dl, MVT::i32),
-      DAG.getConstant(sz, dl, MVT::i32), InFlag
-    };
+                                 DAG.getConstant(sz, dl, MVT::i32), InFlag};
     Chain = DAG.getNode(NVPTXISD::DeclareParam, dl, DeclareParamVTs,
                         DeclareParamOps);
     InFlag = Chain.getValue(1);
     for (unsigned j = 0, je = vtparts.size(); j != je; ++j) {
       EVT elemtype = vtparts[j];
       int curOffset = Offsets[j];
       unsigned PartAlign = GreatestCommonDivisor64(ArgAlign, curOffset);
       auto PtrVT = getPointerTy(DAG.getDataLayout());

test/CodeGen/NVPTX/param-align.ll

 
 ;;; Need 4-byte alignment on float2* passed byval
 %struct.float2 = type { float, float }
 define ptx_device void @t3(%struct.float2* byval %x) {
 ; CHECK: .func t3
 ; CHECK: .param .align 4 .b8 t3_param_0[8]
   ret void
 }
+
+;;; Need at least 4-byte alignment in order to avoid miscompilation by
+;;; ptxas for sm_50+
+define ptx_device void @t4(i8* byval %x) {
+; CHECK: .func t4
+; CHECK: .param .align 4 .b8 t4_param_0[1]
+  ret void
+}