[IRLinker] Fix mapping of declaration metadata

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

Reviewed By: tejohnson

Differential Revision: https://reviews.llvm.org/D145318

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

@@ -112,8 +112,9 @@ inline RemapFlags operator|(RemapFlags LHS, RemapFlags RHS) {
 /// 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
@@ -175,6 +176,7 @@ public:
 
   void remapInstruction(Instruction &I);
   void remapFunction(Function &F);
+  void remapGlobalObjectMetadata(GlobalObject &GO);
 
   void scheduleMapGlobalInitializer(GlobalVariable &GV, Constant &Init,
                                     unsigned MappingContextID = 0);

llvm/lib/Linker/IRMover.cpp

@@ -409,6 +409,10 @@ class IRLinker {
   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);
@@ -750,8 +754,11 @@ GlobalValue *IRLinker::copyGlobalValueProto(const GlobalValue *SGV,
 
   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
@@ -1056,6 +1063,10 @@ Expected<Constant *> IRLinker::linkGlobalValueProto(GlobalValue *SGV,
   // 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;
@@ -1651,6 +1662,13 @@ Error IRLinker::run() {
   // 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(),

llvm/lib/Transforms/Utils/ValueMapper.cpp

@@ -1181,6 +1181,10 @@ 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) {

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

@@ -10,4 +10,13 @@ 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

@@ -21,6 +21,15 @@ 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"}