[BOLT] Fix address mapping for ICP code (#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.

Author: maksfb

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()) {

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());

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