[llvm] [BOLT] Fix address mapping for ICP code (PR #70136)

[via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/70136

When we create new code for indirect code promotion optimization, we should mark it as originating from the indirect jump instruction for BOLT address translation (BAT) to map it to the original instruction.

>From e2c678e564ca7ac87c29147052b37397403f2362 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Mon, 23 Oct 2023 15:17:44 -0700
Subject: [PATCH] [BOLT] Fix address mapping for ICP code

When we create new code for indirect code promotion optimization, we
should mark it as originating from the indirect jump instruction for
BOLT address translation (BAT) to map it to the original instruction.
---
 bolt/lib/Core/BinaryFunction.cpp          |  2 +-
 bolt/lib/Passes/IndirectCallPromotion.cpp | 19 +++++++++++++++----
 bolt/test/X86/jump-table-icp.test         |  6 ++++--
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 2e1fa739e53f20b..e0727083dfa7f7a 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -533,7 +533,7 @@ void BinaryFunction::print(raw_ostream &OS, std::string Annotation) {
       if (BB->getCFIState() >= 0)
         OS << "  CFI State : " << BB->getCFIState() << '\n';
       if (opts::EnableBAT) {
-        OS << "  Input offset: " << Twine::utohexstr(BB->getInputOffset())
+        OS << "  Input offset: 0x" << Twine::utohexstr(BB->getInputOffset())
            << "\n";
       }
       if (!BB->pred_empty()) {
diff --git a/bolt/lib/Passes/IndirectCallPromotion.cpp b/bolt/lib/Passes/IndirectCallPromotion.cpp
index ea8019431cf52f9..89727233ec78b05 100644
--- a/bolt/lib/Passes/IndirectCallPromotion.cpp
+++ b/bolt/lib/Passes/IndirectCallPromotion.cpp
@@ -754,6 +754,15 @@ IndirectCallPromotion::rewriteCall(
   const bool IsTailCallOrJT =
       (MIB->isTailCall(CallInst) || Function.getJumpTable(CallInst));
 
+  // If we are tracking the indirect call/jump address, propagate the address to
+  // the ICP code.
+  const std::optional<uint32_t> IndirectInstrOffset = MIB->getOffset(CallInst);
+  if (IndirectInstrOffset) {
+    for (auto &[Symbol, Instructions] : ICPcode)
+      for (MCInst &Inst : Instructions)
+        MIB->setOffset(Inst, *IndirectInstrOffset);
+  }
+
   // Move instructions from the tail of the original call block
   // to the merge block.
 
@@ -767,10 +776,12 @@ IndirectCallPromotion::rewriteCall(
       TailInsts.push_back(*++TailInst);
 
   InstructionListType MovedInst = IndCallBlock.splitInstructions(&CallInst);
-  // Link new BBs to the original input offset of the BB where the indirect
-  // call site is, so we can map samples recorded in new BBs back to the
-  // original BB seen in the input binary (if using BAT)
-  const uint32_t OrigOffset = IndCallBlock.getInputOffset();
+  // Link new BBs to the original input offset of the indirect call site or its
+  // containing BB, so we can map samples recorded in new BBs back to the
+  // original BB seen in the input binary (if using BAT).
+  const uint32_t OrigOffset = IndirectInstrOffset
+                                  ? *IndirectInstrOffset
+                                  : IndCallBlock.getInputOffset();
 
   IndCallBlock.eraseInstructions(MethodFetchInsns.begin(),
                                  MethodFetchInsns.end());
diff --git a/bolt/test/X86/jump-table-icp.test b/bolt/test/X86/jump-table-icp.test
index 34339dc327fae52..5b989d18018b050 100644
--- a/bolt/test/X86/jump-table-icp.test
+++ b/bolt/test/X86/jump-table-icp.test
@@ -37,12 +37,14 @@ CHECK:   Successors: .Ltmp{{.*}} (mispreds: 189, count: 189), .LFT{{.*}} (mispre
 CHECK: .LFT{{.*}} (4 instructions, align : 1)
 CHECK-NEXT:   Exec Count : 881
 CHECK:   Predecessors: .LBB{{.*}}
-CHECK:   Successors: .Ltmp{{.*}} (mispreds: 138, count: 155), .Ltmp{{.*}} (mispreds: 0, count: 726)
+CHECK:     je {{.*}} # Offset: 28
+CHECK-NEXT: Successors: .Ltmp{{.*}} (mispreds: 138, count: 155), .Ltmp{{.*}} (mispreds: 0, count: 726)
 
 CHECK: .Ltmp{{.*}} (1 instructions, align : 1)
 CHECK-NEXT:   Exec Count : 726
 CHECK:   Predecessors: .LFT{{.*}}
-CHECK:   Successors: .L{{.*}} (mispreds: 126, count: 157), .L{{.*}} (mispreds: 140, count: 156), .L{{.*}} (mispreds: 134, count: 152), .L{{.*}} (mispreds: 137, count: 150), .L{{.*}} (mispreds: 129, count: 148), .L{{.*}} (mispreds: 0, count: 0)
+CHECK:     jmpq {{.*}} # Offset: 28
+CHECK-NEXT: Successors: .L{{.*}} (mispreds: 126, count: 157), .L{{.*}} (mispreds: 140, count: 156), .L{{.*}} (mispreds: 134, count: 152), .L{{.*}} (mispreds: 137, count: 150), .L{{.*}} (mispreds: 129, count: 148), .L{{.*}} (mispreds: 0, count: 0)
 
 CHECK: .Ltmp{{.*}} (5 instructions, align : 1)
 CHECK-NEXT:  Exec Count : 167