[llvm] [MergeFunctions] Preserve instruction-level profile metadata during merging (PR #208009)

[Alok Kumar Sharma via llvm-commits]

================
@@ -858,28 +889,235 @@ void MergeFunctions::writeAlias(Function *F, Function *G) {
 
 // If needed, replace G with an alias to F if possible, or a thunk to F if
 // profitable. Returns false if neither is the case. If \p G is not needed (i.e.
-// it is discardable and unused), \p G is removed directly.
-bool MergeFunctions::writeThunkOrAliasIfNeeded(Function *F, Function *G) {
-  if (G->isDiscardableIfUnused() && G->use_empty() && !MergeFunctionsPDI) {
+// it is discardable and unused), \p G is removed directly. If \p MergeProfile
+// is set, G's profile metadata is merged into F.
+bool MergeFunctions::writeThunkOrAliasIfNeeded(Function *F, Function *G,
+                                               bool MergeProfile) {
+  bool ShouldErase =
+      G->isDiscardableIfUnused() && G->use_empty() && !MergeFunctionsPDI;
+  bool ShouldAlias = canCreateAliasFor(G);
+  bool ShouldThunk = canCreateThunkFor(F);
+
+  if (!ShouldErase && !ShouldAlias && !ShouldThunk)
+    return false;
+
+  if (MergeProfile)
+    mergeInstrProfMetadataInto(F, G);
+
+  if (ShouldErase) {
     G->eraseFromParent();
     return true;
   }
-  if (canCreateAliasFor(G)) {
+
+  if (ShouldAlias) {
     writeAlias(F, G);
     return true;
   }
-  if (canCreateThunkFor(F)) {
+  if (ShouldThunk) {
     writeThunk(F, G);
     return true;
   }
-  return false;
+
+  llvm_unreachable("Erase, alias or thunk must apply");
 }
 
 /// Returns true if \p F is either weak_odr or linkonce_odr.
 static bool isODR(const Function *F) {
   return F->hasWeakODRLinkage() || F->hasLinkOnceODRLinkage();
 }
 
+static uint64_t getBlockCountForMerging(const BlockFrequencyInfo *BFI,
+                                        const BasicBlock *BB,
+                                        std::optional<uint64_t> EC) {
+  if (BFI) {
+    if (auto Count = BFI->getBlockProfileCount(BB, /*AllowSynthetic=*/true))
+      return *Count;
+    // BFI present but no count for this block, fall through to EC.
+  }
+  if (EC)
+    return *EC;
+  return 1;
+}
+
+// The branch/value profile weights are relative within a function. Before
+// merge we need to normalize weights to absolute counts.
+// (weight * BlockCount / TotalWeight)
+static uint64_t scaleToBlockCount(uint64_t Weight, uint64_t TotalWeight,
+                                  uint64_t BlockCount) {
+  if (Weight == 0 || TotalWeight == 0 || BlockCount == 0)
+    return 0;
+  APInt Num(128, BlockCount);
+  Num *= APInt(128, Weight);
+  APInt Den(128, TotalWeight);
+  Num = (Num + Den.lshr(1)).udiv(Den);
+  assert(Num.getActiveBits() <= 64 &&
+         "scaleToBlockCount: result exceeds uint64_t; Weight > TotalWeight?");
+  return Num.getLimitedValue();
+}
+
+// Combine the scaled branch_weights of corresponding instructions of F and G.
+static void
+mergeBranchWeightsOnInstructions(Instruction *DstI, const Instruction *SrcI,
+                                 const BlockFrequencyInfo *DstBFI,
+                                 const BlockFrequencyInfo *SrcBFI,
+                                 std::optional<uint64_t> DstEntryCount,
+                                 std::optional<uint64_t> SrcEntryCount) {
+  SmallVector<uint32_t, 8> DstWeights, SrcWeights;
+  bool HasDst = extractBranchWeights(*DstI, DstWeights);
+  bool HasSrc = extractBranchWeights(*SrcI, SrcWeights);
+  if (!HasDst && !HasSrc)
+    return;
+
+  uint64_t DstBlockCount =
+      getBlockCountForMerging(DstBFI, DstI->getParent(), DstEntryCount);
+  uint64_t SrcBlockCount =
+      getBlockCountForMerging(SrcBFI, SrcI->getParent(), SrcEntryCount);
+
+  uint64_t DstTotal = 0, SrcTotal = 0;
+  if (HasDst)
+    extractProfTotalWeight(*DstI, DstTotal);
+  if (HasSrc)
+    extractProfTotalWeight(*SrcI, SrcTotal);
+
+  size_t NumWeights = std::max(HasDst ? DstWeights.size() : size_t{0},
+                               HasSrc ? SrcWeights.size() : size_t{0});
+  SmallVector<uint64_t, 8> MergedWeights;
+  MergedWeights.reserve(NumWeights);
+  for (size_t I = 0; I < NumWeights; ++I) {
+    uint64_t DstW = (HasDst && I < DstWeights.size()) ? DstWeights[I] : 0;
+    uint64_t SrcW = (HasSrc && I < SrcWeights.size()) ? SrcWeights[I] : 0;
+    uint64_t DstAbs =
+        HasDst ? scaleToBlockCount(DstW, DstTotal, DstBlockCount) : 0;
+    uint64_t SrcAbs =
+        HasSrc ? scaleToBlockCount(SrcW, SrcTotal, SrcBlockCount) : 0;
+    MergedWeights.push_back(SaturatingAdd(DstAbs, SrcAbs));
+  }
+
+  bool IsExpected =
+      hasBranchWeightOrigin(*DstI) && hasBranchWeightOrigin(*SrcI);
+  setFittedBranchWeights(*DstI, MergedWeights, IsExpected);
+}
+
+// Accumulate scaled value profile counts of Instruction I into Merged.
+static void addScaledValueProfile(const Instruction &I, InstrProfValueKind Kind,
+                                  uint64_t BlockCount,
+                                  DenseMap<uint64_t, uint64_t> &Merged) {
+  uint64_t Total = 0;
+  SmallVector<InstrProfValueData, 4> VDs =
+      getValueProfDataFromInst(I, Kind, /*MaxNumValueData=*/UINT32_MAX, Total);
+  if (VDs.empty() || Total == 0)
+    return;
+  for (const InstrProfValueData &VD : VDs) {
+    uint64_t Abs = scaleToBlockCount(VD.Count, Total, BlockCount);
+    Merged[VD.Value] = SaturatingAdd(Merged[VD.Value], Abs);
+  }
+}
+
+// Merge (union) scaled value profiles of Dst and Src.
+static void
+mergeValueProfileOnInstructions(Instruction *DstI, const Instruction *SrcI,
+                                const BlockFrequencyInfo *DstBFI,
+                                const BlockFrequencyInfo *SrcBFI,
+                                std::optional<uint64_t> DstEntryCount,
+                                std::optional<uint64_t> SrcEntryCount) {
+  MDNode *DstProf = DstI->getMetadata(LLVMContext::MD_prof);
+  MDNode *SrcProf = SrcI->getMetadata(LLVMContext::MD_prof);
+  bool HasDst = DstProf && isValueProfileMD(DstProf);
+  bool HasSrc = SrcProf && isValueProfileMD(SrcProf);
+  if (!HasDst && !HasSrc)
+    return;
+
+  auto *DstKind =
+      HasDst ? mdconst::dyn_extract<ConstantInt>(DstProf->getOperand(1))
+             : nullptr;
+  auto *SrcKind =
+      HasSrc ? mdconst::dyn_extract<ConstantInt>(SrcProf->getOperand(1))
+             : nullptr;
+  if (HasDst && HasSrc && DstKind && SrcKind &&
+      DstKind->getZExtValue() != SrcKind->getZExtValue()) {
+    DstI->setMetadata(LLVMContext::MD_prof, nullptr);
+    return;
+  }
+
+  const ConstantInt *KindCI = DstKind ? DstKind : SrcKind;
+  if (!KindCI) {
+    DstI->setMetadata(LLVMContext::MD_prof, nullptr);
+    return;
+  }
+
+  InstrProfValueKind Kind =
+      static_cast<InstrProfValueKind>(KindCI->getZExtValue());
+
+  DenseMap<uint64_t, uint64_t> Merged;
+  uint64_t DstBlockCount =
+      getBlockCountForMerging(DstBFI, DstI->getParent(), DstEntryCount);
+  uint64_t SrcBlockCount =
+      getBlockCountForMerging(SrcBFI, SrcI->getParent(), SrcEntryCount);
+  if (HasDst)
+    addScaledValueProfile(*DstI, Kind, DstBlockCount, Merged);
----------------
alokkrsharma wrote:

Agreed, I will drop the scaling and just add VP counts by key in the next revision, since they're already absolute.

https://github.com/llvm/llvm-project/pull/208009