[IRLinker] Fix mapping of declaration metadata

[critson]

Summary

Ensure metadata for declarations copied during materialization
is properly mapped if declarations do not become definitions.

[tejohnson]

Is this a case where we didn't in the past have metadata attachments on declarations requiring mapping, and now we do? Also, the example uses a "!types" metadata - I'm familiar with !type metadata on definitions, but haven't seen !types before - is this new metadata? I can't find any uses of it in the test suite or any documentation.

llvm/lib/Linker/IRMover.cpp

1148

GlobalVariable are also GlobalObjects, should they get the same treatment in linkGlobalVariable?

But is this removal needed at all, since the loop over UnmappedMetadata later on ignores definitions?

[critson]

In D145318#4169797, @tejohnson wrote:

Is this a case where we didn't in the past have metadata attachments on declarations requiring mapping, and now we do? Also, the example uses a "!types" metadata - I'm familiar with !type metadata on definitions, but haven't seen !types before - is this new metadata? I can't find any uses of it in the test suite or any documentation.

!types is custom metadata we are investigating to handle DXIL bitcode.
An example of its generation can be seen in the tests for D127728.

llvm/lib/Linker/IRMover.cpp

1148

I think removing all the erase operations is probably sensible.

[critson]

• Remove unnecessary erase calls, simplifying implementation
• Extend comment

[tejohnson]

Thanks for the context. lgtm with a comment suggestion.

llvm/lib/Linker/IRMover.cpp

412

Add doxygen comment.

[critson]

• Add missing erasure of functions which are erased through renaming.

[tejohnson]

can you re-upload with context?

[critson]

• Add missing context

[critson]

In D145318#4181503, @tejohnson wrote:

can you re-upload with context?

Sorry, context should be fixed now.
The only change over the accepted revision is IRMover.cpp:1066.

[tejohnson]

llvm/lib/Linker/IRMover.cpp

1066

Looking at remangleIntrinsicFunction, it doesn't appear to copy any metadata from F. I'm not sure if this is because intrinsic function declarations can never have metadata, or it just hasn't been needed thus far. I think it would be safer to add NewGV to the UnmappedMetadata set (if F was in it) in case the latter is true, meaning there is no guarantee we couldn't have any that needs to be mapped in the future.

Also, it would be good to expand your test case to test this case (i.e. I assume we end up with a crash if an intrinsic that was renamed had ended up in the UnmappedMetadata set without this change?).

Actually, in thinking through the case that presumably had an issue without this change (which I'm assuming was an intrinsic that was added to the UnmappedMetadata that did not in fact have any metadata, since I don't see remangleIntrinsicFunction copying any metadata), I'm wondering if we can reduce the amount of unnecessary stuff initially added to UnmappedMetadata, but guarding it with a check of NewGO->hasMetadata() ?

[tejohnson]

llvm/lib/Linker/IRMover.cpp

1066

Actually, in thinking through the case that presumably had an issue without this change (which I'm assuming was an intrinsic that was added to the UnmappedMetadata that did not in fact have any metadata, since I don't see remangleIntrinsicFunction copying any metadata), I'm wondering if we can reduce the amount of unnecessary stuff initially added to UnmappedMetadata, but guarding it with a check of NewGO->hasMetadata() ?

And if you do that, it should avoid the need to remove F from the UnmappedMetadata set here (under the assumption that we don't currently have intrinsic functions with any metadata at this point). So you can simply replace this line with an assert that F is not in UnmappedMetadata, with a comment noting that remangleIntrinsicFunction does not copy any metadata over, so we assume that F could not have metadata to start with. If the assert is ever hit, then not only does this code need to change but remangleIntrinsicFunction needs to be updated as well.

[critson]

• Exclude global objects without metadata from unmapped set
• Add assertion to cover the case of intrinsics with metadata

[critson]

llvm/lib/Linker/IRMover.cpp

1066

This seems like a good solution to me.
I have not extended the test cases as existing linker tests cover this with the assertion added.
(Assertion catches what was only being found by MSAN tests.)

[tejohnson]

lgtm

llvm/lib/Linker/IRMover.cpp

1066

Great, an existing test hitting this is good.

llvm/include/llvm/Transforms/Utils/ValueMapper.h

 ///
 /// A shared context used for mapping and remapping of Value and Metadata
 /// instances using \a ValueToValueMapTy, \a RemapFlags, \a
 /// ValueMapTypeRemapper, and \a ValueMaterializer.
 ///
 /// There are a number of top-level entry points:
 /// - \a mapValue() (and \a mapConstant());
 /// - \a mapMetadata() (and \a mapMDNode());
-/// - \a remapInstruction(); and
-/// - \a remapFunction().
+/// - \a remapInstruction();
+/// - \a remapFunction(); and
+/// - \a remapGlobalObjectMetadata().
 ///
 /// The \a ValueMaterializer can be used as a callback, but cannot invoke any
 /// of these top-level functions recursively.  Instead, callbacks should use
 /// one of the following to schedule work lazily in the \a ValueMapper
 /// instance:
 /// - \a scheduleMapGlobalInitializer()
 /// - \a scheduleMapAppendingVariable()
 /// - \a scheduleMapGlobalAlias()
public:
   Metadata *mapMetadata(const Metadata &MD);
   MDNode *mapMDNode(const MDNode &N);
 
   Value *mapValue(const Value &V);
   Constant *mapConstant(const Constant &C);
 
   void remapInstruction(Instruction &I);
   void remapFunction(Function &F);
+  void remapGlobalObjectMetadata(GlobalObject &GO);
 
   void scheduleMapGlobalInitializer(GlobalVariable &GV, Constant &Init,
                                     unsigned MappingContextID = 0);
   void scheduleMapAppendingVariable(GlobalVariable &GV, Constant *InitPrefix,
                                     bool IsOldCtorDtor,
                                     ArrayRef<Constant *> NewMembers,
                                     unsigned MappingContextID = 0);
   void scheduleMapGlobalAlias(GlobalAlias &GA, Constant &Aliasee,

llvm/lib/Linker/IRMover.cpp

class IRLinker {
   /// but this allows us to reuse the ValueMapper code.
   ValueToValueMapTy ValueMap;
   ValueToValueMapTy IndirectSymbolValueMap;
 
   DenseSet<GlobalValue *> ValuesToLink;
   std::vector<GlobalValue *> Worklist;
   std::vector<std::pair<GlobalValue *, Value*>> RAUWWorklist;
 
+  /// Set of globals with eagerly copied metadata that may require remapping.
+  /// This remapping is performed after metadata linking.
+  DenseSet<GlobalObject *> UnmappedMetadata;
+
   void maybeAdd(GlobalValue *GV) {
     if (ValuesToLink.insert(GV).second)
       Worklist.push_back(GV);
   }
 
   /// Whether we are importing globals for ThinLTO, as opposed to linking the
   /// source module. If this flag is set, it means that we can rely on some
   /// other object file to define any non-GlobalValue entities defined by the
GlobalValue *IRLinker::copyGlobalValueProto(const GlobalValue *SGV,
 
   if (ForDefinition)
     NewGV->setLinkage(SGV->getLinkage());
   else if (SGV->hasExternalWeakLinkage())
     NewGV->setLinkage(GlobalValue::ExternalWeakLinkage);
 
   if (auto *NewGO = dyn_cast<GlobalObject>(NewGV)) {
     // Metadata for global variables and function declarations is copied eagerly.
-    if (isa<GlobalVariable>(SGV) || SGV->isDeclaration())
+    if (isa<GlobalVariable>(SGV) || SGV->isDeclaration()) {
       NewGO->copyMetadata(cast<GlobalObject>(SGV), 0);
+      if (SGV->isDeclaration() && NewGO->hasMetadata())
+        UnmappedMetadata.insert(NewGO);
+    }
   }
 
   // Remove these copied constants in case this stays a declaration, since
   // they point to the source module. If the def is linked the values will
   // be mapped in during linkFunctionBody.
   if (auto *NewF = dyn_cast<Function>(NewGV)) {
     NewF->setPersonalityFn(nullptr);
     NewF->setPrefixData(nullptr);
if (ShouldLink || !ForIndirectSymbol)
       NeedsRenaming = true;
   }
 
   // Overloaded intrinsics have overloaded types names as part of their
   // names. If we renamed overloaded types we should rename the intrinsic
   // as well.
   if (Function *F = dyn_cast<Function>(NewGV))
     if (auto Remangled = Intrinsic::remangleIntrinsicFunction(F)) {
+      // Note: remangleIntrinsicFunction does not copy metadata and as such
+      // F should not occur in the set of objects with unmapped metadata.
+      // If this assertion fails then remangleIntrinsicFunction needs updating.
+      assert(!UnmappedMetadata.count(F) && "intrinsic has unmapped metadata");
       NewGV->eraseFromParent();
       NewGV = *Remangled;
       NeedsRenaming = false;
     }
 
   if (NeedsRenaming)
     forceRenaming(NewGV, SGV->getName());
 
Error IRLinker::linkFunctionBody(Function &Dst, Function &Src) {
   Dst.copyMetadata(&Src, 0);
 
   // Steal arguments and splice the body of Src into Dst.
   Dst.stealArgumentListFrom(Src);
   Dst.splice(Dst.end(), &Src);
 
   // Everything has been moved over.  Remap it.
   Mapper.scheduleRemapFunction(Dst);
   return Error::success();
 }
 
 void IRLinker::linkAliasAliasee(GlobalAlias &Dst, GlobalAlias &Src) {
   Mapper.scheduleMapGlobalAlias(Dst, *Src.getAliasee(), IndirectSymbolMCID);
 }
 
 void IRLinker::linkIFuncResolver(GlobalIFunc &Dst, GlobalIFunc &Src) {
   Mapper.scheduleMapGlobalIFunc(Dst, *Src.getResolver(), IndirectSymbolMCID);
Error IRLinker::run() {
   DoneLinkingBodies = true;
   Mapper.addFlags(RF_NullMapMissingGlobalValues);
 
   // Remap all of the named MDNodes in Src into the DstM module. We do this
   // after linking GlobalValues so that MDNodes that reference GlobalValues
   // are properly remapped.
   linkNamedMDNodes();
 
+  // Clean up any global objects with potentially unmapped metadata.
+  // Specifically declarations which did not become definitions.
+  for (GlobalObject *NGO : UnmappedMetadata) {
+    if (NGO->isDeclaration())
+      Mapper.remapGlobalObjectMetadata(*NGO);
+  }
+
   if (!IsPerformingImport && !SrcM->getModuleInlineAsm().empty()) {
     // Append the module inline asm string.
     DstM.appendModuleInlineAsm(adjustInlineAsm(SrcM->getModuleInlineAsm(),
                                                SrcTriple));
   } else if (IsPerformingImport) {
     // Import any symver directives for symbols in DstM.
     ModuleSymbolTable::CollectAsmSymvers(*SrcM,
                                          [&](StringRef Name, StringRef Alias) {

llvm/lib/Transforms/Utils/ValueMapper.cpp

 void ValueMapper::remapInstruction(Instruction &I) {
   FlushingMapper(pImpl)->remapInstruction(&I);
 }
 
 void ValueMapper::remapFunction(Function &F) {
   FlushingMapper(pImpl)->remapFunction(F);
 }
 
+void ValueMapper::remapGlobalObjectMetadata(GlobalObject &GO) {
+  FlushingMapper(pImpl)->remapGlobalObjectMetadata(GO);
+}
+
 void ValueMapper::scheduleMapGlobalInitializer(GlobalVariable &GV,
                                                Constant &Init,
                                                unsigned MCID) {
   getAsMapper(pImpl)->scheduleMapGlobalInitializer(GV, Init, MCID);
 }
 
 void ValueMapper::scheduleMapAppendingVariable(GlobalVariable &GV,
                                                Constant *InitPrefix,

llvm/test/Linker/Inputs/metadata-function.ll

 define weak void @foo() !weak !0 {
   unreachable
 }
 
 define void @baz() !baz !0 {
   unreachable
 }
 
 define void @b() !b !0 {
   unreachable
 }
 
+%AltHandle = type { i8* }
+declare !types !1 %AltHandle @init.AltHandle()
+
+define void @uses.AltHandle() {
+  %.res = call %AltHandle @init.AltHandle()
+  unreachable
+}
+
 !0 = !{!"b"}
+!1 = !{%AltHandle undef}

llvm/test/Linker/metadata-function.ll

 declare void @baz()
 
 ; CHECK-DAG: define void @a() !a ![[A]] {
 ; CHECK-DAG: define void @b() !b ![[B]] {
 define void @a() !a !0 {
   unreachable
 }
 
+; CHECK-DAG: define %[[HandleType:[A-Za-z]+]] @init.Handle() {
+; CHECK-DAG: declare !types ![[C:[0-9]+]] %[[HandleType]] @init.AltHandle()
+; CHECK-DAG: define void @uses.AltHandle() {
+%Handle = type { i8* }
+define %Handle @init.Handle() {
+  unreachable
+}
+
 ; CHECK-DAG: ![[A]] = !{!"a"}
 ; CHECK-DAG: ![[B]] = !{!"b"}
+; CHECK-DAG: ![[C]] = !{%[[HandleType]] undef}
 !0 = !{!"a"}