[IR] Add support for metadata attachments for Function Arguments

llvm/docs/LangRef.rst

@@ -319,7 +319,7 @@ added in the future:
     not be used lightly but only for specific situations such as an
     alternative to the *register pinning* performance technique often
     used when implementing functional programming languages. At the
-    moment only X86, AArch64, and RISCV support this convention. The 
+    moment only X86, AArch64, and RISCV support this convention. The
     following limitations exist:
 
     -  On *X86-32* only up to 4 bit type parameters are supported. No
@@ -643,10 +643,10 @@ implementation defined, the optimizer can't do the latter.  The former is
 challenging as many commonly expected properties, such as
 ``ptrtoint(v)-ptrtoint(v) == 0``, don't hold for non-integral types.
 Similar restrictions apply to intrinsics that might examine the pointer bits,
-such as :ref:`llvm.ptrmask<int_ptrmask>`. 
+such as :ref:`llvm.ptrmask<int_ptrmask>`.
 
 The alignment information provided by the frontend for a non-integral pointer
-(typically using attributes or metadata) must be valid for every possible 
+(typically using attributes or metadata) must be valid for every possible
 representation of the pointer.
 
 .. _globalvars:
@@ -824,7 +824,8 @@ an optional :ref:`calling convention <callingconv>`,
 an optional ``unnamed_addr`` attribute, a return type, an optional
 :ref:`parameter attribute <paramattrs>` for the return type, a function
 name, a (possibly empty) argument list (each with optional :ref:`parameter
-attributes <paramattrs>`), optional :ref:`function attributes <fnattrs>`,
+attributes <paramattrs>` and an optional list of attached :ref:`metadata <metadata>`),
+optional :ref:`function attributes <fnattrs>`,
 an optional address space, an optional section, an optional partition,
 an optional alignment, an optional :ref:`comdat <langref_comdats>`,
 an optional :ref:`garbage collector name <gc>`, an optional :ref:`prefix <prefixdata>`,
@@ -848,7 +849,7 @@ argument is of the following form:
 
 Syntax::
 
-   <type> [parameter Attrs] [name]
+   <type> [parameter Attrs] (!name !N)* [name]
 
 LLVM function declarations consist of the "``declare``" keyword, an
 optional :ref:`linkage type <linkage>`, an optional :ref:`visibility style

llvm/include/llvm/AsmParser/LLParser.h

@@ -285,6 +285,8 @@ namespace llvm {
     };
     bool parseEnumAttribute(Attribute::AttrKind Attr, AttrBuilder &B,
                             bool InAttrGroup);
+    bool parseOptionalParamMetadata(
+        SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs);
     bool parseOptionalParamOrReturnAttrs(AttrBuilder &B, bool IsParam);
     bool parseOptionalParamAttrs(AttrBuilder &B) {
       return parseOptionalParamOrReturnAttrs(B, true);
@@ -607,8 +609,10 @@ namespace llvm {
       Type *Ty;
       AttributeSet Attrs;
       std::string Name;
-      ArgInfo(LocTy L, Type *ty, AttributeSet Attr, const std::string &N)
-          : Loc(L), Ty(ty), Attrs(Attr), Name(N) {}
+      SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
+      ArgInfo(LocTy L, Type *ty, AttributeSet Attr, const std::string &N,
+              const SmallVector<std::pair<unsigned, MDNode *>, 8> &mds)
+          : Loc(L), Ty(ty), Attrs(Attr), Name(N), MDs(mds) {}
     };
     bool parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
                            SmallVectorImpl<unsigned> &UnnamedArgNums,

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -369,6 +369,7 @@ enum MetadataCodes {
   METADATA_GENERIC_SUBRANGE = 45, // [distinct, count, lo, up, stride]
   METADATA_ARG_LIST = 46,         // [n x [type num, value num]]
   METADATA_ASSIGN_ID = 47,        // [distinct, ...]
+  METADATA_PARAM_ATTACHMENT = 48, // [ [valueid ,] [valueid, [n x [id, mdnode]]]
 };
 
 // The constants block (CONSTANTS_BLOCK_ID) describes emission for each

llvm/include/llvm/IR/Argument.h

@@ -177,6 +177,14 @@ class Argument final : public Value {
   static bool classof(const Value *V) {
     return V->getValueID() == ArgumentVal;
   }
+
+  using Value::addMetadata;
+  using Value::clearMetadata;
+  using Value::eraseMetadata;
+  using Value::getAllMetadata;
+  using Value::getMetadata;
+  using Value::hasMetadata;
+  using Value::setMetadata;
 };
 
 } // End llvm namespace

llvm/lib/AsmParser/LLParser.cpp

@@ -1929,6 +1929,19 @@ bool LLParser::parseOptionalParamOrReturnAttrs(AttrBuilder &B, bool IsParam) {
   }
 }
 
+/// Parse a potentially empty list of parameter or return attributes.
+bool LLParser::parseOptionalParamMetadata(
+    SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) {
+  while (Lex.getKind() == lltok::MetadataVar) {
+    unsigned MDK;
+    MDNode *N;
+    if (parseMetadataAttachment(MDK, N))
+      return true;
+    MDs.emplace_back(MDK, N);
+  }
+  return false;
+}
+
 static unsigned parseOptionalLinkageAux(lltok::Kind Kind, bool &HasLinkage) {
   HasLinkage = true;
   switch (Kind) {
@@ -3058,7 +3071,9 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
       LocTy TypeLoc = Lex.getLoc();
       Type *ArgTy = nullptr;
       AttrBuilder Attrs(M->getContext());
-      if (parseType(ArgTy) || parseOptionalParamAttrs(Attrs))
+      SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
+      if (parseType(ArgTy) || parseOptionalParamAttrs(Attrs) ||
+          parseOptionalParamMetadata(MDs))
         return true;
 
       if (ArgTy->isVoidTy())
@@ -3087,7 +3102,7 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
 
       ArgList.emplace_back(TypeLoc, ArgTy,
                            AttributeSet::get(ArgTy->getContext(), Attrs),
-                           std::move(Name));
+                           std::move(Name), std::move(MDs));
     } while (EatIfPresent(lltok::comma));
   }
 
@@ -3115,6 +3130,8 @@ bool LLParser::parseFunctionType(Type *&Result) {
     if (ArgList[i].Attrs.hasAttributes())
       return error(ArgList[i].Loc,
                    "argument attributes invalid in function type");
+    if (!ArgList[i].MDs.empty())
+      return error(ArgList[i].Loc, "metadata invalid in function type");
   }
 
   SmallVector<Type*, 16> ArgListTy;
@@ -6223,6 +6240,8 @@ bool LLParser::parseFunctionHeader(Function *&Fn, bool IsDefine,
   // Add all of the arguments we parsed to the function.
   Function::arg_iterator ArgIt = Fn->arg_begin();
   for (unsigned i = 0, e = ArgList.size(); i != e; ++i, ++ArgIt) {
+    for (const auto &MD : ArgList[i].MDs)
+      ArgIt->addMetadata(MD.first, *MD.second);
     // If the argument has a name, insert it into the argument symbol table.
     if (ArgList[i].Name.empty()) continue;
 

llvm/lib/Bitcode/Reader/MetadataLoader.cpp

@@ -26,6 +26,7 @@
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Bitcode/LLVMBitCodes.h"
 #include "llvm/Bitstream/BitstreamReader.h"
+#include "llvm/IR/Argument.h"
 #include "llvm/IR/AutoUpgrade.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
@@ -60,9 +61,6 @@
 #include <type_traits>
 #include <utility>
 #include <vector>
-namespace llvm {
-class Argument;
-}
 
 using namespace llvm;
 
@@ -477,6 +475,7 @@ class MetadataLoader::MetadataLoaderImpl {
                              function_ref<void(StringRef)> CallBack);
   Error parseGlobalObjectAttachment(GlobalObject &GO,
                                     ArrayRef<uint64_t> Record);
+  Error parseParamAttachment(Function &F, ArrayRef<uint64_t> Record);
   Error parseMetadataKindRecord(SmallVectorImpl<uint64_t> &Record);
 
   void resolveForwardRefsAndPlaceholders(PlaceholderQueue &Placeholders);
@@ -2263,6 +2262,19 @@ Error MetadataLoader::MetadataLoaderImpl::parseOneMetadata(
     NextMetadataNo++;
     break;
   }
+  case bitc::METADATA_PARAM_ATTACHMENT: {
+    unsigned RecordLength = Record.size();
+    if (Record.empty() || (RecordLength - 2) % 2 == 1)
+      return error("Invalid record");
+    unsigned ValueID = Record[0];
+    if (ValueID >= ValueList.size())
+      return error("Invalid record");
+    if (auto *F = dyn_cast<Function>(ValueList[ValueID]))
+      if (Error Err =
+              parseParamAttachment(*F, ArrayRef<uint64_t>(Record).slice(1)))
+        return Err;
+    break;
+  }
   }
   return Error::success();
 #undef GET_OR_DISTINCT
@@ -2321,6 +2333,24 @@ Error MetadataLoader::MetadataLoaderImpl::parseGlobalObjectAttachment(
   return Error::success();
 }
 
+Error MetadataLoader::MetadataLoaderImpl::parseParamAttachment(
+    Function &F, ArrayRef<uint64_t> Record) {
+  assert((Record.size() - 1) % 2 == 0);
+
+  auto *A = F.getArg(Record[0]);
+  for (unsigned I = 1, E = Record.size(); I != E; I += 2) {
+    auto K = MDKindMap.find(Record[I]);
+    if (K == MDKindMap.end())
+      return error("Invalid ID");
+    MDNode *MD =
+        dyn_cast_or_null<MDNode>(getMetadataFwdRefOrLoad(Record[I + 1]));
+    if (!MD)
+      return error("Invalid metadata attachment: expect fwd ref to MDNode");
+    A->addMetadata(K->second, *MD);
+  }
+  return Error::success();
+}
+
 /// Parse metadata attachments.
 Error MetadataLoader::MetadataLoaderImpl::parseMetadataAttachment(
     Function &F, ArrayRef<Instruction *> InstructionList) {
@@ -2406,6 +2436,14 @@ Error MetadataLoader::MetadataLoaderImpl::parseMetadataAttachment(
       }
       break;
     }
+    case bitc::METADATA_PARAM_ATTACHMENT: {
+      unsigned RecordLength = Record.size();
+      if (Record.empty() || (RecordLength - 1) % 2 == 1)
+        return error("Invalid record");
+      if (Error Err = parseParamAttachment(F, Record))
+        return Err;
+      break;
+    }
     }
   }
 }

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -29,6 +29,7 @@
 #include "llvm/Bitstream/BitCodes.h"
 #include "llvm/Bitstream/BitstreamWriter.h"
 #include "llvm/Config/llvm-config.h"
+#include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Comdat.h"
@@ -378,6 +379,8 @@ class ModuleBitcodeWriter : public ModuleBitcodeWriterBase {
   void writeFunctionMetadataAttachment(const Function &F);
   void pushGlobalMetadataAttachment(SmallVectorImpl<uint64_t> &Record,
                                     const GlobalObject &GO);
+  void pushParamMetadataAttachment(SmallVectorImpl<uint64_t> &Record,
+                                   const Argument &A);
   void writeModuleMetadataKinds();
   void writeOperandBundleTags();
   void writeSyncScopeNames();
@@ -2393,8 +2396,17 @@ void ModuleBitcodeWriter::writeModuleMetadata() {
     Stream.EmitRecord(bitc::METADATA_GLOBAL_DECL_ATTACHMENT, Record);
   };
   for (const Function &F : M)
-    if (F.isDeclaration() && F.hasMetadata())
-      AddDeclAttachedMetadata(F);
+    if (F.isDeclaration()) {
+      if (F.hasMetadata())
+        AddDeclAttachedMetadata(F);
+      for (const auto &A : F.args())
+        if (A.hasMetadata()) {
+          Record.push_back(VE.getValueID(&F));
+          pushParamMetadataAttachment(Record, A);
+          Stream.EmitRecord(bitc::METADATA_PARAM_ATTACHMENT, Record);
+          Record.clear();
+        }
+    }
   // FIXME: Only store metadata for declarations here, and move data for global
   // variable definitions to a separate block (PR28134).
   for (const GlobalVariable &GV : M.globals())
@@ -2426,11 +2438,30 @@ void ModuleBitcodeWriter::pushGlobalMetadataAttachment(
   }
 }
 
+void ModuleBitcodeWriter::pushParamMetadataAttachment(
+    SmallVectorImpl<uint64_t> &Record, const Argument &A) {
+  Record.push_back(A.getArgNo());
+  // [n x [id, mdnode]]
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  A.getAllMetadata(MDs);
+  for (const auto &I : MDs) {
+    Record.push_back(I.first);
+    Record.push_back(VE.getMetadataID(I.second));
+  }
+}
+
 void ModuleBitcodeWriter::writeFunctionMetadataAttachment(const Function &F) {
   Stream.EnterSubblock(bitc::METADATA_ATTACHMENT_ID, 3);
 
   SmallVector<uint64_t, 64> Record;
 
+  for (const auto &A : F.args())
+    if (A.hasMetadata()) {
+      pushParamMetadataAttachment(Record, A);
+      Stream.EmitRecord(bitc::METADATA_PARAM_ATTACHMENT, Record);
+      Record.clear();
+    }
+
   if (F.hasMetadata()) {
     pushGlobalMetadataAttachment(Record, F);
     Stream.EmitRecord(bitc::METADATA_ATTACHMENT, Record, 0);
@@ -3476,6 +3507,9 @@ void ModuleBitcodeWriter::writeFunction(
 
   bool NeedsMetadataAttachment = F.hasMetadata();
 
+  for (auto IA = F.arg_begin(), E = F.arg_end(); IA != E; ++IA)
+    NeedsMetadataAttachment |= IA->hasMetadata();
+
   DILocation *LastDL = nullptr;
   SmallSetVector<Function *, 4> BlockAddressUsers;
 

llvm/lib/Bitcode/Writer/ValueEnumerator.cpp

@@ -398,8 +398,13 @@ ValueEnumerator::ValueEnumerator(const Module &M,
 
   // Enumerate types used by function bodies and argument lists.
   for (const Function &F : M) {
-    for (const Argument &A : F.args())
+    for (const Argument &A : F.args()) {
       EnumerateType(A.getType());
+      MDs.clear();
+      A.getAllMetadata(MDs);
+      for (const auto &I : MDs)
+        EnumerateMetadata(F.isDeclaration() ? nullptr : &F, I.second);
+    }
 
     // Enumerate metadata attached to this function.
     MDs.clear();

llvm/lib/IR/AsmWriter.cpp

@@ -857,6 +857,9 @@ class SlotTracker : public AbstractSlotTrackerStorage {
   /// Add all of the metadata from a function.
   void processFunctionMetadata(const Function &F);
 
+  /// Add all of the metadata from a argument.
+  void processParamMetadata(const Argument &A);
+
   /// Add all of the metadata from an instruction.
   void processInstructionMetadata(const Instruction &I);
 
@@ -1128,6 +1131,9 @@ void SlotTracker::processGlobalObjectMetadata(const GlobalObject &GO) {
 
 void SlotTracker::processFunctionMetadata(const Function &F) {
   processGlobalObjectMetadata(F);
+  for (const auto &A : F.args()) {
+    processParamMetadata(A);
+  }
   for (auto &BB : F) {
     for (auto &I : BB) {
       for (const DPValue &DPV : I.getDbgValueRange())
@@ -1162,6 +1168,14 @@ void SlotTracker::processInstructionMetadata(const Instruction &I) {
     CreateMetadataSlot(MD.second);
 }
 
+void SlotTracker::processParamMetadata(const Argument &A) {
+  // Process metadata attached to this instruction.
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  A.getAllMetadata(MDs);
+  for (auto &MD : MDs)
+    CreateMetadataSlot(MD.second);
+}
+
 /// Clean up after incorporating a function. This is the only way to get out of
 /// the function incorporation state that affects get*Slot/Create*Slot. Function
 /// incorporation state is indicated by TheFunction != 0.
@@ -3918,6 +3932,12 @@ void AssemblyWriter::printFunction(const Function *F) {
         Out << ' ';
         writeAttributeSet(ArgAttrs);
       }
+      auto *Arg = F->getArg(I);
+      if (Arg->hasMetadata()) {
+        SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+        Arg->getAllMetadata(MDs);
+        printMetadataAttachments(MDs, " ");
+      }
     }
   } else {
     // The arguments are meaningful here, print them in detail.
@@ -4010,6 +4030,10 @@ void AssemblyWriter::printArgument(const Argument *Arg, AttributeSet Attrs) {
     writeAttributeSet(Attrs);
   }
 
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  Arg->getAllMetadata(MDs);
+  printMetadataAttachments(MDs, " ");
+
   // Output name, if available...
   if (Arg->hasName()) {
     Out << ' ';