[lld] [llvm] [llvm][lld] Respect temporal profile weights in balanced partitioning (PR #213837)

[Karim Alweheshy via llvm-commits]

================
@@ -200,34 +262,73 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
     }
   }
 
+  SmallVector<BPFunctionNode::UtilityNodeWeightT, 4> RenumberedWeights;
+  if constexpr (!std::is_same_v<UtilityNodeWeightsPtrT, std::nullptr_t>) {
+    if (!UtilityNodeWeights->empty()) {
+      RenumberedWeights.assign(Signatures.size(), 1);
+      for (auto [UN, Weight] : *UtilityNodeWeights)
+        RenumberedWeights[UN] = Weight;
+    }
+  }
+
   for (unsigned I = 0; I < Config.IterationsPerSplit; I++) {
-    unsigned NumMovedNodes =
-        runIteration(Nodes, LeftBucket, RightBucket, Signatures, RNG);
+    unsigned NumMovedNodes;
+    if (RenumberedWeights.empty())
+      NumMovedNodes = runIteration(Nodes, LeftBucket, RightBucket, Signatures,
+                                   nullptr, RNG);
+    else
+      NumMovedNodes = runIteration(
+          Nodes, LeftBucket, RightBucket, Signatures,
+          ArrayRef<BPFunctionNode::UtilityNodeWeightT>(RenumberedWeights), RNG);
     if (NumMovedNodes == 0)
       break;
   }
+  return UtilityNodeWeights;
 }
 
-unsigned BalancedPartitioning::runIteration(const FunctionNodeRange Nodes,
-                                            unsigned LeftBucket,
-                                            unsigned RightBucket,
-                                            SignaturesT &Signatures,
-                                            std::mt19937 &RNG) const {
+template <typename UtilityNodeWeightsRefT>
+unsigned BalancedPartitioning::runIteration(
+    const FunctionNodeRange Nodes, unsigned LeftBucket, unsigned RightBucket,
+    SignaturesT &Signatures, UtilityNodeWeightsRefT UtilityNodeWeights,
+    std::mt19937 &RNG) const {
   // Init signature cost caches
-  for (auto &Signature : Signatures) {
-    if (Signature.CachedGainIsValid)
-      continue;
-    unsigned L = Signature.LeftCount;
-    unsigned R = Signature.RightCount;
-    assert((L > 0 || R > 0) && "incorrect signature");
-    float Cost = logCost(L, R);
-    Signature.CachedGainLR = 0.f;
-    Signature.CachedGainRL = 0.f;
-    if (L > 0)
-      Signature.CachedGainLR = Cost - logCost(L - 1, R + 1);
-    if (R > 0)
-      Signature.CachedGainRL = Cost - logCost(L + 1, R - 1);
-    Signature.CachedGainIsValid = true;
+  if constexpr (std::is_same_v<UtilityNodeWeightsRefT, std::nullptr_t>) {
----------------
karim-alweheshy wrote:

Agreed. fa2f28cdc1d0 collapses this back to one source loop. A compile-time conditional emits the multiply only in the weighted instantiation, so the unweighted instantiation keeps the original operations without duplicating the critical calculation. The focused Support/ProfileData unit tests and all three llvm-profdata/ELF/Mach-O weighted integration tests pass.

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