[TLI] Fix replace-with-veclib crash with invalid arguments

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/CMakeLists.txt

@@ -40,6 +40,7 @@ set(ANALYSIS_TEST_SOURCES
   PluginInlineAdvisorAnalysisTest.cpp
   PluginInlineOrderAnalysisTest.cpp
   ProfileSummaryInfoTest.cpp
+  ReplaceWithVecLibTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
   SparsePropagation.cpp

llvm/unittests/Analysis/ReplaceWithVecLibTest.cpp

@@ -0,0 +1,85 @@
+//===--- ReplaceWithVecLibTest.cpp - replace-with-veclib unit tests -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/ReplaceWithVeclib.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("ReplaceWithVecLibTest", errs());
+  return Mod;
+}
+
+/// Runs ReplaceWithVecLib with different TLIIs that have custom VecDescs. This
+/// allows checking that the pass won't crash when the function to replace (from
+/// the input IR) does not match the replacement function (derived from the
+/// VecDesc mapping).
+class ReplaceWithVecLibTest : public ::testing::Test {
+protected:
+  LLVMContext Ctx;
+
+  /// Creates TLII using the given \p VD, and then runs the ReplaceWithVeclib
+  /// pass. The pass should not crash even when the replacement function
+  /// (derived from the \p VD mapping) does not match the function to be
+  /// replaced (from the input \p IR).
+  bool run(const VecDesc &VD, const char *IR) {
+    // Create TLII and register it with FAM so it's preserved when
+    // ReplaceWithVeclib pass runs.
+    TargetLibraryInfoImpl TLII = TargetLibraryInfoImpl(Triple());
+    TLII.addVectorizableFunctions({VD});
+    FunctionAnalysisManager FAM;
+    FAM.registerPass([&TLII]() { return TargetLibraryAnalysis(TLII); });
+
+    // Register and run the pass on the 'foo' function from the input IR.
+    FunctionPassManager FPM;
+    FPM.addPass(ReplaceWithVeclib());
+    std::unique_ptr<Module> M = parseIR(Ctx, IR);
+    PassBuilder PB;
+    PB.registerFunctionAnalyses(FAM);
+    FPM.run(*M->getFunction("foo"), FAM);
+
+    return true;
+  }
+};
+
+} // end anonymous namespace
+
+static const char *IR = R"IR(
+define <vscale x 4 x float> @foo(<vscale x 4 x float> %in){
+  %call = call <vscale x 4 x float> @llvm.powi.f32.i32(<vscale x 4 x float> %in, i32 3)
+  ret <vscale x 4 x float> %call
+}
+
+declare <vscale x 4 x float> @llvm.powi.f32.i32(<vscale x 4 x float>, i32) #0
+)IR";
+
+// LLVM intrinsic 'powi' (in IR) has the same signature with the VecDesc.
+TEST_F(ReplaceWithVecLibTest, TestValidMapping) {
+  VecDesc CorrectVD = {"llvm.powi.f32.i32", "_ZGVsMxvu_powi",
+                       ElementCount::getScalable(4), true, "_ZGVsMxvu"};
+  EXPECT_TRUE(run(CorrectVD, IR));
+}
+
+// LLVM intrinsic 'powi' (in IR) has different signature with the VecDesc.
+TEST_F(ReplaceWithVecLibTest, TestInvalidMapping) {
+  VecDesc IncorrectVD = {"llvm.powi.f32.i32", "_ZGVsMxvv_powi",
+                         ElementCount::getScalable(4), true, "_ZGVsMxvv"};
+  EXPECT_TRUE(run(IncorrectVD, IR));
+}