[llvm] [InstrProf] Adding utility weights to BalancedPartitioning (PR #72717)

[Ellis Hoag via llvm-commits]

================
@@ -167,36 +176,43 @@ void BalancedPartitioning::runIterations(const FunctionNodeRange Nodes,
                                          unsigned RightBucket,
                                          std::mt19937 &RNG) const {
   unsigned NumNodes = std::distance(Nodes.begin(), Nodes.end());
-  DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeDegree;
+  DenseMap<uint32_t, unsigned> UtilityNodeDegree;
   for (auto &N : Nodes)
     for (auto &UN : N.UtilityNodes)
-      ++UtilityNodeDegree[UN];
+      ++UtilityNodeDegree[UN.id];
   // Remove utility nodes if they have just one edge or are connected to all
   // functions
   for (auto &N : Nodes)
     llvm::erase_if(N.UtilityNodes, [&](auto &UN) {
-      return UtilityNodeDegree[UN] <= 1 || UtilityNodeDegree[UN] >= NumNodes;
+      return UtilityNodeDegree[UN.id] <= 1 ||
+             UtilityNodeDegree[UN.id] >= NumNodes;
     });
 
   // Renumber utility nodes so they can be used to index into Signatures
-  DenseMap<BPFunctionNode::UtilityNodeT, unsigned> UtilityNodeIndex;
+  DenseMap<uint32_t, unsigned> UtilityNodeIndex;
   for (auto &N : Nodes)
     for (auto &UN : N.UtilityNodes)
-      if (!UtilityNodeIndex.count(UN))
-        UtilityNodeIndex[UN] = UtilityNodeIndex.size();
+      if (!UtilityNodeIndex.count(UN.id))
+        UtilityNodeIndex[UN.id] = UtilityNodeIndex.size();
   for (auto &N : Nodes)
     for (auto &UN : N.UtilityNodes)
-      UN = UtilityNodeIndex[UN];
+      UN.id = UtilityNodeIndex[UN.id];
 
   // Initialize signatures
   SignaturesT Signatures(/*Size=*/UtilityNodeIndex.size());
   for (auto &N : Nodes) {
     for (auto &UN : N.UtilityNodes) {
-      assert(UN < Signatures.size());
-      if (N.Bucket == LeftBucket) {
-        Signatures[UN].LeftCount++;
-      } else {
-        Signatures[UN].RightCount++;
+      assert(UN.id < Signatures.size());
+      if (N.Bucket == LeftBucket)
+        Signatures[UN.id].LeftCount++;
+      else
+        Signatures[UN.id].RightCount++;
+      // Identical utility nodes (having the same UN.id) are guaranteed to have
+      // the same weight; thus, we can get a new weight only when the signature
+      // is uninitialized.
+      if (Signatures[UN.id].Weight != UN.weight) {
+        assert(Signatures[UN.id].Weight == 1);
+        Signatures[UN.id].Weight = UN.weight;
----------------
ellishg wrote:

Since utility nodes usually come from hashes, it's possible utility nodes with the same id have different weights due to hash collisions. I don't think we need to worry too much about this case, but I'd like to remove assert and the word "guaranteed".

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