[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;
----------------
alokkrsharma wrote:

Fair point. I will remove those ternary guards in the next revision — scaleToBlockCount already returns 0 for a 0 weight/total, so they're not needed.

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