[TLI] Fix replace-with-veclib crashes with invalid arguments.

replace-with-veclib used to crash when the arguments of the TLI mapping
did not match the arguments of the mapping. Now, it simply ignores such
cases.

Author: paschalis-mpeis

llvm/lib/CodeGen/ReplaceWithVeclib.cpp

@@ -108,15 +108,17 @@ static bool replaceWithCallToVeclib(const TargetLibraryInfo &TLI,
   // Compute the argument types of the corresponding scalar call and the scalar
   // function name. For calls, it additionally finds the function to replace
   // and checks that all vector operands match the previously found EC.
-  SmallVector<Type *, 8> ScalarArgTypes;
+  SmallVector<Type *, 8> ScalarArgTypes, OrigArgTypes;
   std::string ScalarName;
   Function *FuncToReplace = nullptr;
-  if (auto *CI = dyn_cast<CallInst>(&I)) {
+  auto *CI = dyn_cast<CallInst>(&I);
+  if (CI) {
     FuncToReplace = CI->getCalledFunction();
     Intrinsic::ID IID = FuncToReplace->getIntrinsicID();
     assert(IID != Intrinsic::not_intrinsic && "Not an intrinsic");
     for (auto Arg : enumerate(CI->args())) {
       auto *ArgTy = Arg.value()->getType();
+      OrigArgTypes.push_back(ArgTy);
       if (isVectorIntrinsicWithScalarOpAtArg(IID, Arg.index())) {
         ScalarArgTypes.push_back(ArgTy);
       } else if (auto *VectorArgTy = dyn_cast<VectorType>(ArgTy)) {
@@ -174,6 +176,24 @@ static bool replaceWithCallToVeclib(const TargetLibraryInfo &TLI,
 
   Function *TLIFunc = getTLIFunction(I.getModule(), VectorFTy,
                                      VD->getVectorFnName(), FuncToReplace);
+
+  // For calls, bail out when their arguments do not match with the TLI mapping.
+  if (CI) {
+    int IdxNonPred = 0;
+    for (auto [OrigTy, VFParam] :
+         zip(OrigArgTypes, OptInfo->Shape.Parameters)) {
+      if (VFParam.ParamKind == VFParamKind::GlobalPredicate)
+        continue;
+      ++IdxNonPred;
+      if (OrigTy->isVectorTy() != (VFParam.ParamKind == VFParamKind::Vector)) {
+        LLVM_DEBUG(dbgs() << DEBUG_TYPE
+                          << ": Will not replace: wrong type at index: "
+                          << IdxNonPred << ": " << *OrigTy << "\n");
+        return false;
+      }
+    }
+  }
+
   replaceWithTLIFunction(I, *OptInfo, TLIFunc);
   LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Replaced call to `" << ScalarName
                     << "` with call to `" << TLIFunc->getName() << "`.\n");

llvm/unittests/Analysis/ReplaceWithVecLibTest.cpp

@@ -81,6 +81,5 @@ TEST_F(ReplaceWithVecLibTest, TestValidMapping) {
 TEST_F(ReplaceWithVecLibTest, TestInvalidMapping) {
   VecDesc IncorrectVD = {"llvm.powi.f32.i32", "_ZGVsMxvv_powi",
                          ElementCount::getScalable(4), true, "_ZGVsMxvv"};
-  /// TODO: test should avoid and not crash.
-  EXPECT_DEATH(run(IncorrectVD, IR), "");
+  EXPECT_TRUE(run(IncorrectVD, IR));
 }