[DeadArgElim][ArgumentPromotion] Preserve MD_Associated metadata when re-creating function

[Aaron1011]

Summary

Currently, DeadArgElim and ArgumentPromotion fail to preserve MD_Associated metadata pointing at the old function when a replacement function is created.
Since the replacement function is intended to be 'the same' as the old one (except for the changed arguments), any globals with MD_Associated metadata pointing to the old function should have it pointed at the new function.

I've added two helper functions to ModuleUtils which perform the metadata updating. It should be very easy to insert the proper calls to these functions into any other passes found to require the same logic.

[Aaron1011]

I accidentally excluded my test cases from the first patch

[Aaron1011]

Ping

[vsk]

Why not use RAUW to fix up any metadata which refers to a function?

include/llvm/Transforms/IPO/ArgumentPromotion.h

 //===- ArgumentPromotion.h - Promote by-reference arguments -----*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_IPO_ARGUMENTPROMOTION_H
 #define LLVM_TRANSFORMS_IPO_ARGUMENTPROMOTION_H
 
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
 
 namespace llvm {
 
 /// Argument promotion pass.
 ///
 /// This pass walks the functions in each SCC and for each one tries to
 /// transform it and all of its callers to replace indirect arguments with
 /// direct (by-value) arguments.
 class ArgumentPromotionPass : public PassInfoMixin<ArgumentPromotionPass> {
   unsigned MaxElements;
+  // Tracks globals with metadata pointing at functions.
+  // If we end up re-creating a function, we point
+  // the metadata at the new function.
+  AssociatedGlobalsMap AssociatedGlobals;
 
 public:
   ArgumentPromotionPass(unsigned MaxElements = 3u) : MaxElements(MaxElements) {}
 
   PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_ARGUMENTPROMOTION_H

include/llvm/Transforms/IPO/DeadArgumentElimination.h

 
 #ifndef LLVM_TRANSFORMS_IPO_DEADARGUMENTELIMINATION_H
 #define LLVM_TRANSFORMS_IPO_DEADARGUMENTELIMINATION_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <map>
 #include <set>
 #include <string>
 #include <tuple>
 
 namespace llvm {
 
 class Module;
public:
   ///    argument to F.
   ///  - Uses[ret F] = arg F
   ///    This means that F returns one of its own arguments.
   ///  - Uses[arg F] = arg G
   ///    This means that G calls F and passes one of its own (G's) arguments
   ///    directly to F.
   UseMap Uses;
 
+  /// Maps Functions to GlobalObjects which reference them via MDNodes
+  /// (e.g. MD_associated metadata). When we modify a function, we need
+  /// to change any metadata which point at the old function
+  /// to point at the new function (with dead args/return values removed)
+  AssociatedGlobalsMap AssociatedGlobals;
+
   using LiveSet = std::set<RetOrArg>;
   using LiveFuncSet = std::set<const Function *>;
 
   /// This set contains all values that have been determined to be live.
   LiveSet LiveValues;
 
   /// This set contains all values that are cannot be changed in any way.
   LiveFuncSet LiveFunctions;

include/llvm/Transforms/Utils/ModuleUtils.h

 // This family of functions perform manipulations on Modules.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_UTILS_MODULEUTILS_H
 #define LLVM_TRANSFORMS_UTILS_MODULEUTILS_H
 
 #include "llvm/ADT/StringRef.h"
+#include <llvm/IR/GlobalObject.h>
+#include <llvm/IR/LLVMContext.h>
+#include <llvm/IR/Metadata.h>
+#include <llvm/Support/Debug.h>
+#include <map>
 #include <utility> // for std::pair
 
 namespace llvm {
 
 template <typename T> class ArrayRef;
 class Module;
 class Function;
 class GlobalValue;
 class GlobalVariable;
 class Constant;
 class StringRef;
 class Value;
 class Type;
 
+using AssociatedGlobalsMap = std::multimap<const Function *, GlobalObject *>;
+
 /// Append F to the list of global ctors of module M with the given Priority.
 /// This wraps the function in the appropriate structure and stores it along
 /// side other global constructors. For details see
 /// http://llvm.org/docs/LangRef.html#intg_global_ctors
 void appendToGlobalCtors(Module &M, Function *F, int Priority,
                          Constant *Data = nullptr);
 
 /// Same as appendToGlobalCtors(), but for global dtors.

 /// 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);
 
+/// Creates a map from Functions to GlobalObjects which reference them via
+/// MDNodes (e.g. MD_associated metadata). When we modify a function, we need to
+/// change any metadata which point at the old function to point at the new
+/// function (with dead args/return values removed)
+void collectAssociatedGlobals(Module *M,
+                              AssociatedGlobalsMap *AssociatedGlobals);
+
+/// Changes GlobalObject metadata that points at OldFn to point at NewFn, using
+/// the data collected from collectAssociatedGlobals.
+/// This should be used by transformations that 'modify' functions by deleting
+/// the original and replacing it with a modified version (e.g.
+/// DeadArgumentElimination, ArgumentPromotion). This should only be used when
+/// the new function should be conisdered to be 'the same' as the old one.
+void FixupMetadataReferences(AssociatedGlobalsMap *AssociatedGlobals,
+                             Function *OldFn, Function *NewFn);
+
 } // End llvm namespace
 
 #endif //  LLVM_TRANSFORMS_UTILS_MODULEUTILS_H

lib/Transforms/IPO/ArgumentPromotion.cpp

 }
 
 PreservedAnalyses ArgumentPromotionPass::run(LazyCallGraph::SCC &C,
                                              CGSCCAnalysisManager &AM,
                                              LazyCallGraph &CG,
                                              CGSCCUpdateResult &UR) {
   bool Changed = false, LocalChange;
 
+  collectAssociatedGlobals(C.begin()->getFunction().getParent(),
+                           &AssociatedGlobals);
+
   // Iterate until we stop promoting from this SCC.
   do {
     LocalChange = false;
 
     for (LazyCallGraph::Node &N : C) {
       Function &OldF = N.getFunction();
 
       FunctionAnalysisManager &FAM =
for (LazyCallGraph::Node &N : C) {
       LocalChange = true;
 
       // Directly substitute the functions in the call graph. Note that this
       // requires the old function to be completely dead and completely
       // replaced by the new function. It does no call graph updates, it merely
       // swaps out the particular function mapped to a particular node in the
       // graph.
       C.getOuterRefSCC().replaceNodeFunction(N, *NewF);
+      FixupMetadataReferences(&AssociatedGlobals, &OldF, NewF);
       OldF.eraseFromParent();
     }
 
     Changed |= LocalChange;
   } while (LocalChange);
 
+  AssociatedGlobals.clear();
+
   if (!Changed)
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }
 
 namespace {
 
 /// ArgPromotion - The 'by reference' to 'by value' argument promotion pass.
 struct ArgPromotion : public CallGraphSCCPass {
   // Pass identification, replacement for typeid
   static char ID;
+  AssociatedGlobalsMap AssociatedGlobals;
 
   explicit ArgPromotion(unsigned MaxElements = 3)
       : CallGraphSCCPass(ID), MaxElements(MaxElements) {
     initializeArgPromotionPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();

 Pass *llvm::createArgumentPromotionPass(unsigned MaxElements) {
   return new ArgPromotion(MaxElements);
 }
 
 bool ArgPromotion::runOnSCC(CallGraphSCC &SCC) {
   if (skipSCC(SCC))
     return false;
 
+  collectAssociatedGlobals(&SCC.getCallGraph().getModule(), &AssociatedGlobals);
+
   // Get the callgraph information that we need to update to reflect our
   // changes.
   CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
 
   LegacyAARGetter AARGetter(*this);
 
   bool Changed = false, LocalChange;
 
for (CallGraphNode *OldNode : SCC) {
         // Update the call graph for the newly promoted function.
         CallGraphNode *NewNode = CG.getOrInsertFunction(NewF);
         NewNode->stealCalledFunctionsFrom(OldNode);
         if (OldNode->getNumReferences() == 0)
           delete CG.removeFunctionFromModule(OldNode);
         else
           OldF->setLinkage(Function::ExternalLinkage);
 
+        FixupMetadataReferences(&AssociatedGlobals, OldF, NewF);
+
         // And updat ethe SCC we're iterating as well.
         SCC.ReplaceNode(OldNode, NewNode);
       }
     }
     // Remember that we changed something.
     Changed |= LocalChange;
   } while (LocalChange);
 
   return Changed;
 }
 
 bool ArgPromotion::doInitialization(CallGraph &CG) {
   return CallGraphSCCPass::doInitialization(CG);
 }

lib/Transforms/IPO/DeadArgumentElimination.cpp

bool DeadArgumentEliminationPass::DeleteDeadVarargs(Function &Fn) {
   }
 
   // Clone metadatas from the old function, including debug info descriptor.
   SmallVector<std::pair<unsigned, MDNode *>, 1> MDs;
   Fn.getAllMetadata(MDs);
   for (auto MD : MDs)
     NF->addMetadata(MD.first, *MD.second);
 
+  FixupMetadataReferences(&AssociatedGlobals, &Fn, NF);
+
   // Fix up any BlockAddresses that refer to the function.
   Fn.replaceAllUsesWith(ConstantExpr::getBitCast(NF, Fn.getType()));
   // Delete the bitcast that we just created, so that NF does not
   // appear to be address-taken.
   NF->removeDeadConstantUsers();
   // Finally, nuke the old function.
   Fn.eraseFromParent();
   return true;
for (BasicBlock &BB : *NF)
       }
 
   // Clone metadatas from the old function, including debug info descriptor.
   SmallVector<std::pair<unsigned, MDNode *>, 1> MDs;
   F->getAllMetadata(MDs);
   for (auto MD : MDs)
     NF->addMetadata(MD.first, *MD.second);
 
+  FixupMetadataReferences(&AssociatedGlobals, F, NF);
+
   // Now that the old function is dead, delete it.
   F->eraseFromParent();
 
   return true;
 }
 
 PreservedAnalyses DeadArgumentEliminationPass::run(Module &M,
                                                    ModuleAnalysisManager &) {
   bool Changed = false;
 
+  // Save information about MD_Associated metadata so that we can
+  // modify it to point at our re-created functions
+  collectAssociatedGlobals(&M, &AssociatedGlobals);
+
   // First pass: Do a simple check to see if any functions can have their "..."
   // removed.  We can do this if they never call va_start.  This loop cannot be
   // fused with the next loop, because deleting a function invalidates
   // information computed while surveying other functions.
   LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - Deleting dead varargs\n");
   for (Module::iterator I = M.begin(), E = M.end(); I != E; ) {
     Function &F = *I++;
     if (F.getFunctionType()->isVarArg())
for (Module::iterator I = M.begin(), E = M.end(); I != E; ) {
     Changed |= RemoveDeadStuffFromFunction(F);
   }
 
   // Finally, look for any unused parameters in functions with non-local
   // linkage and replace the passed in parameters with undef.
   for (auto &F : M)
     Changed |= RemoveDeadArgumentsFromCallers(F);
 
+  AssociatedGlobals.clear();
+
   if (!Changed)
     return PreservedAnalyses::all();
   return PreservedAnalyses::none();
 }

lib/Transforms/Utils/ModuleUtils.cpp

std::string llvm::getUniqueModuleId(Module *M) {
 
   MD5::MD5Result R;
   Md5.final(R);
 
   SmallString<32> Str;
   MD5::stringifyResult(R, Str);
   return ("$" + Str).str();
 }
+
+void llvm::collectAssociatedGlobals(Module *M,
+                                    AssociatedGlobalsMap *AssociatedGlobals) {
+  for (GlobalObject &GO : M->global_objects()) {
+    MDNode *MD = GO.getMetadata(LLVMContext::MD_associated);
+    if (!MD)
+      continue;
+
+    const MDOperand &Op = MD->getOperand(0);
+    if (!Op.get())
+      continue;
+
+    auto *VM = dyn_cast<ValueAsMetadata>(Op);
+    if (!VM)
+      report_fatal_error("MD_associated operand is not ValueAsMetadata");
+
+    Function *OtherFunc = dyn_cast<Function>(VM->getValue());
+    if (OtherFunc) {
+      AssociatedGlobals->insert(std::make_pair(OtherFunc, &GO));
+    }
+  }
+}
+
+void llvm::FixupMetadataReferences(AssociatedGlobalsMap *AssociatedGlobals,
+                                   Function *OldFn, Function *NewFn) {
+  std::pair<AssociatedGlobalsMap::iterator, AssociatedGlobalsMap::iterator>
+      result = AssociatedGlobals->equal_range(OldFn);
+  for (AssociatedGlobalsMap::iterator I = result.first; I != result.second;
+       I++) {
+    MDNode *MD = MDNode::get(NewFn->getContext(), ValueAsMetadata::get(NewFn));
+    I->second->setMetadata(LLVMContext::MD_associated, MD);
+  }
+}