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

[Mircea Trofin via llvm-commits]

================
@@ -858,28 +902,236 @@ 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)
+    mergeProfMetadataInto(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);
+  return Num.getLimitedValue();
+}
+
+// Combine the scaled branch_weights of corresponding instructions of F and G.
+static void mergeBranchWeightsOnInstructions(Instruction *FI,
+                                             const Instruction *GI,
+                                             const BlockFrequencyInfo *BfiF,
+                                             const BlockFrequencyInfo *BfiG,
+                                             std::optional<uint64_t> FEC,
+                                             std::optional<uint64_t> GEC) {
+  SmallVector<uint32_t, 8> FWeights, GWeights;
+  bool HasF = extractBranchWeights(*FI, FWeights);
+  bool HasG = extractBranchWeights(*GI, GWeights);
+  if (!HasF && !HasG)
+    return;
+
+  uint64_t BlockCountF = getBlockCountForMerging(BfiF, FI->getParent(), FEC);
+  uint64_t BlockCountG = getBlockCountForMerging(BfiG, GI->getParent(), GEC);
+
+  uint64_t TotalF = 0, TotalG = 0;
+  if (HasF)
+    extractProfTotalWeight(*FI, TotalF);
+  if (HasG)
+    extractProfTotalWeight(*GI, TotalG);
+
+  size_t NumWeights = std::max(HasF ? FWeights.size() : size_t{0},
+                               HasG ? GWeights.size() : size_t{0});
+  SmallVector<uint64_t, 8> MergedWeights;
+  MergedWeights.reserve(NumWeights);
+  for (size_t I = 0; I < NumWeights; ++I) {
+    uint64_t FW = (HasF && I < FWeights.size()) ? FWeights[I] : 0;
+    uint64_t GW = (HasG && I < GWeights.size()) ? GWeights[I] : 0;
+    uint64_t AbsF = HasF ? scaleToBlockCount(FW, TotalF, BlockCountF) : 0;
+    uint64_t AbsG = HasG ? scaleToBlockCount(GW, TotalG, BlockCountG) : 0;
+    MergedWeights.push_back(SaturatingAdd(AbsF, AbsG));
+  }
+
+  bool IsExpected = hasBranchWeightOrigin(*FI) || hasBranchWeightOrigin(*GI);
+  setFittedBranchWeights(*FI, 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 F and G.
+static void mergeValueProfileOnInstructions(Instruction *FI,
+                                            const Instruction *GI,
+                                            const BlockFrequencyInfo *BfiF,
+                                            const BlockFrequencyInfo *BfiG,
+                                            std::optional<uint64_t> FEC,
+                                            std::optional<uint64_t> GEC) {
+  MDNode *ProfF = FI->getMetadata(LLVMContext::MD_prof);
+  MDNode *ProfG = GI->getMetadata(LLVMContext::MD_prof);
+  bool HasF = ProfF && isValueProfileMD(ProfF);
+  bool HasG = ProfG && isValueProfileMD(ProfG);
+  if (!HasF && !HasG)
+    return;
+
+  auto *KindF =
+      HasF ? mdconst::dyn_extract<ConstantInt>(ProfF->getOperand(1)) : nullptr;
+  auto *KindG =
+      HasG ? mdconst::dyn_extract<ConstantInt>(ProfG->getOperand(1)) : nullptr;
+  if (HasF && HasG && KindF && KindG &&
+      KindF->getZExtValue() != KindG->getZExtValue()) {
+    FI->setMetadata(LLVMContext::MD_prof, nullptr);
+    return;
+  }
+
+  const ConstantInt *KindCI = KindF ? KindF : KindG;
+  if (!KindCI) {
+    FI->setMetadata(LLVMContext::MD_prof, nullptr);
+    return;
+  }
+
+  InstrProfValueKind Kind =
+      static_cast<InstrProfValueKind>(KindCI->getZExtValue());
+
+  DenseMap<uint64_t, uint64_t> Merged;
+  uint64_t BCF = getBlockCountForMerging(BfiF, FI->getParent(), FEC);
+  uint64_t BCG = getBlockCountForMerging(BfiG, GI->getParent(), GEC);
+  if (HasF)
+    addScaledValueProfile(*FI, Kind, BCF, Merged);
+  if (HasG)
+    addScaledValueProfile(*GI, Kind, BCG, Merged);
+
+  if (Merged.empty())
+    return;
+
+  SmallVector<InstrProfValueData, 8> VDs;
+  VDs.reserve(Merged.size());
+  uint64_t Sum = 0;
+  for (auto &[Value, Count] : Merged) {
+    VDs.push_back({Value, Count});
+    Sum = SaturatingAdd(Sum, Count);
+  }
+  llvm::sort(VDs, [](const InstrProfValueData &A, const InstrProfValueData &B) {
+    return A.Count > B.Count;
+  });
+  annotateValueSite(*FI->getFunction()->getParent(), *FI, VDs, Sum, Kind,
+                    VDs.size());
+}
+
+const BlockFrequencyInfo &MergeFunctions::getBFI(Function &F) {
+  if (FAM)
+    return FAM->getResult<BlockFrequencyAnalysis>(F);
+  auto &Entry = CachedBFI[&F];
+  if (!Entry)
+    Entry = std::make_unique<CachedProfileAnalyses>(F);
+  return Entry->BFI;
+}
+
+void MergeFunctions::mergeProfMetadataInto(Function *F, Function *G) {
----------------
mtrofin wrote:

sorry for the spam, github had an issue that kept saying "can't save your comment", so looks like my hitting "retry" did... this.

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