[llvm] [BOLT] Remove mutable from BB:LayoutIndex (PR #93224)

[shaw young via llvm-commits]

https://github.com/shawbyoung updated https://github.com/llvm/llvm-project/pull/93224

>From 4486b9c1a135fd1c758ac30379a133366fcd9917 Mon Sep 17 00:00:00 2001
From: shawbyoung <shawbyoung at gmail.com>
Date: Thu, 23 May 2024 10:58:35 -0700
Subject: [PATCH 1/2] [BOLT][NFC] Remove mutable from BB:LayoutIndex

---
 bolt/include/bolt/Core/BinaryBasicBlock.h |  4 ++--
 bolt/lib/Core/BinaryFunction.cpp          | 10 +++-------
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryBasicBlock.h b/bolt/include/bolt/Core/BinaryBasicBlock.h
index bc95e2c4de3a1..a57b70714fe38 100644
--- a/bolt/include/bolt/Core/BinaryBasicBlock.h
+++ b/bolt/include/bolt/Core/BinaryBasicBlock.h
@@ -115,7 +115,7 @@ class BinaryBasicBlock {
   unsigned Index{InvalidIndex};
 
   /// Index in the current layout.
-  mutable unsigned LayoutIndex{InvalidIndex};
+  unsigned LayoutIndex{InvalidIndex};
 
   /// Number of pseudo instructions in this block.
   uint32_t NumPseudos{0};
@@ -891,7 +891,7 @@ class BinaryBasicBlock {
   }
 
   /// Set layout index. To be used by BinaryFunction.
-  void setLayoutIndex(unsigned Index) const { LayoutIndex = Index; }
+  void setLayoutIndex(unsigned Index) { LayoutIndex = Index; }
 
   /// Needed by graph traits.
   BinaryFunction *getParent() const { return getFunction(); }
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 10b93e702984f..51dc5a2cbf0d6 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3636,8 +3636,8 @@ bool BinaryFunction::forEachEntryPoint(EntryPointCallbackTy Callback) const {
 
 BinaryFunction::BasicBlockListType BinaryFunction::dfs() const {
   BasicBlockListType DFS;
-  unsigned Index = 0;
   std::stack<BinaryBasicBlock *> Stack;
+  std::set<BinaryBasicBlock *> Visited;
 
   // Push entry points to the stack in reverse order.
   //
@@ -3654,17 +3654,13 @@ BinaryFunction::BasicBlockListType BinaryFunction::dfs() const {
   for (BinaryBasicBlock *const BB : reverse(EntryPoints))
     Stack.push(BB);
 
-  for (BinaryBasicBlock &BB : blocks())
-    BB.setLayoutIndex(BinaryBasicBlock::InvalidIndex);
-
   while (!Stack.empty()) {
     BinaryBasicBlock *BB = Stack.top();
     Stack.pop();
 
-    if (BB->getLayoutIndex() != BinaryBasicBlock::InvalidIndex)
+    if (Visited.find(BB) != Visited.end())
       continue;
-
-    BB->setLayoutIndex(Index++);
+    Visited.insert(BB);
     DFS.push_back(BB);
 
     for (BinaryBasicBlock *SuccBB : BB->landing_pads()) {

>From c269abc0dc72bbe45b86a7ad79f8a68eb1ecef8d Mon Sep 17 00:00:00 2001
From: shawbyoung <shawbyoung at gmail.com>
Date: Fri, 24 May 2024 14:46:58 -0700
Subject: [PATCH 2/2] [BOLT][NFC] Fix invalid BB::getLayoutIndex calls

---
 bolt/lib/Passes/IdenticalCodeFolding.cpp | 28 ++++++++++++++++++++++--
 bolt/lib/Profile/YAMLProfileWriter.cpp   | 13 +++++++++--
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/bolt/lib/Passes/IdenticalCodeFolding.cpp b/bolt/lib/Passes/IdenticalCodeFolding.cpp
index 87eba10354a37..e43c7e1a32c20 100644
--- a/bolt/lib/Passes/IdenticalCodeFolding.cpp
+++ b/bolt/lib/Passes/IdenticalCodeFolding.cpp
@@ -180,11 +180,26 @@ static bool isIdenticalWith(const BinaryFunction &A, const BinaryFunction &B,
 
   const BinaryContext &BC = A.getBinaryContext();
 
+  std::unordered_map<const BinaryBasicBlock*, unsigned> LayoutIndiciesA;
+  for (unsigned I = 0; I < OrderA.size(); I++)
+    LayoutIndiciesA[OrderA[I]] = I;
+
+  std::unordered_map<const BinaryBasicBlock*, unsigned> LayoutIndiciesB;
+  for (unsigned I = 0; I < OrderB.size(); I++)
+    LayoutIndiciesB[OrderB[I]] = I;
+
   auto BBI = OrderB.begin();
   for (const BinaryBasicBlock *BB : OrderA) {
     const BinaryBasicBlock *OtherBB = *BBI;
 
-    if (BB->getLayoutIndex() != OtherBB->getLayoutIndex())
+    auto LayoutIndiciesAIt = LayoutIndiciesA.find(BB);
+    assert(LayoutIndiciesAIt != LayoutIndiciesA.end());
+    unsigned BBLayoutIndex = LayoutIndiciesAIt->second;
+
+    auto LayoutIndiciesBIt = LayoutIndiciesB.find(OtherBB);
+    assert(LayoutIndiciesBIt != LayoutIndiciesB.end());
+    unsigned OtherBBLayoutIndex = LayoutIndiciesBIt->second;
+    if (BBLayoutIndex != OtherBBLayoutIndex)
       return false;
 
     // Compare successor basic blocks.
@@ -195,7 +210,16 @@ static bool isIdenticalWith(const BinaryFunction &A, const BinaryFunction &B,
     auto SuccBBI = OtherBB->succ_begin();
     for (const BinaryBasicBlock *SuccBB : BB->successors()) {
       const BinaryBasicBlock *SuccOtherBB = *SuccBBI;
-      if (SuccBB->getLayoutIndex() != SuccOtherBB->getLayoutIndex())
+
+      LayoutIndiciesAIt = LayoutIndiciesA.find(SuccBB);
+      assert(LayoutIndiciesAIt != LayoutIndiciesA.end());
+      unsigned SuccBBLayoutIndex = LayoutIndiciesAIt->second;
+
+      LayoutIndiciesBIt = LayoutIndiciesB.find(SuccOtherBB);
+      assert(LayoutIndiciesBIt != LayoutIndiciesB.end());
+      unsigned SuccOtherBBLayoutIndex = LayoutIndiciesBIt->second;
+
+      if (SuccBBLayoutIndex != SuccOtherBBLayoutIndex)
         return false;
       ++SuccBBI;
     }
diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index ef04ba0d21ad7..f0721a2a5d16b 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -69,9 +69,15 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
   llvm::copy(UseDFS ? BF.dfs() : BF.getLayout().blocks(),
              std::back_inserter(Order));
 
+  std::unordered_map<const BinaryBasicBlock*, unsigned> LayoutIndicies;
+  for (unsigned I = 0; I < Order.size(); I++)
+    LayoutIndicies[Order[I]] = I;
+
   for (const BinaryBasicBlock *BB : Order) {
     yaml::bolt::BinaryBasicBlockProfile YamlBB;
-    YamlBB.Index = BB->getLayoutIndex();
+    auto It = LayoutIndicies.find(BB);
+    assert(It != LayoutIndicies.end());
+    YamlBB.Index = It->second;
     YamlBB.NumInstructions = BB->getNumNonPseudos();
     YamlBB.Hash = BB->getHash();
 
@@ -160,7 +166,10 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
     auto BranchInfo = BB->branch_info_begin();
     for (const BinaryBasicBlock *Successor : BB->successors()) {
       yaml::bolt::SuccessorInfo YamlSI;
-      YamlSI.Index = Successor->getLayoutIndex();
+
+      auto It = LayoutIndicies.find(Successor);
+      assert(It != LayoutIndicies.end());
+      YamlSI.Index = It->second;
       YamlSI.Count = BranchInfo->Count;
       YamlSI.Mispreds = BranchInfo->MispredictedCount;