diff --git a/llvm/include/llvm/Analysis/VectorUtils.h b/llvm/include/llvm/Analysis/VectorUtils.h --- a/llvm/include/llvm/Analysis/VectorUtils.h +++ b/llvm/include/llvm/Analysis/VectorUtils.h @@ -161,7 +161,13 @@ /// /// \param MangledName -> input string in the format /// _ZGV_[()]. -Optional tryDemangleForVFABI(StringRef MangledName); +/// \param M -> Module used to retrive informations about the vector +/// function that are not possible to retrieve from the mangled +/// name. This parameter is needed only to retrive the Vectorization +/// Factor of scalable vector functions from their respective IR +/// declarations. +Optional tryDemangleForVFABI(StringRef MangledName, + const Module *M = nullptr); /// Retrieve the `VFParamKind` from a string token. VFParamKind getVFParamKindFromString(const StringRef Token); @@ -200,7 +206,8 @@ SmallVector ListOfStrings; VFABI::getVectorVariantNames(CI, ListOfStrings); for (const auto &MangledName : ListOfStrings) { - const Optional Shape = VFABI::tryDemangleForVFABI(MangledName); + const Optional Shape = + VFABI::tryDemangleForVFABI(MangledName, CI.getModule()); // A match is found via scalar and vector names, and also by // ensuring that the variant described in the attribute has a // corresponding definition or declaration of the vector diff --git a/llvm/lib/Analysis/VFABIDemangling.cpp b/llvm/lib/Analysis/VFABIDemangling.cpp --- a/llvm/lib/Analysis/VFABIDemangling.cpp +++ b/llvm/lib/Analysis/VFABIDemangling.cpp @@ -287,11 +287,57 @@ return ParseRet::None; } +#ifndef NDEBUG +// Verify the assumtion that all vectors in the signature of a vector +// function have the same number of elements. +bool verifyAllVectorsHaveSameWidth(FunctionType *Signature) { + using Container = SmallVector; + Container VecTys; + auto RetTy = dyn_cast(Signature->getReturnType()); + if (RetTy) + VecTys.push_back(RetTy); + for (auto Ty : Signature->params()) { + auto VTy = dyn_cast(Ty); + if (VTy) + VecTys.push_back(VTy); + } + + if (VecTys.size() <= 1) + return true; + + assert(VecTys.size() > 1 && "Invalid number of elements."); + const ElementCount EC = VecTys[0]->getElementCount(); + return std::all_of(VecTys.begin() + 1, VecTys.end(), [&EC](VectorType *VTy) { + return (EC == VTy->getElementCount()); + }); +} + +#endif // NDEBUG + +// Extract the VectorizationFactor from a given function signature, +// under the assumtion that all vectors have the same number of +// elements, i.e. same ElementCount.Min. +ElementCount getECFromSignature(FunctionType *Signature) { + assert(verifyAllVectorsHaveSameWidth(Signature) && + "Invalid vector signature."); + + auto RetTy = dyn_cast(Signature->getReturnType()); + if (RetTy) + return RetTy->getElementCount(); + for (auto Ty : Signature->params()) { + auto VTy = dyn_cast(Ty); + if (VTy) + return VTy->getElementCount(); + } + + return {1 /*Min*/, false /*Scalable*/}; +} } // namespace // Format of the ABI name: // _ZGV_[()] -Optional VFABI::tryDemangleForVFABI(StringRef MangledName) { +Optional VFABI::tryDemangleForVFABI(StringRef MangledName, + const Module *M) { const StringRef OriginalName = MangledName; // Assume there is no custom name , and therefore the // vector name consists of @@ -402,6 +448,22 @@ assert(Parameters.back().ParamKind == VFParamKind::GlobalPredicate && "The global predicate must be the last parameter"); + // Adjust the VF for scalable signatures. The EC.Min is not encoded + // in the name of the function, but it is encoded in the IR + // signature of the function. We need to extract this information + // because it is needed by the loop vectorizer, which reasons in + // terms of VectorizationFactor or ElementCount. In particular, we + // need to make sure that the VF field of the VFShape class is never + // set to 0. + if (IsScalable) { + assert(M && "Access to the vector function declaration is needed."); + const Function *F = M->getFunction(VectorName); + assert(F && "Vector function definition is missing."); + const ElementCount EC = getECFromSignature(F->getFunctionType()); + VF = EC.Min; + } + + assert(VF && "VF must be a positive value."); const VFShape Shape({VF, IsScalable, Parameters}); return VFInfo({Shape, ScalarName, VectorName, ISA}); } diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp --- a/llvm/lib/Analysis/VectorUtils.cpp +++ b/llvm/lib/Analysis/VectorUtils.cpp @@ -1175,7 +1175,7 @@ for (auto &S : SetVector(ListAttr.begin(), ListAttr.end())) { #ifndef NDEBUG LLVM_DEBUG(dbgs() << "VFABI: adding mapping '" << S << "'\n"); - Optional Info = VFABI::tryDemangleForVFABI(S); + Optional Info = VFABI::tryDemangleForVFABI(S, CI.getModule()); assert(Info.hasValue() && "Invalid name for a VFABI variant."); assert(CI.getModule()->getFunction(Info.getValue().VectorName) && "Vector function is missing."); diff --git a/llvm/lib/Transforms/Utils/ModuleUtils.cpp b/llvm/lib/Transforms/Utils/ModuleUtils.cpp --- a/llvm/lib/Transforms/Utils/ModuleUtils.cpp +++ b/llvm/lib/Transforms/Utils/ModuleUtils.cpp @@ -301,7 +301,7 @@ #ifndef NDEBUG for (const std::string &VariantMapping : VariantMappings) { LLVM_DEBUG(dbgs() << "VFABI: adding mapping '" << VariantMapping << "'\n"); - Optional VI = VFABI::tryDemangleForVFABI(VariantMapping); + Optional VI = VFABI::tryDemangleForVFABI(VariantMapping, M); assert(VI.hasValue() && "Cannot add an invalid VFABI name."); assert(M->getNamedValue(VI.getValue().VectorName) && "Cannot add variant to attribute: " diff --git a/llvm/tools/vfabi-demangle-fuzzer/vfabi-demangler-fuzzer.cpp b/llvm/tools/vfabi-demangle-fuzzer/vfabi-demangler-fuzzer.cpp --- a/llvm/tools/vfabi-demangle-fuzzer/vfabi-demangler-fuzzer.cpp +++ b/llvm/tools/vfabi-demangle-fuzzer/vfabi-demangler-fuzzer.cpp @@ -16,7 +16,7 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) { const StringRef MangledName((const char *)Data, Size); - const auto Info = VFABI::tryDemangleForVFABI(MangledName); + const auto Info = VFABI::tryDemangleForVFABI(MangledName, nullptr /*M*/); // Do not optimize away the return value. Inspired by // https://github.com/google/benchmark/blob/master/include/benchmark/benchmark.h#L307-L345 diff --git a/llvm/unittests/Analysis/VectorFunctionABITest.cpp b/llvm/unittests/Analysis/VectorFunctionABITest.cpp --- a/llvm/unittests/Analysis/VectorFunctionABITest.cpp +++ b/llvm/unittests/Analysis/VectorFunctionABITest.cpp @@ -86,6 +86,11 @@ // Reset the parser output references. void reset() { Info = VFInfo(); } + // Data needed to load the optional IR passed to invokeParser + LLVMContext Ctx; + SMDiagnostic Err; + std::unique_ptr M; + // CallInst *CI; protected: // Referencies to the parser output field. unsigned &VF = Info.Shape.VF; @@ -95,9 +100,11 @@ std::string &VectorName = Info.VectorName; bool &IsScalable = Info.Shape.IsScalable; // Invoke the parser. - bool invokeParser(const StringRef MangledName) { + bool invokeParser(const StringRef MangledName, const char *IR = nullptr) { reset(); - const auto OptInfo = VFABI::tryDemangleForVFABI(MangledName); + if (IR) + M = parseAssemblyString(IR, Err, Ctx); + const auto OptInfo = VFABI::tryDemangleForVFABI(MangledName, M.get()); if (OptInfo.hasValue()) { Info = OptInfo.getValue(); return true; @@ -168,13 +175,25 @@ } TEST_F(VFABIParserTest, ParseScalableSVE) { - EXPECT_TRUE(invokeParser("_ZGVsMxv_sin")); - EXPECT_EQ(VF, (unsigned)0); + const char *IR = "define i32 @f(i32 %a) {\n" + " %1 = call i32 @sin(i32 %a) #0\n" + " ret i32 %1\n" + "}\n" + "declare i32 @sin(i32)\n" + "declare @custom_vg(, " + ")\n" + "attributes #0 = { " + "\"vector-function-abi-variant\"=\"" + "_ZGVsMxv_sin(custom_vg)" + "\" }"; + + EXPECT_TRUE(invokeParser("_ZGVsMxv_sin(custom_vg)", IR)); + EXPECT_EQ(VF, (unsigned)2); EXPECT_TRUE(IsMasked()); EXPECT_TRUE(IsScalable); EXPECT_EQ(ISA, VFISAKind::SVE); EXPECT_EQ(ScalarName, "sin"); - EXPECT_EQ(VectorName, "_ZGVsMxv_sin"); + EXPECT_EQ(VectorName, "custom_vg"); } TEST_F(VFABIParserTest, ParseFixedWidthSVE) { @@ -467,8 +486,21 @@ } TEST_F(VFABIParserTest, ParseScalableMaskingLLVM) { - EXPECT_TRUE(invokeParser("_ZGV_LLVM_Mxv_sin(custom_vector_sin)")); + const char *IR = "define i32 @f(i32 %a) {\n" + " %1 = call i32 @sin(i32 %a) #0\n" + " ret i32 %1\n" + "}\n" + "declare i32 @sin(i32)\n" + "declare @custom_vector_sin(, )\n" + "attributes #0 = { " + "\"vector-function-abi-variant\"=\"" + "_ZGVsMxv_sin(custom_vector_sin)" + "\" }"; + EXPECT_TRUE(invokeParser("_ZGV_LLVM_Mxv_sin(custom_vector_sin)", IR)); EXPECT_TRUE(IsMasked()); + EXPECT_EQ(VF, (unsigned)2); + EXPECT_TRUE(IsScalable); EXPECT_EQ(ISA, VFISAKind::LLVM); EXPECT_EQ(Parameters.size(), (unsigned)2); @@ -479,7 +511,22 @@ } TEST_F(VFABIParserTest, ParseScalableMaskingLLVMSincos) { - EXPECT_TRUE(invokeParser("_ZGV_LLVM_Mxvl8l8_sincos(custom_vector_sincos)")); + const char *IR = + "define void @f(double %a, double * %sin, double * %cos) {\n" + " call void @sincos(double %a, double * %sin, double * %cos) #0\n" + " ret void \n" + "}\n" + "declare void @sincos(double, double *, double *)\n" + "declare void @custom_vector_sincos(, double *, " + "double *)\n" + "attributes #0 = { " + "\"vector-function-abi-variant\"=\"" + "_ZGV_LLVM_Mxvl8l8_sincos(custom_vector_sincos)" + "\" }"; + + EXPECT_TRUE( + invokeParser("_ZGV_LLVM_Mxvl8l8_sincos(custom_vector_sincos)", IR)); + EXPECT_EQ(VF, (unsigned)2); EXPECT_TRUE(IsMasked()); EXPECT_TRUE(IsScalable); EXPECT_EQ(ISA, VFISAKind::LLVM);