[llvm] [InjectTLIMappings] Remove SExt/ZExt attributes from vector functions. (PR #80546)

[Yeting Kuo via llvm-commits]

https://github.com/yetingk created https://github.com/llvm/llvm-project/pull/80546

Previously, when declaring vector functions, its attributes are copied by its scalar function. It may be illegal, since signext/zeroext could not be used for vector types.

>From dd62dc6adfb413fd0b8a1dfe9e383cbb3a016d02 Mon Sep 17 00:00:00 2001
From: Yeting Kuo <46629943+yetingk at users.noreply.github.com>
Date: Sat, 3 Feb 2024 21:51:08 +0800
Subject: [PATCH] [InjectTLIMappings] Remove SExt/ZExt attributes from vector
 functions.

Previously, when declaring vector functions, its attributes are copied by its
scalar function. It may be illegal, since signext/zeroext could not be used for
vector types.
---
 llvm/lib/Transforms/Utils/InjectTLIMappings.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
index 9bfac2ac9167e..27bfee4dfc7a9 100644
--- a/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
+++ b/llvm/lib/Transforms/Utils/InjectTLIMappings.cpp
@@ -34,6 +34,19 @@ STATISTIC(NumVFDeclAdded,
 STATISTIC(NumCompUsedAdded,
           "Number of `@llvm.compiler.used` operands that have been added.");
 
+static void removeIllegalAttributes(Function *Func, const FunctionType *FTy) {
+  // Vector types could not have attributes signext/zeroext.
+  if (FTy->getReturnType()->isVectorTy()) {
+    Func->removeRetAttr(Attribute::SExt);
+    Func->removeRetAttr(Attribute::ZExt);
+  }
+  for (const auto I : enumerate(FTy->params()))
+    if (I.value()->isVectorTy()) {
+      Func->removeParamAttr(I.index(), Attribute::SExt);
+      Func->removeParamAttr(I.index(), Attribute::ZExt);
+    }
+}
+
 /// A helper function that adds the vector variant declaration for vectorizing
 /// the CallInst \p CI with a vectorization factor of \p VF lanes. For each
 /// mapping, TLI provides a VABI prefix, which contains all information required
@@ -56,6 +69,8 @@ static void addVariantDeclaration(CallInst &CI, const ElementCount &VF,
   Function *VecFunc =
       Function::Create(VectorFTy, Function::ExternalLinkage, VFName, M);
   VecFunc->copyAttributesFrom(CI.getCalledFunction());
+  removeIllegalAttributes(VecFunc, VectorFTy);
+
   ++NumVFDeclAdded;
   LLVM_DEBUG(dbgs() << DEBUG_TYPE << ": Added to the module: `" << VFName
                     << "` of type " << *VectorFTy << "\n");