[BOLT][ICP] Propagate parent hash to newly inserted BBs for BAT

bolt/include/bolt/Core/BinaryContext.h

@@ -64,6 +64,7 @@ using namespace object;
 namespace bolt {
 
 class BinaryFunction;
+class BoltAddressTranslation;
 
 /// Information on loadable part of the file.
 struct SegmentInfo {
@@ -280,6 +281,8 @@ class BinaryContext {
   /// Internal helper for removing section name from a lookup table.
   void deregisterSectionName(const BinarySection &Section);
 
+  BoltAddressTranslation *BAT = nullptr;
+
 public:
   static Expected<std::unique_ptr<BinaryContext>> createBinaryContext(
       Triple TheTriple, std::shared_ptr<orc::SymbolStringPool> SSP,
@@ -1559,6 +1562,10 @@ class BinaryContext {
     return *IOAddressMap;
   }
 
+  void setBAT(BoltAddressTranslation *B) { BAT = B; }
+
+  BoltAddressTranslation *getBAT() const { return BAT; }
+
   raw_ostream &outs() const { return Logger.Out; }
 
   raw_ostream &errs() const { return Logger.Err; }

bolt/include/bolt/Core/BinaryFunction.h

@@ -2498,6 +2498,15 @@ class BinaryFunction {
   void
   computeBlockHashes(HashFunction HashFunction = HashFunction::Default) const;
 
+  // Updates the BAT BBHashMap for a list of newly inserted \p BBs by
+  // mapping the InputOffset of each new block to the Index and Hash of the
+  // \p OriginBB. This ensures that new blocks generated by optimizations (like
+  // ICP) are treated as part of the original block for address translation
+  // purposes.
+  void updateInsertedBBHashMap(
+      const BinaryBasicBlock &OriginBB,
+      const std::vector<std::unique_ptr<BinaryBasicBlock>> &BBs);
+
   void addDWARFUnit(DWARFUnit *Unit) { DwarfUnitMap[Unit->getOffset()] = Unit; }
 
   void removeDWARFUnit(DWARFUnit *Unit) {

bolt/lib/Passes/IndirectCallPromotion.cpp

@@ -1401,6 +1401,7 @@ Error IndirectCallPromotion::runOnFunctions(BinaryContext &BC) {
         std::vector<std::unique_ptr<BinaryBasicBlock>> NewBBs =
             rewriteCall(*BB, Inst, std::move(ICPcode), MethodInfo.second);
 
+        Function.updateInsertedBBHashMap(*BB, NewBBs);
         // Fix the CFG after inserting the new basic blocks.
         BinaryBasicBlock *MergeBlock =
             fixCFG(*BB, IsTailCall, IsJumpTable, std::move(NewBBs), Targets);

bolt/lib/Profile/StaleProfileMatching.cpp

@@ -26,7 +26,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Core/HashUtilities.h"
+#include "bolt/Profile/BoltAddressTranslation.h"
 #include "bolt/Profile/YAMLProfileReader.h"
+#include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/ADT/Bitfields.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/MC/MCPseudoProbe.h"
@@ -440,6 +442,26 @@ void BinaryFunction::computeBlockHashes(HashFunction HashFunction) const {
     BB->setHash(BlendedHashes[I].combine());
   }
 }
+
+void BinaryFunction::updateInsertedBBHashMap(
+    const BinaryBasicBlock &OriginBB,
+    const std::vector<std::unique_ptr<BinaryBasicBlock>> &BBs) {
+  if (!opts::EnableBAT || !OriginBB.getFunction() || !OriginBB.getHash())
+    return;
+
+  auto &BBHashMap =
+      BC.getBAT()->getBBHashMap(OriginBB.getFunction()->getAddress());
+
+  for (const auto &BB : BBs) {
+    if (!BB)
+      continue;
+    // Update BB and BBHashMap using the original BB's index and hash.
+    BB->setHash(OriginBB.getHash());
+    BB->setIndex(OriginBB.getIndex());
+    BBHashMap.addEntry(BB->getInputOffset(), BB->getIndex(), BB->getHash());
+  }
+}
+
 // TODO: mediate the difference between flow function construction here in BOLT
 // and in the compiler by splitting blocks with exception throwing calls at the
 // call and adding the landing pad as the successor.

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -447,6 +447,7 @@ RewriteInstance::RewriteInstance(ELFObjectFileBase *File, const int Argc,
                           BC->MII.get(), BC->MRI.get(), BC->STI.get())));
 
   BAT = std::make_unique<BoltAddressTranslation>();
+  BC->setBAT(BAT.get());
 
   if (opts::UpdateDebugSections)
     DebugInfoRewriter = std::make_unique<DWARFRewriter>(*BC);

bolt/test/X86/icp-inline.s

@@ -36,6 +36,17 @@
 # CHECK-ICP-INLINE:     callq *(%rcx,%rax,8)
 # CHECK-ICP-INLINE-NOT: callq bar
 # CHECK-ICP-INLINE:     End of Function "main"
+
+# Checks BB hash has been updated correctly after ICP optimization
+# RUN: llvm-bolt %t.exe --icp=calls --icp-calls-topn=1 --inline-small-functions\
+# RUN:   -o %t.null --lite=0 \
+# RUN:   --inline-small-functions-bytes=4 --icp-inline --print-icp \
+# RUN:   --data %t.fdata --enable-bat
+# RUN: llvm-bat-dump --dump-all %t.null | FileCheck %s --check-prefix=CHECK-BB-HASH
+# CHECK-BB-HASH: 0x17 -> 0xd hash: 0x5b24d9a60000
+# CHECK-BB-HASH: 0x1d -> 0xd hash: 0x5b24d9a60000
+# CHECK-BB-HASH: 0x1f -> 0xd hash: 0x5b24d9a60000
+
   .globl bar
 bar:
 	.cfi_startproc