[NVPTX] Fix a segfault for bitcasted calls with byval params

[ldrumm]

Summary

`getFunctionParamOptimizedAlign` was being passed a null function
  argument when getting the callee of a bitcasted function symbol. This is
  because `CallBase::getCalledFunction` does not look through bitcasts.
  
  There is already code to handle this case in
  `NVPTXTargetLowering::getArgumentAlignment`, which is now hoisted into
  an NVPTX util.
  
  The alignment computation now gracefully handles computing alignment of
  virtual functions with a check for null.

[tra]

LGTM with few nits.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1429

Would we ever have an instruction other than Call which would also carry explicit alignment metadata?

4329

Nit: if (!(F && F->hasLocalLinkage())) would be a bit easier to understand, IMO.

llvm/lib/Target/NVPTX/NVPTXUtilities.cpp

328

Style nit: no braces for single-statement if bodies. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

[ldrumm]

Code review suggestions

[ldrumm]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1429

allocas, loads, and stores, though I'm not sure if that's what you're asking. Could you expand a little on what you mean? The code here is only dealing with argument alignment so in isolation the answer is "no", but I have a feeling you might be hinting at something deeper ?

llvm/lib/Target/NVPTX/NVPTXUtilities.cpp

328

Yes. I actually prefer that style too. Too many competing style guides in my head!

[tra]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1429

If CallBase can be an instruction other than Call, will we be able to get argument alignment via getAlign() and if so, should we?

Another observation here is that previously digging through bitcasts was done under if ( dyn_cast<CallInst>(CB)) while now it's done for any CB instruction. It looks like it should work (getMaybeBitcastedCallee will bail out on other instructions, if it can't get ), but was that done that intentionally? If looking through bitcasts is applicable to the Call instruction only, I'd prefer to keep it under the if to make it obvious.

[nikic]

llvm/lib/Target/NVPTX/NVPTXUtilities.cpp

349

The whole function can be reduced to dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts()).

[ldrumm]

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

1429

I don't think this can ever be anything other than a call. It's within if (const auto *CI = dyn_cast<CallInst>(CB)) {. The getMaybeBitcastedCallee is indeed used in more possible places but I figured this to be safe in the current usage.

llvm/lib/Target/NVPTX/NVPTXUtilities.cpp

349

Thanks. TIL about stripPointerCasts

[asavonic]

LGTM.

[ldrumm]

Merged in https://reviews.llvm.org/rG940fa35ece5294a115a2fdba89ef6c095d90df0f wherein I forgot to note the differential revision

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

Align NVPTXTargetLowering::getArgumentAlignment(SDValue Callee,
 
   if (!DirectCallee) {
     // We don't have a direct function symbol, but that may be because of
     // constant cast instructions in the call.
 
     // With bitcast'd call targets, the instruction will be the call
     if (const auto *CI = dyn_cast<CallInst>(CB)) {
       // Check if we have call alignment metadata
       if (getAlign(*CI, Idx, Alignment))
         return Align(Alignment);
-
-      const Value *CalleeV = CI->getCalledOperand();
-      // Ignore any bitcast instructions
-      while (isa<ConstantExpr>(CalleeV)) {
-        const ConstantExpr *CE = cast<ConstantExpr>(CalleeV);
-        if (!CE->isCast())
-          break;
-        // Look through the bitcast
-        CalleeV = cast<ConstantExpr>(CalleeV)->getOperand(0);
-      }
-
-      // We have now looked past all of the bitcasts.  Do we finally have a
-      // Function?
-      if (const auto *CalleeF = dyn_cast<Function>(CalleeV))
-        DirectCallee = CalleeF;
     }
+    DirectCallee = getMaybeBitcastedCallee(CB);
   }
 
   // Check for function alignment information if we found that the
   // ultimate target is a Function
   if (DirectCallee) {
     if (getAlign(*DirectCallee, Idx, Alignment))
       return Align(Alignment);
     // If alignment information is not available, fall back to the
for (unsigned i = 0, e = Args.size(); i != e; ++i, ++OIdx) {
     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.
       ArgAlign = std::max(ArgAlign, getFunctionParamOptimizedAlign(
-                                        CB->getCalledFunction(), ETy, DL));
+                                        getMaybeBitcastedCallee(CB), 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));
     } else {
       ArgAlign = getArgumentAlignment(Callee, CB, Ty, ParamCount + 1, DL);
     }

 /// alignment. To allow using 128-bit vectorized loads/stores, this function
 /// ensures that alignment is 16 or greater.
 Align NVPTXTargetLowering::getFunctionParamOptimizedAlign(
     const Function *F, Type *ArgTy, const DataLayout &DL) const {
   const uint64_t ABITypeAlign = DL.getABITypeAlign(ArgTy).value();
 
   // If a function has linkage different from internal or private, we
   // must use default ABI alignment as external users rely on it.
-  if (!F->hasLocalLinkage())
+  if (!F || !F->hasLocalLinkage())
     return Align(ABITypeAlign);
 
   assert(!isKernelFunction(*F) && "Expect kernels to have non-local linkage");
   return Align(std::max(uint64_t(16), ABITypeAlign));
 }
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.

llvm/lib/Target/NVPTX/NVPTXUtilities.h

 bool getReqNTIDz(const Function &, unsigned &);
 
 bool getMinCTASm(const Function &, unsigned &);
 bool getMaxNReg(const Function &, unsigned &);
 bool isKernelFunction(const Function &);
 
 bool getAlign(const Function &, unsigned index, unsigned &);
 bool getAlign(const CallInst &, unsigned index, unsigned &);
+Function *getMaybeBitcastedCallee(const CallBase *CB);
 
 // PTX ABI requires all scalar argument/return values to have
 // bit-size as a power of two of at least 32 bits.
 inline unsigned promoteScalarArgumentSize(unsigned size) {
   if (size <= 32)
     return 32;
   else if (size <= 64)
     return 64;
   else
     return size;
 }
 }
 
 #endif

llvm/lib/Target/NVPTX/NVPTXUtilities.cpp

for (int i = 0, n = alignNode->getNumOperands(); i < n; i++) {
           return false;
         }
       }
     }
   }
   return false;
 }
 
+Function *getMaybeBitcastedCallee(const CallBase *CB) {
+  if (Function *DirectCallee = CB->getCalledFunction()) {
+    return DirectCallee;
+  }
+
+  // We don't have a direct function symbol, but that may be because of
+  // constant cast instructions in the call.
+  // With bitcast'd call targets, the instruction will be the call
+  auto *CI = dyn_cast<CallInst>(CB);
+  if (!CI)
+    return nullptr;
+
+  Value *CalleeV = CI->getCalledOperand();
+  // Ignore any bitcast instructions
+  while (isa<ConstantExpr>(CalleeV)) {
+    const ConstantExpr *CE = cast<ConstantExpr>(CalleeV);
+    if (!CE->isCast())
+      break;
+    // Look through the bitcast
+    CalleeV = cast<ConstantExpr>(CalleeV)->getOperand(0);
+  }
+  // We have now looked past all of the bitcasts.  Do we finally have a
+  // Function?
+  return dyn_cast<Function>(CalleeV);
+}
+
 } // namespace llvm

llvm/test/CodeGen/NVPTX/call_bitcast_byval.ll

+; RUN: llc < %s -march=nvptx -mcpu=sm_50 -verify-machineinstrs | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_50 -verify-machineinstrs | %ptxas-verify %}
+
+; calls with a bitcasted function symbol should be fine, but in combination with
+; a byval attribute were causing a segfault during isel. This testcase was
+; reduced from a SYCL kernel using aggregate types which ended up being passed
+; `byval`
+
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+%"class.complex" = type { %"class.sycl::_V1::detail::half_impl::half", %"class.sycl::_V1::detail::half_impl::half" }
+%"class.sycl::_V1::detail::half_impl::half" = type { half }
+%complex_half = type { half, half }
+
+define weak_odr void @foo() {
+entry:
+  %call.i.i.i = tail call %"class.complex" bitcast (%complex_half ()* @_Z20__spirv_GroupCMulKHRjjN5__spv12complex_halfE to %"class.complex" (i32, i32, %"class.complex"*)*)(i32 0, i32 0, %"class.complex"* byval(%"class.complex") null)
+  ret void
+}
+
+declare %complex_half @_Z20__spirv_GroupCMulKHRjjN5__spv12complex_halfE()
+
+; CHECK: .param .align 4 .b8 param2[4];
+; CHECK: st.param.v2.b16         [param2+0], {%h2, %h1};
+; CHECK: .param .align 2 .b8 retval0[4];
+; CHECK: call.uni (retval0),
+; CHECK-NEXT: _Z20__spirv_GroupCMulKHRjjN5__spv12complex_halfE,