[Analysis] Change VFABI::mangleTLIVectorName to use ElementCount

[david-arm]

Summary

Adds support for mangling TLI vector names for scalable vectors.

[sdesmalen]

LGTM!

llvm/include/llvm/Analysis/VectorUtils.h

 /// This returned string holds scalar-to-vector mapping:
 ///    _ZGV<isa><mask><vlen><vparams>_<scalarname>(<vectorname>)
 ///
 /// where:
 ///
 /// <isa> = "_LLVM_"
 /// <mask> = "N". Note: TLI does not support masked interfaces.
 /// <vlen> = Number of concurrent lanes, stored in the `VectorizationFactor`
-///          field of the `VecDesc` struct.
+///          field of the `VecDesc` struct. If the number of lanes is scalable
+///          then 'x' is printed instead.
 /// <vparams> = "v", as many as are the numArgs.
 /// <scalarname> = the name of the scalar function.
 /// <vectorname> = the name of the vector function.
 std::string mangleTLIVectorName(StringRef VectorName, StringRef ScalarName,
-                                unsigned numArgs, unsigned VF);
+                                unsigned numArgs, ElementCount VF);
 
 /// Retrieve the `VFParamKind` from a string token.
 VFParamKind getVFParamKindFromString(const StringRef Token);
 
 // Name of the attribute where the variant mappings are stored.
 static constexpr char const *MappingsAttrName = "vector-function-abi-variant";
 
 /// Populates a set of strings representing the Vector Function ABI variants

llvm/lib/Analysis/VectorUtils.cpp

void InterleaveGroup<Instruction>::addMetadata(Instruction *NewInst) const {
   std::transform(Members.begin(), Members.end(), std::back_inserter(VL),
                  [](std::pair<int, Instruction *> p) { return p.second; });
   propagateMetadata(NewInst, VL);
 }
 }
 
 std::string VFABI::mangleTLIVectorName(StringRef VectorName,
                                        StringRef ScalarName, unsigned numArgs,
-                                       unsigned VF) {
+                                       ElementCount VF) {
   SmallString<256> Buffer;
   llvm::raw_svector_ostream Out(Buffer);
-  Out << "_ZGV" << VFABI::_LLVM_ << "N" << VF;
+  Out << "_ZGV" << VFABI::_LLVM_ << "N";
+  if (VF.isScalable())
+    Out << 'x';
+  else
+    Out << VF.getFixedValue();
   for (unsigned I = 0; I < numArgs; ++I)
     Out << "v";
   Out << "_" << ScalarName << "(" << VectorName << ")";
   return std::string(Out.str());
 }
 
 void VFABI::getVectorVariantNames(
     const CallInst &CI, SmallVectorImpl<std::string> &VariantMappings) {

llvm/lib/Transforms/Utils/InjectTLIMappings.cpp

const SetVector<StringRef> OriginalSetOfMappings(Mappings.begin(),
                                                    Mappings.end());
   //  All VFs in the TLI are powers of 2.
   for (unsigned VF = 2, WidestVF = TLI.getWidestVF(ScalarName); VF <= WidestVF;
        VF *= 2) {
     const std::string TLIName =
         std::string(TLI.getVectorizedFunction(ScalarName, VF));
     if (!TLIName.empty()) {
       std::string MangledName = VFABI::mangleTLIVectorName(
-          TLIName, ScalarName, CI.getNumArgOperands(), VF);
+          TLIName, ScalarName, CI.getNumArgOperands(),
+          ElementCount::getFixed(VF));
       if (!OriginalSetOfMappings.count(MangledName)) {
         Mappings.push_back(MangledName);
         ++NumCallInjected;
       }
       Function *VariantF = M->getFunction(TLIName);
       if (!VariantF)
         addVariantDeclaration(CI, VF, TLIName);
     }

llvm/unittests/Analysis/VectorFunctionABITest.cpp

bool IsMasked() const {
     return NGlobalPreds == 1 && Info.Shape.Parameters.back().ParamKind ==
                                     VFParamKind::GlobalPredicate;
   }
 };
 } // unnamed namespace
 
 // This test makes sure correct mangling occurs for given string.
 TEST_F(VFABIParserTest, ManglingVectorTLINames) {
-  EXPECT_EQ(VFABI::mangleTLIVectorName("vec", "scalar", 3, 4),
-            "_ZGV_LLVM_N4vvv_scalar(vec)");
-  EXPECT_EQ(VFABI::mangleTLIVectorName("custom.call.v5", "custom.call", 1, 5),
+  EXPECT_EQ(
+      VFABI::mangleTLIVectorName("vec", "scalar", 3, ElementCount::getFixed(4)),
+  EXPECT_EQ(VFABI::mangleTLIVectorName("vec", "scalar", 3,
+                                       ElementCount::getScalable(4)),
+            "_ZGV_LLVM_Nxvvv_scalar(vec)");
+  EXPECT_EQ(VFABI::mangleTLIVectorName("custom.call.v5", "custom.call", 1,
+                                       ElementCount::getFixed(5)),
             "_ZGV_LLVM_N5v_custom.call(custom.call.v5)");
 }
 
 // This test makes sure that the demangling method succeeds only on
 // valid values of the string.
 TEST_F(VFABIParserTest, OnlyValidNames) {
   // Incomplete string.
   EXPECT_FALSE(invokeParser(""));