[OpenMP][FIX] Use name + type checks not only name checks for calls

[jdoerfert]

Summary

A call that is analyzed in an optimization needs to be verified against
the name and type of the runtime function to avoid that we look at
arguments that do not exist (anymore). This can happen if the signature
was rewritten.

[tianshilei1992]

I didn’t get the idea in the description from the code clearly. The old comparison is still address based. In what scenario that RFI->Declaration can be nullptr? The only difference might be when CB->getCalledFunction() and RFI.Declarion are both nullptr, which will fail the check now.

[jdoerfert]

In D106341#2889391, @tianshilei1992 wrote:

I didn’t get the idea in the description from the code clearly. The old comparison is still address based. In what scenario that RFI->Declaration can be nullptr? The only difference might be when CB->getCalledFunction() and RFI.Declarion are both nullptr, which will fail the check now.

Yes. We do not set RFI.Declaration if the type doesn't match what we expect. getCallIfRegularCall should have always checked RFI has a declaration, now it does and that makes sure we see a call to the right function with the right type when we access the operand.

[tianshilei1992]

LGTM. Probably good to put the explanation in description as well.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct OpenMPOpt {
   }
 
   /// Return the call if \p U is a callee use in a regular call. If \p RFI is
   /// given it has to be the callee or a nullptr is returned.
   static CallInst *getCallIfRegularCall(
       Use &U, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr) {
     CallInst *CI = dyn_cast<CallInst>(U.getUser());
     if (CI && CI->isCallee(&U) && !CI->hasOperandBundles() &&
-        (!RFI || CI->getCalledFunction() == RFI->Declaration))
+        (!RFI ||
+         (RFI->Declaration && CI->getCalledFunction() == RFI->Declaration)))
       return CI;
     return nullptr;
   }
 
   /// Return the call if \p V is a regular call. If \p RFI is given it has to be
   /// the callee or a nullptr is returned.
   static CallInst *getCallIfRegularCall(
       Value &V, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr) {
     CallInst *CI = dyn_cast<CallInst>(&V);
     if (CI && !CI->hasOperandBundles() &&
-        (!RFI || CI->getCalledFunction() == RFI->Declaration))
+        (!RFI ||
+         (RFI->Declaration && CI->getCalledFunction() == RFI->Declaration)))
       return CI;
     return nullptr;
   }
 
 private:
   /// Merge parallel regions when it is safe.
   bool mergeParallelRegions() {
     const unsigned CallbackCalleeOperand = 2;
auto IsInitialThreadOnly = [&](BranchInst *Edge, BasicBlock *SuccessorBB) {
 
     ConstantInt *C = dyn_cast<ConstantInt>(Cmp->getOperand(1));
     if (!C)
       return false;
 
     // Match:  -1 == __kmpc_target_init (for non-SPMD kernels only!)
     if (C->isAllOnesValue()) {
       auto *CB = dyn_cast<CallBase>(Cmp->getOperand(0));
-      if (!CB || CB->getCalledFunction() != RFI.Declaration)
+      CB = CB ? OpenMPOpt::getCallIfRegularCall(*CB, &RFI) : nullptr;
+      if (!CB)
         return false;
       const int InitIsSPMDArgNo = 1;
       auto *IsSPMDModeCI =
           dyn_cast<ConstantInt>(CB->getOperand(InitIsSPMDArgNo));
       return IsSPMDModeCI && IsSPMDModeCI->isZero();
     }
 
     return false;