[VFABI] Read/Write functions for the VFABI attribute.

[fpetrogalli]

Summary

The attribute is stored at the FunctionIndex attribute set, with the name "vector-function-abi-variant".

The get/set methods of the attribute have assertion to verify that:

1. Each name in the attribute is a valid VFABI mangled name.
2. Each name in the attribute correspond to a function declared in the module.

[jdoerfert]

There are some unrelated changes once could commit seperatly as NFC and I left some generic code design comments to further simplify this. Though this already looks pretty good to me.

llvm/lib/Analysis/VectorUtils.cpp

1184

To get the attribute:

`CI->getAttribute(AttributeList::FunctionIndex, "vector-function-abi-variant")`

To get the module:

`CI->getModule()`

Is the container choice (set) on purpose? Do we need to worry about non-determinism (I usually do!)? Do we use it as a set or would a vector do as well?

llvm/lib/Transforms/Utils/ModuleUtils.cpp

317

I would have made it something like this, though it is mostly style:

if (VariantMappings.empty())
  return;
for (const std::string &VariantMapping : VariantMappings)
  Out << VariantMapping << ",";
Buffer.pop_back();
... Buffer.str();

In the NDEBUG:
for (const std::string &VariantMapping : VariantMappings)

We need a variable to hold this string "vector-...".

Attribute addition:
CI->addAttribute(AttributeList::FunctionIndex, Attribute::get(C, MISSING_VAR_FOR_STR, Buffer.out())

llvm/unittests/Analysis/VectorFunctionABITest.cpp

482

A static LLVMContext seems risky if the test run in parallel (idk if they do but they could maybe).
Other unit tests have a setup and teardown method and own their module instead of using the constructor.
And now I see there is also a Ctx member, I'm confused.

[sdesmalen]

Thanks for this patch @fpetrogalli, I appreciate you splitting this off from the other patch!
The patch overall looks good to me when you implement some of @jdoerfert suggestions.

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

124

Can these use StringRef?

[fpetrogalli]

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

124

Indeed they can. Thank you.

llvm/lib/Analysis/VectorUtils.cpp

1184

Is the container choice (set) on purpose? Do we need to worry about non-determinism (I usually do!)? Do we use it as a set or would a vector do as well?

Could do with a vector as well. It's just that I didn't want to take care of the accidental possibility of seeing duplicates in the IR. Not a functional problem itself, just something without any particular meaning.

llvm/lib/Transforms/Utils/ModuleUtils.cpp

317

We need a variable to hold this string "vector-...".

Sure. Where? If I add it in one header file, it complains of not being used in other places where the header file is not used. I have added it in the compilation unit, but in there it replaces only a single use...

llvm/unittests/Analysis/VectorFunctionABITest.cpp

482

Bad coding, sorry. I fixed both tests.

[fpetrogalli]

@jdoerfert , @sdesmalen , thank you for the reviews!

[fpetrogalli]

[spam] Testing arcanist.

[fpetrogalli]

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

124

It turned out that the assertion on invalid vector ABI names fire when running the code on the TLI tests in the loop vectorizer. It looks like the underlying memory of the StringRefs is modified. I have added some debugging and it shows that the mappings values are modified somewhere in the pipeline. I'll revert back to std::string for now.

FWIW, StringRef is not supposed to store data (see doxygen docs), but just to operate on buffers that owned by other component. This means that StringRef cannot be used as a return parameter of a function, because the buffer set up in the function is destroyed when leaving the function. We are not using StringRef as a return value here, but I guess that we incur in a similar problem by using StringRef in a reference parameter that is in fact an output value.

[jdoerfert]

I have only minor comments left (below), @sdesmalen you wanna take a look and accept?

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

124

FWIW, It will not work here because the "strings" you put in there do not outlive the function on their own. One can return strings in a vector of StringRefs in different situations.

llvm/lib/Analysis/VectorUtils.cpp

1184

If you need deduplication -> SetVector :)

llvm/lib/Transforms/Utils/ModuleUtils.cpp

286

Now this is better but having it defined multiple times is bad.
Single location (static constexpr char*) should work or single location (extern const char *) with a single definition somewhere.

317

See above :)

[fpetrogalli]

Minor changes for last review from @jdoerfert .

[sdesmalen]

llvm/include/llvm/Analysis/VectorUtils.h

138

It would be better passed as an ArrayRef<std::string>, or otherwise a SmallVectorImpl<std::string> &VariantMappings.

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

124

Great, thanks for clarifying!

I would however recommend using ArrayRef or SmallVectorImpl, instead of SmallVector<std::string, 8>.

llvm/lib/Analysis/VectorUtils.cpp

1172

The code would be simpler to read if we let the compiler hoist M out of the loop below, so please move it to the second NDEBUG block.

llvm/lib/Transforms/Utils/ModuleUtils.cpp

304

nit: unnecessary curly braces.

327

nit: is this assert really necessary?

353

are these asserts necessary?

365

nit: the comment doesn't really clarify more than the code already expresses, so can be removed.

376

if (!TLIName.empty())

378

remove?

391

is this assert necessary?

[fpetrogalli]

llvm/lib/Analysis/VectorUtils.cpp

1172

I invoked getModule directly in the assertion, as it is the only place where it is used.

llvm/lib/Transforms/Utils/ModuleUtils.cpp

304

This function shouldn't be here, I have removed it.

327

This method shouldn't be here. I have removed it.

353

This function shouldn't be here. I have removed it.

365

This function shouldn't be here. I have removed it.

376

This function shouldn't be here. I have removed it.

378

This function shouldn't be here. I have removed it.

[fpetrogalli]

Address review from @sdesmalen .

Notice that I removed a couple of stray methods that where related to the handling of the TLI changes.

[sdesmalen]

LGTM!

llvm/include/llvm/Analysis/VectorUtils.h

 //===- llvm/Analysis/VectorUtils.h - Vector utilities -----------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines some vectorizer utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_VECTORUTILS_H
 #define LLVM_ANALYSIS_VECTORUTILS_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/Support/CheckedArithmetic.h"
 
 namespace llvm {
 
 /// Describes the type of Parameters
 enum class VFParamKind {

 ///
 /// This methods support demangling rules for the following <isa>:
 ///
 /// * AArch64: https://developer.arm.com/docs/101129/latest
 ///
 /// * x86 (libmvec): https://sourceware.org/glibc/wiki/libmvec and
 ///  https://sourceware.org/glibc/wiki/libmvec?action=AttachFile&do=view&target=VectorABI.txt
 ///
-///
-///
 /// \param MangledName -> input string in the format
 /// _ZGV<isa><mask><vlen><parameters>_<scalarname>[(<redirection>)].
 Optional<VFInfo> tryDemangleForVFABI(StringRef MangledName);
 
 /// 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
+/// associated to the CallInst CI.
+void getVectorVariantNames(const CallInst &CI,
+                           SmallVector<std::string, 8> &VariantMappings);
 } // end namespace VFABI
 
 template <typename T> class ArrayRef;
 class DemandedBits;
 class GetElementPtrInst;
 template <typename InstTy> class InterleaveGroup;
 class Loop;
 class ScalarEvolution;
-class TargetLibraryInfo;
 class TargetTransformInfo;
 class Type;
 class Value;
 
 namespace Intrinsic {
 enum ID : unsigned;
 }
 

llvm/include/llvm/Transforms/Utils/ModuleUtils.h

 //===-- ModuleUtils.h - Functions to manipulate Modules ---------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This family of functions perform manipulations on Modules.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_MODULEUTILS_H
 #define LLVM_TRANSFORMS_UTILS_MODULEUTILS_H
 
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringRef.h"
 #include <utility> // for std::pair
 
 namespace llvm {
 
 template <typename T> class ArrayRef;
 class Module;
 class Function;

 /// This identifier is normally guaranteed to be unique, or the program would
 /// fail to link due to multiply defined symbols.
 ///
 /// If the module has no strong external symbols (such a module may still have a
 /// semantic effect if it performs global initialization), we cannot produce a
 /// unique identifier for this module, so we return the empty string.
 std::string getUniqueModuleId(Module *M);
 
+class TargetLibraryInfo;
+class CallInst;
+namespace VFABI {
+
+/// \defgroup Vector Function ABI (VABI) Module functions.
+///
+/// Utility functions for VFABI data that can modify the module.
+///
+/// @{
+/// Overwrite the Vector Function ABI variants attribute with the names provide
+/// in \p VariantMappings.
+void setVectorVariantNames(CallInst *CI,
+                           const SmallVector<std::string, 8> &VariantMappings);
+
+/// @}
+} // End VFABI namespace
+
 } // End llvm namespace
 
 #endif //  LLVM_TRANSFORMS_UTILS_MODULEUTILS_H

llvm/lib/Analysis/VFABIDemangling.cpp

 //===- VFABIDemangling.cpp - Vector Function ABI demangling utilities. ---===//
 //
 // 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/ADT/SmallSet.h"
+#include "llvm/ADT/SmallString.h"
 #include "llvm/Analysis/VectorUtils.h"
 
 using namespace llvm;
 
 namespace {
 /// Utilities for the Vector Function ABI name parser.
 
 /// Return types for the parser functions.
ParseRet tryParseISA(StringRef &MangledName, VFISAKind &ISA) {
   ISA = StringSwitch<VFISAKind>(MangledName.take_front(1))
             .Case("n", VFISAKind::AdvancedSIMD)
             .Case("s", VFISAKind::SVE)
             .Case("b", VFISAKind::SSE)
             .Case("c", VFISAKind::AVX)
             .Case("d", VFISAKind::AVX2)
             .Case("e", VFISAKind::AVX512)
             .Default(VFISAKind::Unknown);
-
   MangledName = MangledName.drop_front(1);
 
   return ParseRet::OK;
 }
 
 /// Extracts the `<mask>` information from the mangled string, and
 /// sets `IsMasked` accordingly. The input string `MangledName` is
 /// left unmodified.
if (ParamFound == ParseRet::OK) {
       if (AlignFound == ParseRet::Error)
         return None;
 
       // Add the parameter.
       Parameters.push_back({ParameterPos, PKind, StepOrPos, Alignment});
     }
   } while (ParamFound == ParseRet::OK);
 
-  // A valid MangledName mus have at least one valid entry in the
+  // A valid MangledName must have at least one valid entry in the
   // <parameters>.
   if (Parameters.empty())
     return None;
 
   // Check for the <scalarname> and the optional <redirection>, which
   // are separated from the prefix with "_"
   if (!MangledName.consume_front("_"))
     return None;

llvm/lib/Analysis/VectorUtils.cpp

 template <>
 void InterleaveGroup<Instruction>::addMetadata(Instruction *NewInst) const {
   SmallVector<Value *, 4> VL;
   std::transform(Members.begin(), Members.end(), std::back_inserter(VL),
                  [](std::pair<int, Instruction *> p) { return p.second; });
   propagateMetadata(NewInst, VL);
 }
 }
+
+void VFABI::getVectorVariantNames(
+    const CallInst &CI, SmallVector<std::string, 8> &VariantMappings) {
+  const StringRef S =
+      CI.getAttribute(AttributeList::FunctionIndex, VFABI::MappingsAttrName)
+          .getValueAsString();
+  SmallVector<StringRef, 8> ListAttr;
+  S.split(ListAttr, ",");
+
+#ifndef NDEBUG
+  const Module *M = CI.getModule();
+#endif
+  for (auto &S : SetVector<StringRef>(ListAttr.begin(), ListAttr.end())) {
+#ifndef NDEBUG
+    Optional<VFInfo> Info = VFABI::tryDemangleForVFABI(S);
+    assert(Info.hasValue() && "Invalid name for a VFABI variant.");
+    assert(M->getFunction(Info.getValue().VectorName) &&
+           "Vector function is missing.");
+#endif
+    VariantMappings.push_back(S);
+  }
+}

llvm/lib/Transforms/Utils/ModuleUtils.cpp

 //===-- ModuleUtils.cpp - Functions to manipulate Modules -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This family of functions perform manipulations on Modules.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 
std::string llvm::getUniqueModuleId(Module *M) {
 
   MD5::MD5Result R;
   Md5.final(R);
 
   SmallString<32> Str;
   MD5::stringifyResult(R, Str);
   return ("$" + Str).str();
 }
+
+/// Helper function to map the TLI name to a strings that holds
+/// scalar-to-vector mapping.
+///
+///    _ZGV<isa><mask><vlen><vparams>_<scalarname>(<vectorname>)
+///
+/// where:
+///
+/// <isa> = "_LLVM_TLI_"
+/// <mask> = "N". Note: TLI does not support masked interfaces.
+/// <vlen> = Number of concurrent lanes, stored in the `VectorizationFactor`
+///          field of the `VecDesc` struct.
+/// <vparams> = "v", as many as are the number of parameters of CI.
+/// <scalarname> = the name of the scalar function called by CI.
+/// <vectorname> = the name of the vector function mapped by the TLI.
+static std::string mangleTLIName(StringRef VectorName, CallInst *CI,
+                                 unsigned VF) {
+  SmallString<256> Buffer;
+  llvm::raw_svector_ostream Out(Buffer);
+  Out << "_ZGV" << _LLVM_TLI_ << "N" << VF;
+  for (unsigned I = 0; I < CI->getNumArgOperands(); ++I) {
+    Out << "v";
+  }
+  Out << "_" << CI->getCalledFunction()->getName() << "(" << VectorName << ")";
+  return Out.str();
+}
+
+/// A helper function for converting Scalar types to vector types.
+/// If the incoming type is void, we return void. If the VF is 1, we return
+/// the scalar type.
+static Type *ToVectorTy(Type *Scalar, unsigned VF) {
+  if (Scalar->isVoidTy() || VF == 1)
+    return Scalar;
+  return VectorType::get(Scalar, VF);
+}
+
+/// A helper function that adds the vector function declaration that
+/// vectorizes the CallInst CI with a vectorization factor of VF
+/// lanes. The TLI assumes that all parameters and the return type of
+/// CI (other than void) need to be widened to a VectorType of VF
+/// lanes.
+static void addVariantDeclaration(CallInst *CI, const unsigned VF,
+                                  const StringRef VFName) {
+  assert(CI && "Invalid CallInst.");
+  Module *M = CI->getParent()->getParent()->getParent();
+  llvm::GlobalValue *Global = M->getNamedValue(VFName);
+  // Nothing to do if the function already exists in the module.
+  if (Global)
+    return;
+
+  Type *RetTy = ToVectorTy(CI->getType(), VF);
+  SmallVector<Type *, 4> Tys;
+  for (Value *ArgOperand : CI->arg_operands())
+    Tys.push_back(ToVectorTy(ArgOperand->getType(), VF));
+  FunctionType *FTy = FunctionType::get(RetTy, Tys, /*isVarArg=*/false);
+  Function *VectorF =
+      Function::Create(FTy, Function::ExternalLinkage, VFName, M);
+  VectorF->copyAttributesFrom(CI->getCalledFunction());
+
+  // Make function declaration (without a body) "sticky" in the IR by
+  // listing them in the @llvm.compiler.used attribute
+  if (VectorF->size() == 0) {
+    Global = M->getNamedValue(VFName);
+    assert(Global && "Missing function declaration.");
+    appendToCompilerUsed(*M, {Global});
+  }
+}
+
+void VFABI::addMappingsFromTLI(const TargetLibraryInfo *TLI, CallInst *CI) {
+  assert(TLI && "Invalid TLI.");
+  assert(CI && "Invalid CallInst.");
+
+  // This is needed to make sure we don't query the TLI for calls to
+  // bitcast of function pointers, like `%call = call i32 (i32*, ...)
+  // bitcast (i32 (...)* @goo to i32 (i32*, ...)*)(i32* nonnull %i)`,
+  // as such calls make the `isFunctionVectorizable` raise an
+  // exception.
+  if (CI->isNoBuiltin() || !CI->getCalledFunction())
+    return;
+
+  const std::string ScalarName = CI->getCalledFunction()->getName();
+  // Nothing to be done if the TLI things the function is not
+  // vectorizable.
+  if (!TLI->isFunctionVectorizable(ScalarName))
+    return;
+
+  SmallSet<std::string, 8> SetOfMangledNames;
+  VFABI::getVectorVariantNames(CI, SetOfMangledNames);
+  Module *M = CI->getParent()->getParent()->getParent();
+
+  for (unsigned VF = 2; VF <= 16; VF *= 2) {
+    const std::string TLIName = TLI->getVectorizedFunction(ScalarName, VF);
+    if (TLIName != "") {
+      std::string MangledName = mangleTLIName(TLIName, CI, VF);
+      //      List.push_back(MangledName);
+      SetOfMangledNames.insert(MangledName);
+      Function *VariantF = M->getFunction(TLIName);
+      if (!VariantF)
+        addVariantDeclaration(CI, VF, TLIName);
+    }
+  }
+
+  VFABI::setVectorVariantNames(CI, SetOfMangledNames);
+}
+
+void VFABI::setVectorVariantNames(
+    CallInst *CI, const SmallVector<std::string, 8> &VariantMappings) {
+  assert(CI && "Invalid CallInst");
+
+  if (VariantMappings.empty())
+    return;
+
+  SmallString<256> Buffer;
+  llvm::raw_svector_ostream Out(Buffer);
+  for (const std::string &VariantMapping : VariantMappings)
+    Out << VariantMapping << ",";
+  // Get rid of the trailing ','.
+  assert(!Buffer.str().empty() && "Must have at least one char.");
+  Buffer.pop_back();
+
+  Module *M = CI->getModule();
+#ifndef NDEBUG
+  for (const std::string &VariantMapping : VariantMappings) {
+    Optional<VFInfo> VI = VFABI::tryDemangleForVFABI(VariantMapping);
+    assert(VI.hasValue() && "Canno add an invalid VFABI name.");
+    assert(M->getNamedValue(VI.getValue().VectorName) &&
+           "Cannot add variant to attribute: "
+           "vector function declaration is missing.");
+  }
+#endif
+  CI->addAttribute(
+      AttributeList::FunctionIndex,
+      Attribute::get(M->getContext(), MappingsAttrName, Buffer.str()));
+}

llvm/unittests/Analysis/VectorFunctionABITest.cpp

 //===------- VectorFunctionABITest.cpp - VFABI Unittests  ---------===//
 //
 // 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/Analysis/VectorUtils.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/InstIterator.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 // This test makes sure that the getFromVFABI method succeeds only on
 // valid values of the string.
 TEST(VectorFunctionABITests, OnlyValidNames) {
   // Incomplete string.
TEST_F(VFABIParserTest, ParseMaskingAVX512) {
   EXPECT_TRUE(IsMasked());
   EXPECT_FALSE(IsScalable);
   EXPECT_EQ(ISA, VFISAKind::AVX512);
   EXPECT_EQ(Parameters.size(), (unsigned)2);
   EXPECT_EQ(Parameters[0], VFParameter({0, VFParamKind::Vector}));
   EXPECT_EQ(Parameters[1], VFParameter({1, VFParamKind::GlobalPredicate}));
   EXPECT_EQ(ScalarName, "sin");
 }
+
+class VFABIAttrTest : public testing::Test {
+protected:
+  void SetUp() override {
+    M = parseAssemblyString(IR, Err, Ctx);
+    // Get the only call instruction in the block, which is the first
+    // instruction.
+    CI = dyn_cast<CallInst>(&*(instructions(M->getFunction("f")).begin()));
+  }
+  const char *IR = "define i32 @f(i32 %a) {\n"
+                   " %1 = call i32 @g(i32 %a) #0\n"
+                   "  ret i32 %1\n"
+                   "}\n"
+                   "declare i32 @g(i32)\n"
+                   "declare <2 x i32> @custom_vg(<2 x i32>)"
+                   "declare <4 x i32> @_ZGVnN4v_g(<4 x i32>)"
+                   "declare <8 x i32> @_ZGVnN8v_g(<8 x i32>)"
+                   "attributes #0 = { "
+                   "\"vector-function-abi-variant\"=\""
+                   "_ZGVnN2v_g(custom_vg),_ZGVnN4v_g\" }";
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M;
+  CallInst *CI;
+  SmallVector<std::string, 8> Mappings;
+};
+
+TEST_F(VFABIAttrTest, Read) {
+  VFABI::getVectorVariantNames(*CI, Mappings);
+  SmallVector<std::string, 8> Exp;
+  Exp.push_back("_ZGVnN2v_g(custom_vg)");
+  Exp.push_back("_ZGVnN4v_g");
+  EXPECT_EQ(Mappings, Exp);
+}

llvm/unittests/Transforms/Utils/CMakeLists.txt

add_llvm_unittest(UtilsTests
   CodeExtractorTest.cpp
   FunctionComparatorTest.cpp
   IntegerDivisionTest.cpp
   LocalTest.cpp
   SizeOptsTest.cpp
   SSAUpdaterBulkTest.cpp
   UnrollLoopTest.cpp
   ValueMapperTest.cpp
+  VFABIUtils.cpp
   )

llvm/unittests/Transforms/Utils/VFABIUtils.cpp

+//===------- VFABIUtils.cpp - VFABI Unittests  ----------------------------===//
+//
+// 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/AsmParser/Parser.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+class VFABIAttrTest : public testing::Test {
+protected:
+  void SetUp() override {
+    M = parseAssemblyString(IR, Err, Ctx);
+    // Get the only call instruction in the block, which is the first
+    // instruction.
+    CI = dyn_cast<CallInst>(&*(instructions(M->getFunction("f")).begin()));
+  }
+  const char *IR = "define i32 @f(i32 %a) {\n"
+                   " %1 = call i32 @g(i32 %a) #0\n"
+                   "  ret i32 %1\n"
+                   "}\n"
+                   "declare i32 @g(i32)\n"
+                   "declare <2 x i32> @custom_vg(<2 x i32>)"
+                   "declare <4 x i32> @_ZGVnN4v_g(<4 x i32>)"
+                   "declare <8 x i32> @_ZGVnN8v_g(<8 x i32>)"
+                   "attributes #0 = { "
+                   "\"vector-function-abi-variant\"=\""
+                   "_ZGVnN2v_g(custom_vg),_ZGVnN4v_g\" }";
+  LLVMContext Ctx;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M;
+  CallInst *CI;
+  SmallVector<std::string, 8> Mappings;
+};
+
+TEST_F(VFABIAttrTest, Write) {
+  Mappings.push_back("_ZGVnN8v_g");
+  Mappings.push_back("_ZGVnN2v_g(custom_vg)");
+  VFABI::setVectorVariantNames(CI, Mappings);
+  const AttributeList Attrs = CI->getAttributes();
+  const AttributeSet FnAttrs = Attrs.getFnAttributes();
+  const StringRef S = CI->getAttribute(AttributeList::FunctionIndex,
+                                       "vector-function-abi-variant")
+                          .getValueAsString();
+  EXPECT_EQ(S, "_ZGVnN8v_g,_ZGVnN2v_g(custom_vg)");
+}