[MemProf][PGO] Prevent dropping of profile metadata during optimization

llvm/include/llvm/IR/Metadata.h

@@ -1464,6 +1464,8 @@ class MDNode : public Metadata {
   static MDNode *getMergedProfMetadata(MDNode *A, MDNode *B,
                                        const Instruction *AInstr,
                                        const Instruction *BInstr);
+  static MDNode *getMergedMemProfMetadata(MDNode *A, MDNode *B);
+  static MDNode *getMergedCallsiteMetadata(MDNode *A, MDNode *B);
 };
 
 /// Tuple of metadata.

llvm/include/llvm/Transforms/Utils/Local.h

@@ -412,6 +412,11 @@ Instruction *removeUnwindEdge(BasicBlock *BB, DomTreeUpdater *DTU = nullptr);
 bool removeUnreachableBlocks(Function &F, DomTreeUpdater *DTU = nullptr,
                              MemorySSAUpdater *MSSAU = nullptr);
 
+/// DO NOT CALL EXTERNALLY.
+/// FIXME: https://github.com/llvm/llvm-project/issues/121495
+/// Once external callers of this function are removed, either inline into
+/// combineMetadataForCSE, or internalize and remove KnownIDs parameter.
+///
 /// Combine the metadata of two instructions so that K can replace J. Some
 /// metadata kinds can only be kept if K does not move, meaning it dominated
 /// J in the original IR.

llvm/lib/Analysis/MemoryProfileInfo.cpp

@@ -347,3 +347,24 @@ template <> uint64_t CallStack<MDNode, MDNode::op_iterator>::back() const {
   return mdconst::dyn_extract<ConstantInt>(N->operands().back())
       ->getZExtValue();
 }
+
+MDNode *MDNode::getMergedMemProfMetadata(MDNode *A, MDNode *B) {
+  if (!(A && B)) {
+    return A ? A : B;
+  }
+
+  // TODO: Support more sophisticated merging, such as selecting the one with
+  // more bytes allocated, or implement support for carrying multiple allocation
+  // leaf contexts. For now, keep the first one.
+  return A;
+}
+
+MDNode *MDNode::getMergedCallsiteMetadata(MDNode *A, MDNode *B) {
+  if (!(A && B)) {
+    return A ? A : B;
+  }
+
+  // TODO: Support more sophisticated merging, which will require support for
+  // carrying multiple contexts. For now, keep the first one.
+  return A;
+}

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

@@ -788,6 +788,9 @@ Instruction *InstCombinerImpl::foldPHIArgLoadIntoPHI(PHINode &PN) {
     BasicBlock *BB = std::get<0>(Incoming);
     Value *V = std::get<1>(Incoming);
     LoadInst *LI = cast<LoadInst>(V);
+    // FIXME: https://github.com/llvm/llvm-project/issues/121495
+    // Call combineMetadataForCSE instead, so that an explicit set of KnownIDs
+    // doesn't need to be maintained here.
     combineMetadata(NewLI, LI, KnownIDs, true);
     Value *NewInVal = LI->getOperand(0);
     if (NewInVal != InVal)

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

@@ -350,6 +350,9 @@ static void combineAAMetadata(Instruction *ReplInst, Instruction *I) {
       LLVMContext::MD_noalias,      LLVMContext::MD_invariant_group,
       LLVMContext::MD_access_group, LLVMContext::MD_prof,
       LLVMContext::MD_memprof,      LLVMContext::MD_callsite};
+  // FIXME: https://github.com/llvm/llvm-project/issues/121495
+  // Use custom AA metadata combining handling instead of combineMetadata, which
+  // is meant for CSE and will drop any metadata not in the KnownIDs list.
   combineMetadata(ReplInst, I, KnownIDs, true);
 }
 

llvm/lib/Transforms/Utils/Local.cpp

@@ -3308,6 +3308,9 @@ bool llvm::removeUnreachableBlocks(Function &F, DomTreeUpdater *DTU,
   return Changed;
 }
 
+// FIXME: https://github.com/llvm/llvm-project/issues/121495
+// Once external callers of this function are removed, either inline into
+// combineMetadataForCSE, or internalize and remove KnownIDs parameter.
 void llvm::combineMetadata(Instruction *K, const Instruction *J,
                            ArrayRef<unsigned> KnownIDs, bool DoesKMove) {
   SmallVector<std::pair<unsigned, MDNode *>, 4> Metadata;
@@ -3320,6 +3323,10 @@ void llvm::combineMetadata(Instruction *K, const Instruction *J,
 
     switch (Kind) {
       default:
+        // FIXME: https://github.com/llvm/llvm-project/issues/121495
+        // Change to removing only explicitly listed other metadata, and assert
+        // on unknown metadata, to avoid inadvertently dropping newly added
+        // metadata types.
         K->setMetadata(Kind, nullptr); // Remove unknown metadata
         break;
       case LLVMContext::MD_dbg:
@@ -3380,8 +3387,10 @@ void llvm::combineMetadata(Instruction *K, const Instruction *J,
             MDNode::getMostGenericAlignmentOrDereferenceable(JMD, KMD));
         break;
       case LLVMContext::MD_memprof:
+        K->setMetadata(Kind, MDNode::getMergedMemProfMetadata(KMD, JMD));
+        break;
       case LLVMContext::MD_callsite:
-        // Preserve !memprof and !callsite metadata on K.
+        K->setMetadata(Kind, MDNode::getMergedCallsiteMetadata(KMD, JMD));
         break;
       case LLVMContext::MD_preserve_access_index:
         // Preserve !preserve.access.index in K.