[BOLT][NFC] Move Offset annotation to Group 1

Summary:
Move the annotation to avoid dynamic memory allocations.
Improves the CPU time of instrumenting a large binary by 1% (+-0.8%, p-value 0.01)

Test Plan: NFC

Reviewers: maksfb

FBD30091656

14 files changed, 69 insertions, 34 deletions

diff --git a/bolt/include/bolt/Core/MCPlus.h b/bolt/include/bolt/Core/MCPlus.h
index 67430bb..0abb70a 100644
--- a/bolt/include/bolt/Core/MCPlus.h
+++ b/bolt/include/bolt/Core/MCPlus.h
@@ -65,6 +65,7 @@ public:
     kJumpTable,           /// Jump Table.
     kTailCall,            /// Tail call.
     kConditionalTailCall, /// CTC.
+    kOffset,              /// Offset in the function.
     kGeneric              /// First generic annotation.
   };
 
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 6576263..8a4497b 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1075,6 +1075,18 @@ public:
   /// branch. Return true if the instruction was converted.
   bool unsetConditionalTailCall(MCInst &Inst);
 
+  /// Return offset of \p Inst in the original function, if available.
+  Optional<uint32_t> getOffset(const MCInst &Inst) const;
+
+  /// Return the offset if the annotation is present, or \p Default otherwise.
+  uint32_t getOffsetWithDefault(const MCInst &Inst, uint32_t Default) const;
+
+  /// Set offset of \p Inst in the original function.
+  bool setOffset(MCInst &Inst, uint32_t Offset, AllocatorIdTy AllocatorId = 0);
+
+  /// Remove offset annotation.
+  bool clearOffset(MCInst &Inst);
+
   /// Return MCSymbol that represents a target of this instruction at a given
   /// operand number \p OpNum. If there's no symbol associated with
   /// the operand - return nullptr.
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 5ba0dcd..5368d2a 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -1632,6 +1632,8 @@ void BinaryContext::printInstruction(raw_ostream &OS, const MCInst &Instruction,
       OS << " # UNKNOWN CONTROL FLOW";
     }
   }
+  if (Optional<uint32_t> Offset = MIB->getOffset(Instruction))
+    OS << " # Offset: " << *Offset;
 
   MIB->printAnnotations(Instruction, OS);
 
diff --git a/bolt/lib/Core/BinaryEmitter.cpp b/bolt/lib/Core/BinaryEmitter.cpp
index 61ebeae..6bf7ab0 100644
--- a/bolt/lib/Core/BinaryEmitter.cpp
+++ b/bolt/lib/Core/BinaryEmitter.cpp
@@ -468,8 +468,8 @@ void BinaryEmitter::emitFunctionBody(BinaryFunction &BF, bool EmitColdPart,
       // Prepare to tag this location with a label if we need to keep track of
       // the location of calls/returns for BOLT address translation maps
       if (!EmitCodeOnly && BF.requiresAddressTranslation() &&
-          BC.MIB->hasAnnotation(Instr, "Offset")) {
-        const auto Offset = BC.MIB->getAnnotationAs<uint32_t>(Instr, "Offset");
+          BC.MIB->getOffset(Instr)) {
+        const uint32_t Offset = *BC.MIB->getOffset(Instr);
         MCSymbol *LocSym = BC.Ctx->createTempSymbol();
         Streamer.emitLabel(LocSym);
         BB->getLocSyms().emplace_back(Offset, LocSym);
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 9a2344e..5b81dfa 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1406,7 +1406,7 @@ add_instruction:
 
     // Record offset of the instruction for profile matching.
     if (BC.keepOffsetForInstruction(Instruction))
-      MIB->addAnnotation(Instruction, "Offset", static_cast<uint32_t>(Offset));
+      MIB->setOffset(Instruction, static_cast<uint32_t>(Offset));
 
     if (BC.MIB->isNoop(Instruction)) {
       // NOTE: disassembly loses the correct size information for noops.
@@ -1989,9 +1989,8 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
         break;
       }
     }
-    if (LastNonNop && !MIB->hasAnnotation(*LastNonNop, "Offset"))
-      MIB->addAnnotation(*LastNonNop, "Offset", static_cast<uint32_t>(Offset),
-                         AllocatorId);
+    if (LastNonNop && !MIB->getOffset(*LastNonNop))
+      MIB->setOffset(*LastNonNop, static_cast<uint32_t>(Offset), AllocatorId);
   };
 
   for (auto I = Instructions.begin(), E = Instructions.end(); I != E; ++I) {
@@ -2014,9 +2013,8 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
     bool IsLKMarker = BC.LKMarkers.count(InstrInputAddr);
     // Mark all nops with Offset for profile tracking purposes.
     if (MIB->isNoop(Instr) || IsLKMarker) {
-      if (!MIB->hasAnnotation(Instr, "Offset"))
-        MIB->addAnnotation(Instr, "Offset", static_cast<uint32_t>(Offset),
-                           AllocatorId);
+      if (!MIB->getOffset(Instr))
+        MIB->setOffset(Instr, static_cast<uint32_t>(Offset), AllocatorId);
       if (IsSDTMarker || IsLKMarker)
         HasSDTMarker = true;
       else
@@ -2216,7 +2214,7 @@ void BinaryFunction::postProcessCFG() {
   if (!requiresAddressTranslation() && !opts::Instrument) {
     for (BinaryBasicBlock *BB : layout())
       for (MCInst &Inst : *BB)
-        BC.MIB->removeAnnotation(Inst, "Offset");
+        BC.MIB->clearOffset(Inst);
   }
 
   assert((!isSimple() || validateCFG()) &&
@@ -2233,8 +2231,7 @@ void BinaryFunction::calculateMacroOpFusionStats() {
 
     // Check offset of the second instruction.
     // FIXME: arch-specific.
-    const uint32_t Offset =
-        BC.MIB->getAnnotationWithDefault<uint32_t>(*std::next(II), "Offset", 0);
+    const uint32_t Offset = BC.MIB->getOffsetWithDefault(*std::next(II), 0);
     if (!Offset || (getAddress() + Offset) % 64)
       continue;
 
@@ -4325,8 +4322,7 @@ MCInst *BinaryFunction::getInstructionAtOffset(uint64_t Offset) {
 
     for (MCInst &Inst : *BB) {
       constexpr uint32_t InvalidOffset = std::numeric_limits<uint32_t>::max();
-      if (Offset == BC.MIB->getAnnotationWithDefault<uint32_t>(Inst, "Offset",
-                                                               InvalidOffset))
+      if (Offset == BC.MIB->getOffsetWithDefault(Inst, InvalidOffset))
         return &Inst;
     }
 
diff --git a/bolt/lib/Core/MCPlusBuilder.cpp b/bolt/lib/Core/MCPlusBuilder.cpp
index b02e139..d5d135b 100644
--- a/bolt/lib/Core/MCPlusBuilder.cpp
+++ b/bolt/lib/Core/MCPlusBuilder.cpp
@@ -228,6 +228,33 @@ bool MCPlusBuilder::unsetConditionalTailCall(MCInst &Inst) {
   return true;
 }
 
+Optional<uint32_t> MCPlusBuilder::getOffset(const MCInst &Inst) const {
+  Optional<int64_t> Value = getAnnotationOpValue(Inst, MCAnnotation::kOffset);
+  if (!Value)
+    return NoneType();
+  return static_cast<uint32_t>(*Value);
+}
+
+uint32_t MCPlusBuilder::getOffsetWithDefault(const MCInst &Inst,
+                                             uint32_t Default) const {
+  if (Optional<uint32_t> Offset = getOffset(Inst))
+    return *Offset;
+  return Default;
+}
+
+bool MCPlusBuilder::setOffset(MCInst &Inst, uint32_t Offset,
+                              AllocatorIdTy AllocatorId) {
+  setAnnotationOpValue(Inst, MCAnnotation::kOffset, Offset, AllocatorId);
+  return true;
+}
+
+bool MCPlusBuilder::clearOffset(MCInst &Inst) {
+  if (!hasAnnotation(Inst, MCAnnotation::kOffset))
+    return false;
+  removeAnnotation(Inst, MCAnnotation::kOffset);
+  return true;
+}
+
 bool MCPlusBuilder::hasAnnotation(const MCInst &Inst, unsigned Index) const {
   const MCInst *AnnotationInst = getAnnotationInst(Inst);
   if (!AnnotationInst)
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 4ac06e0..fb6e387 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -629,10 +629,9 @@ void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
 
       // Now record preserved annotations separately and then strip annotations.
       for (auto II = BB->begin(); II != BB->end(); ++II) {
-        if (BF.requiresAddressTranslation() &&
-            BC.MIB->hasAnnotation(*II, "Offset"))
-          PreservedOffsetAnnotations.emplace_back(
-              &(*II), BC.MIB->getAnnotationAs<uint32_t>(*II, "Offset"));
+        if (BF.requiresAddressTranslation() && BC.MIB->getOffset(*II))
+          PreservedOffsetAnnotations.emplace_back(&(*II),
+                                                  *BC.MIB->getOffset(*II));
         BC.MIB->stripAnnotations(*II);
       }
     }
@@ -647,7 +646,7 @@ void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
 
   // Reinsert preserved annotations we need during code emission.
   for (const std::pair<MCInst *, uint32_t> &Item : PreservedOffsetAnnotations)
-    BC.MIB->addAnnotation<uint32_t>(*Item.first, "Offset", Item.second);
+    BC.MIB->setOffset(*Item.first, Item.second);
 }
 
 namespace {
diff --git a/bolt/lib/Passes/Instrumentation.cpp b/bolt/lib/Passes/Instrumentation.cpp
index 50cb6f9..c0084e04 100644
--- a/bolt/lib/Passes/Instrumentation.cpp
+++ b/bolt/lib/Passes/Instrumentation.cpp
@@ -377,7 +377,7 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
 
     for (auto I = BB.begin(); I != BB.end(); ++I) {
       const MCInst &Inst = *I;
-      if (!BC.MIB->hasAnnotation(Inst, "Offset"))
+      if (!BC.MIB->getOffset(Inst))
         continue;
 
       const bool IsJumpTable = Function.getJumpTable(Inst);
@@ -389,7 +389,7 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
                BC.MIB->isUnsupportedBranch(Inst.getOpcode()))
         continue;
 
-      uint32_t FromOffset = BC.MIB->getAnnotationAs<uint32_t>(Inst, "Offset");
+      const uint32_t FromOffset = *BC.MIB->getOffset(Inst);
       const MCSymbol *Target = BC.MIB->getTargetSymbol(Inst);
       BinaryBasicBlock *TargetBB = Function.getBasicBlockForLabel(Target);
       uint32_t ToOffset = TargetBB ? TargetBB->getInputOffset() : 0;
@@ -465,9 +465,9 @@ void Instrumentation::instrumentFunction(BinaryFunction &Function,
       // if it was branching to the end of the function as a result of
       // __builtin_unreachable(), in which case it was deleted by fixBranches.
       // Ignore this case. FIXME: force fixBranches() to preserve the offset.
-      if (!BC.MIB->hasAnnotation(*LastInstr, "Offset"))
+      if (!BC.MIB->getOffset(*LastInstr))
         continue;
-      FromOffset = BC.MIB->getAnnotationAs<uint32_t>(*LastInstr, "Offset");
+      FromOffset = *BC.MIB->getOffset(*LastInstr);
 
       // Do not instrument edges in the spanning tree
       if (STOutSet[&BB].find(FTBB) != STOutSet[&BB].end()) {
diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
index 87be768..bb110cdb 100644
--- a/bolt/lib/Profile/DataAggregator.cpp
+++ b/bolt/lib/Profile/DataAggregator.cpp
@@ -914,7 +914,7 @@ bool DataAggregator::recordTrace(
       const MCInst *Instr = BB->getLastNonPseudoInstr();
       uint64_t Offset = 0;
       if (Instr)
-        Offset = BC.MIB->getAnnotationWithDefault<uint32_t>(*Instr, "Offset");
+        Offset = BC.MIB->getOffsetWithDefault(*Instr, 0);
       else
         Offset = BB->getOffset();
 
diff --git a/bolt/lib/Profile/DataReader.cpp b/bolt/lib/Profile/DataReader.cpp
index 65e8e73..28b8629 100644
--- a/bolt/lib/Profile/DataReader.cpp
+++ b/bolt/lib/Profile/DataReader.cpp
@@ -721,7 +721,7 @@ bool DataReader::recordBranch(BinaryFunction &BF, uint64_t From, uint64_t To,
     const MCInst *LastInstr = ToBB->getLastNonPseudoInstr();
     if (LastInstr) {
       const uint32_t LastInstrOffset =
-          BC.MIB->getAnnotationWithDefault<uint32_t>(*LastInstr, "Offset");
+          BC.MIB->getOffsetWithDefault(*LastInstr, 0);
 
       // With old .fdata we are getting FT branches for "jcc,jmp" sequences.
       if (To == LastInstrOffset && BC.MIB->isUnconditionalBranch(*LastInstr))
diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
index 2b72de2..d8d6975 100644
--- a/bolt/lib/Profile/YAMLProfileWriter.cpp
+++ b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -53,10 +53,10 @@ void convert(const BinaryFunction &BF,
         continue;
 
       yaml::bolt::CallSiteInfo CSI;
-      auto Offset = BC.MIB->tryGetAnnotationAs<uint32_t>(Instr, "Offset");
-      if (!Offset || Offset.get() < BB->getInputOffset())
+      Optional<uint32_t> Offset = BC.MIB->getOffset(Instr);
+      if (!Offset || *Offset < BB->getInputOffset())
         continue;
-      CSI.Offset = Offset.get() - BB->getInputOffset();
+      CSI.Offset = *Offset - BB->getInputOffset();
 
       if (BC.MIB->isIndirectCall(Instr) || BC.MIB->isIndirectBranch(Instr)) {
         auto ICSP = BC.MIB->tryGetAnnotationAs<IndirectCallSiteProfile>(
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 9dc7a9e..73288a3 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -2771,7 +2771,6 @@ void RewriteInstance::buildFunctionsCFG() {
                      "Build Binary Functions", opts::TimeBuild);
 
   // Create annotation indices to allow lock-free execution
-  BC->MIB->getOrCreateAnnotationIndex("Offset");
   BC->MIB->getOrCreateAnnotationIndex("JTIndexReg");
   BC->MIB->getOrCreateAnnotationIndex("NOP");
   BC->MIB->getOrCreateAnnotationIndex("Size");
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 1cb0772..976348a 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -805,7 +805,7 @@ public:
 
   bool convertTailCallToJmp(MCInst &Inst) override {
     removeAnnotation(Inst, MCPlus::MCAnnotation::kTailCall);
-    removeAnnotation(Inst, "Offset");
+    clearOffset(Inst);
     if (getConditionalTailCall(Inst))
       unsetConditionalTailCall(Inst);
     return true;
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 21878c9..67efa16 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2020,7 +2020,7 @@ public:
 
     Inst.setOpcode(NewOpcode);
     removeAnnotation(Inst, MCPlus::MCAnnotation::kTailCall);
-    removeAnnotation(Inst, "Offset");
+    clearOffset(Inst);
     return true;
   }
 
@@ -3808,10 +3808,9 @@ public:
 
         if (CallOrJmp.getOpcode() == X86::CALL64r ||
             CallOrJmp.getOpcode() == X86::CALL64pcrel32) {
-          if (hasAnnotation(CallInst, "Offset"))
+          if (Optional<uint32_t> Offset = getOffset(CallInst))
             // Annotated as duplicated call
-            addAnnotation(CallOrJmp, "Offset",
-                          getAnnotationAs<uint32_t>(CallInst, "Offset"));
+            setOffset(CallOrJmp, *Offset);
         }
 
         if (isInvoke(CallInst) && !isInvoke(CallOrJmp)) {