1 files changed, 68 insertions, 21 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPULibFunc.cpp b/llvm/lib/Target/AMDGPU/AMDGPULibFunc.cpp
index 4c596e37476c..264b4d43248c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULibFunc.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULibFunc.cpp
@@ -620,17 +620,17 @@ bool ItaniumParamParser::parseItaniumParam(StringRef& param,
   // parse type
   char const TC = param.front();
   if (isDigit(TC)) {
-    res.ArgType = StringSwitch<AMDGPULibFunc::EType>
-      (eatLengthPrefixedName(param))
-      .Case("ocl_image1darray" , AMDGPULibFunc::IMG1DA)
-      .Case("ocl_image1dbuffer", AMDGPULibFunc::IMG1DB)
-      .Case("ocl_image2darray" , AMDGPULibFunc::IMG2DA)
-      .Case("ocl_image1d"      , AMDGPULibFunc::IMG1D)
-      .Case("ocl_image2d"      , AMDGPULibFunc::IMG2D)
-      .Case("ocl_image3d"      , AMDGPULibFunc::IMG3D)
-      .Case("ocl_event"        , AMDGPULibFunc::DUMMY)
-      .Case("ocl_sampler"      , AMDGPULibFunc::DUMMY)
-      .Default(AMDGPULibFunc::DUMMY);
+    res.ArgType =
+        StringSwitch<AMDGPULibFunc::EType>(eatLengthPrefixedName(param))
+            .StartsWith("ocl_image1d_array", AMDGPULibFunc::IMG1DA)
+            .StartsWith("ocl_image1d_buffer", AMDGPULibFunc::IMG1DB)
+            .StartsWith("ocl_image2d_array", AMDGPULibFunc::IMG2DA)
+            .StartsWith("ocl_image1d", AMDGPULibFunc::IMG1D)
+            .StartsWith("ocl_image2d", AMDGPULibFunc::IMG2D)
+            .StartsWith("ocl_image3d", AMDGPULibFunc::IMG3D)
+            .Case("ocl_event", AMDGPULibFunc::DUMMY)
+            .Case("ocl_sampler", AMDGPULibFunc::DUMMY)
+            .Default(AMDGPULibFunc::DUMMY);
   } else {
     drop_front(param);
     switch (TC) {
@@ -957,8 +957,15 @@ static Type* getIntrinsicParamType(
   case AMDGPULibFunc::EVENT:
     T = PointerType::getUnqual(C);
     break;
-  default:
-    llvm_unreachable("Unhandled param type");
+  case AMDGPULibFunc::B8:
+  case AMDGPULibFunc::B16:
+  case AMDGPULibFunc::B32:
+  case AMDGPULibFunc::B64:
+  case AMDGPULibFunc::SIZE_MASK:
+  case AMDGPULibFunc::FLOAT:
+  case AMDGPULibFunc::INT:
+  case AMDGPULibFunc::UINT:
+  case AMDGPULibFunc::DUMMY:
     return nullptr;
   }
   if (P.VectorSize > 1)
@@ -969,13 +976,18 @@ static Type* getIntrinsicParamType(
   return T;
 }
 
-FunctionType *AMDGPUMangledLibFunc::getFunctionType(Module &M) const {
+FunctionType *AMDGPUMangledLibFunc::getFunctionType(const Module &M) const {
   LLVMContext& C = M.getContext();
   std::vector<Type*> Args;
   ParamIterator I(Leads, manglingRules[FuncId]);
   Param P;
-  while ((P=I.getNextParam()).ArgType != 0)
-    Args.push_back(getIntrinsicParamType(C, P, true));
+  while ((P = I.getNextParam()).ArgType != 0) {
+    Type *ParamTy = getIntrinsicParamType(C, P, true);
+    if (!ParamTy)
+      return nullptr;
+
+    Args.push_back(ParamTy);
+  }
 
   return FunctionType::get(
     getIntrinsicParamType(C, getRetType(FuncId, Leads), true),
@@ -997,9 +1009,44 @@ std::string AMDGPUMangledLibFunc::getName() const {
   return std::string(OS.str());
 }
 
-bool AMDGPULibFunc::isCompatibleSignature(const FunctionType *FuncTy) const {
-  // TODO: Validate types make sense
-  return !FuncTy->isVarArg() && FuncTy->getNumParams() == getNumArgs();
+bool AMDGPULibFunc::isCompatibleSignature(const Module &M,
+                                          const FunctionType *CallTy) const {
+  const FunctionType *FuncTy = getFunctionType(M);
+
+  if (!FuncTy) {
+    // Give up on mangled functions with unexpected types.
+    if (AMDGPULibFuncBase::isMangled(getId()))
+      return false;
+
+    // FIXME: UnmangledFuncInfo does not have any type information other than
+    // the number of arguments.
+    return getNumArgs() == CallTy->getNumParams();
+  }
+
+  // Normally the types should exactly match.
+  if (FuncTy == CallTy)
+    return true;
+
+  const unsigned NumParams = FuncTy->getNumParams();
+  if (NumParams != CallTy->getNumParams())
+    return false;
+
+  for (unsigned I = 0; I != NumParams; ++I) {
+    Type *FuncArgTy = FuncTy->getParamType(I);
+    Type *CallArgTy = CallTy->getParamType(I);
+    if (FuncArgTy == CallArgTy)
+      continue;
+
+    // Some cases permit implicit splatting a scalar value to a vector argument.
+    auto *FuncVecTy = dyn_cast<VectorType>(FuncArgTy);
+    if (FuncVecTy && FuncVecTy->getElementType() == CallArgTy &&
+        allowsImplicitVectorSplat(I))
+      continue;
+
+    return false;
+  }
+
+  return true;
 }
 
 Function *AMDGPULibFunc::getFunction(Module *M, const AMDGPULibFunc &fInfo) {
@@ -1012,7 +1059,7 @@ Function *AMDGPULibFunc::getFunction(Module *M, const AMDGPULibFunc &fInfo) {
   if (F->hasFnAttribute(Attribute::NoBuiltin))
     return nullptr;
 
-  if (!fInfo.isCompatibleSignature(F->getFunctionType()))
+  if (!fInfo.isCompatibleSignature(*M, F->getFunctionType()))
     return nullptr;
 
   return F;
@@ -1028,7 +1075,7 @@ FunctionCallee AMDGPULibFunc::getOrInsertFunction(Module *M,
     if (F->hasFnAttribute(Attribute::NoBuiltin))
       return nullptr;
     if (!F->isDeclaration() &&
-        fInfo.isCompatibleSignature(F->getFunctionType()))
+        fInfo.isCompatibleSignature(*M, F->getFunctionType()))
       return F;
   }