Merge pull request #144 from sx-aurora-dev/feature/merge-upstream-20220120

Feature/merge upstream 20220120

Author: kaz7

bolt/include/bolt/Core/MCPlus.h

@@ -65,6 +65,7 @@ class MCAnnotation {
     kJumpTable,           /// Jump Table.
     kTailCall,            /// Tail call.
     kConditionalTailCall, /// CTC.
+    kOffset,              /// Offset in the function.
     kGeneric              /// First generic annotation.
   };
 

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -78,7 +78,7 @@ class MCPlusBuilder {
     assert(Index < 256 && "annotation index max value exceeded");
     assert((Value == (Value << 8) >> 8) && "annotation value out of range");
 
-    Value &= 0xffffffffffffff;
+    Value &= 0xff'ffff'ffff'ffff;
     Value |= (int64_t)Index << 56;
 
     return Value;
@@ -91,7 +91,7 @@ class MCPlusBuilder {
 
   /// Extract annotation value from immediate operand value.
   static int64_t extractAnnotationValue(int64_t ImmValue) {
-    ImmValue &= 0xffffffffffffff;
+    ImmValue &= 0xff'ffff'ffff'ffff;
     return (ImmValue << 8) >> 8;
   }
 
@@ -1075,6 +1075,18 @@ class MCPlusBuilder {
   /// 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.

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

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

bolt/lib/Core/BinaryFunction.cpp

@@ -1406,7 +1406,7 @@ bool BinaryFunction::disassemble() {
 
     // 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
@@ -2033,7 +2031,8 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
       assert(PrevInstr && "no previous instruction for a fall through");
       if (MIB->isUnconditionalBranch(Instr) &&
           !MIB->isUnconditionalBranch(*PrevInstr) &&
-          !MIB->getConditionalTailCall(*PrevInstr)) {
+          !MIB->getConditionalTailCall(*PrevInstr) &&
+          !MIB->isReturn(*PrevInstr)) {
         // Temporarily restore inserter basic block.
         InsertBB = PrevBB;
       } else {
@@ -2216,7 +2215,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 +2232,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 +4323,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;
     }
 

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)

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 {

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

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

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