[BOLT] Add support for safe-icf

bolt/include/bolt/Core/BinaryContext.h

@@ -276,12 +276,23 @@ class BinaryContext {
   void deregisterSectionName(const BinarySection &Section);
 
 public:
+  enum class ICFLevel {
+    None,
+    Safe,
+    All,
+  };
   static Expected<std::unique_ptr<BinaryContext>>
   createBinaryContext(Triple TheTriple, StringRef InputFileName,
                       SubtargetFeatures *Features, bool IsPIC,
                       std::unique_ptr<DWARFContext> DwCtx,
                       JournalingStreams Logger);
 
+  /// Returns addend of a relocation.
+  static int64_t getRelocationAddend(const ELFObjectFileBase *Obj,
+                                     const RelocationRef &Rel);
+  /// Returns symbol of a relocation.
+  static uint32_t getRelocationSymbol(const ELFObjectFileBase *Obj,
+                                      const RelocationRef &Rel);
   /// Superset of compiler units that will contain overwritten code that needs
   /// new debug info. In a few cases, functions may end up not being
   /// overwritten, but it is okay to re-generate debug info for them.
@@ -660,6 +671,9 @@ class BinaryContext {
 
   std::unique_ptr<MCAsmBackend> MAB;
 
+  /// ICF level to use for this binary.
+  ICFLevel CurrICFLevel{ICFLevel::None};
+
   /// Allows BOLT to print to log whenever it is necessary (with or without
   /// const references)
   mutable JournalingStreams Logger;
@@ -1350,6 +1364,9 @@ class BinaryContext {
     return Code.size();
   }
 
+  /// Processes .text section to identify function references.
+  void processInstructionForFuncReferences(const MCInst &Inst);
+
   /// Compute the native code size for a range of instructions.
   /// Note: this can be imprecise wrt the final binary since happening prior to
   /// relaxation, as well as wrt the original binary because of opcode
@@ -1476,6 +1493,8 @@ class BinaryContext {
     return *IOAddressMap;
   }
 
+  ICFLevel getICFLevel() const { return CurrICFLevel; }
+
   raw_ostream &outs() const { return Logger.Out; }
 
   raw_ostream &errs() const { return Logger.Err; }

bolt/include/bolt/Core/BinaryFunction.h

@@ -428,6 +428,9 @@ class BinaryFunction {
   /// Function order for streaming into the destination binary.
   uint32_t Index{-1U};
 
+  /// Indicates if the function is safe to fold.
+  bool IsSafeToICF{true};
+
   /// Get basic block index assuming it belongs to this function.
   unsigned getIndex(const BinaryBasicBlock *BB) const {
     assert(BB->getIndex() < BasicBlocks.size());
@@ -817,6 +820,12 @@ class BinaryFunction {
     return nullptr;
   }
 
+  /// Indicates if the function is safe to fold.
+  bool isSafeToICF() const { return IsSafeToICF; }
+
+  /// Sets the function is not safe to fold.
+  void setUnsafeICF() { IsSafeToICF = false; }
+
   /// Returns the raw binary encoding of this function.
   ErrorOr<ArrayRef<uint8_t>> getData() const;
 

bolt/include/bolt/Passes/IdenticalCodeFolding.h

@@ -27,6 +27,8 @@ class IdenticalCodeFolding : public BinaryFunctionPass {
       return false;
     if (BF.hasSDTMarker())
       return false;
+    if (!BF.isSafeToICF())
+      return false;
     return BinaryFunctionPass::shouldOptimize(BF);
   }
 
@@ -36,6 +38,15 @@ class IdenticalCodeFolding : public BinaryFunctionPass {
 
   const char *getName() const override { return "identical-code-folding"; }
   Error runOnFunctions(BinaryContext &BC) override;
+
+private:
+  /// Analyses .text section and relocations and marks functions that are not
+  /// safe to fold.
+  Error markFunctionsUnsafeToFold(BinaryContext &BC);
+  /// Processes relocations in the .data section to identify function
+  /// references.
+  void processDataRelocations(BinaryContext &BC,
+                              const SectionRef &SecRefRelData);
 };
 
 } // namespace bolt

bolt/lib/Core/BinaryContext.cpp

@@ -75,8 +75,67 @@ cl::opt<std::string> CompDirOverride(
              "location, which is used with DW_AT_dwo_name to construct a path "
              "to *.dwo files."),
     cl::Hidden, cl::init(""), cl::cat(BoltCategory));
+
+cl::opt<bolt::BinaryContext::ICFLevel> ICF(
+    "icf", cl::desc("fold functions with identical code"),
+    cl::init(bolt::BinaryContext::ICFLevel::None),
+    cl::values(
+        clEnumValN(bolt::BinaryContext::ICFLevel::All, "all", "enable ICF"),
+        clEnumValN(bolt::BinaryContext::ICFLevel::All, "", "enable ICF"),
+        clEnumValN(bolt::BinaryContext::ICFLevel::None, "none", "disable ICF"),
+        clEnumValN(bolt::BinaryContext::ICFLevel::Safe, "safe",
+                   "enable safe ICF")),
+    cl::ZeroOrMore, cl::ValueOptional, cl::cat(BoltOptCategory));
 } // namespace opts
 
+namespace {
+template <typename ELFT>
+int64_t getRelocationAddend(const ELFObjectFile<ELFT> *Obj,
+                            const RelocationRef &RelRef) {
+  using ELFShdrTy = typename ELFT::Shdr;
+  using Elf_Rela = typename ELFT::Rela;
+  int64_t Addend = 0;
+  const ELFFile<ELFT> &EF = Obj->getELFFile();
+  DataRefImpl Rel = RelRef.getRawDataRefImpl();
+  const ELFShdrTy *RelocationSection = cantFail(EF.getSection(Rel.d.a));
+  switch (RelocationSection->sh_type) {
+  default:
+    llvm_unreachable("unexpected relocation section type");
+  case ELF::SHT_REL:
+    break;
+  case ELF::SHT_RELA: {
+    const Elf_Rela *RelA = Obj->getRela(Rel);
+    Addend = RelA->r_addend;
+    break;
+  }
+  }
+
+  return Addend;
+}
+
+template <typename ELFT>
+uint32_t getRelocationSymbol(const ELFObjectFile<ELFT> *Obj,
+                             const RelocationRef &RelRef) {
+  using ELFShdrTy = typename ELFT::Shdr;
+  uint32_t Symbol = 0;
+  const ELFFile<ELFT> &EF = Obj->getELFFile();
+  DataRefImpl Rel = RelRef.getRawDataRefImpl();
+  const ELFShdrTy *RelocationSection = cantFail(EF.getSection(Rel.d.a));
+  switch (RelocationSection->sh_type) {
+  default:
+    llvm_unreachable("unexpected relocation section type");
+  case ELF::SHT_REL:
+    Symbol = Obj->getRel(Rel)->getSymbol(EF.isMips64EL());
+    break;
+  case ELF::SHT_RELA:
+    Symbol = Obj->getRela(Rel)->getSymbol(EF.isMips64EL());
+    break;
+  }
+
+  return Symbol;
+}
+} // anonymous namespace
+
 namespace llvm {
 namespace bolt {
 
@@ -144,6 +203,7 @@ BinaryContext::BinaryContext(std::unique_ptr<MCContext> Ctx,
       Logger(Logger), InitialDynoStats(isAArch64()) {
   RegularPageSize = isAArch64() ? RegularPageSizeAArch64 : RegularPageSizeX86;
   PageAlign = opts::NoHugePages ? RegularPageSize : HugePageSize;
+  CurrICFLevel = opts::ICF;
 }
 
 BinaryContext::~BinaryContext() {
@@ -156,6 +216,16 @@ BinaryContext::~BinaryContext() {
   clearBinaryData();
 }
 
+uint32_t BinaryContext::getRelocationSymbol(const ELFObjectFileBase *Obj,
+                                            const RelocationRef &Rel) {
+  return ::getRelocationSymbol(cast<ELF64LEObjectFile>(Obj), Rel);
+}
+
+int64_t BinaryContext::getRelocationAddend(const ELFObjectFileBase *Obj,
+                                           const RelocationRef &Rel) {
+  return ::getRelocationAddend(cast<ELF64LEObjectFile>(Obj), Rel);
+}
+
 /// Create BinaryContext for a given architecture \p ArchName and
 /// triple \p TripleName.
 Expected<std::unique_ptr<BinaryContext>> BinaryContext::createBinaryContext(
@@ -1945,6 +2015,27 @@ static void printDebugInfo(raw_ostream &OS, const MCInst &Instruction,
     OS << " discriminator:" << Row.Discriminator;
 }
 
+static bool skipInstruction(const MCInst &Inst, const BinaryContext &BC) {
+  const bool IsX86 = BC.isX86();
+  return (BC.MIB->isPseudo(Inst) || BC.MIB->isUnconditionalBranch(Inst) ||
+          (IsX86 && BC.MIB->isConditionalBranch(Inst)) ||
+          BC.MIB->isCall(Inst) || BC.MIB->isBranch(Inst));
+}
+void BinaryContext::processInstructionForFuncReferences(const MCInst &Inst) {
+  if (CurrICFLevel != ICFLevel::Safe || skipInstruction(Inst, *this))
+    return;
+  for (const MCOperand &Op : MCPlus::primeOperands(Inst)) {
+    if (!Op.isExpr())
+      continue;
+    const MCExpr &Expr = *Op.getExpr();
+    if (Expr.getKind() == MCExpr::SymbolRef) {
+      const MCSymbol &Symbol = cast<MCSymbolRefExpr>(Expr).getSymbol();
+      if (BinaryFunction *BF = getFunctionForSymbol(&Symbol))
+        BF->setUnsafeICF();
+    }
+  }
+}
+
 void BinaryContext::printInstruction(raw_ostream &OS, const MCInst &Instruction,
                                      uint64_t Offset,
                                      const BinaryFunction *Function,

bolt/lib/Core/BinaryFunction.cpp

@@ -1626,6 +1626,8 @@ bool BinaryFunction::scanExternalRefs() {
     if (!BC.HasRelocations)
       continue;
 
+    BC.processInstructionForFuncReferences(Instruction);
+
     if (BranchTargetSymbol) {
       BC.MIB->replaceBranchTarget(Instruction, BranchTargetSymbol,
                                   Emitter.LocalCtx.get());

bolt/lib/Passes/IdenticalCodeFolding.cpp

@@ -11,8 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/IdenticalCodeFolding.h"
+#include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ThreadPool.h"
@@ -31,6 +33,7 @@ using namespace bolt;
 namespace opts {
 
 extern cl::OptionCategory BoltOptCategory;
+extern cl::opt<unsigned> Verbosity;
 
 static cl::opt<bool>
     ICFUseDFS("icf-dfs", cl::desc("use DFS ordering when using -icf option"),
@@ -341,6 +344,65 @@ typedef std::unordered_map<BinaryFunction *, std::vector<BinaryFunction *>,
 namespace llvm {
 namespace bolt {
 
+void IdenticalCodeFolding::processDataRelocations(
+    BinaryContext &BC, const SectionRef &SecRefRelData) {
+  for (const RelocationRef &Rel : SecRefRelData.relocations()) {
+    symbol_iterator SymbolIter = Rel.getSymbol();
+    const ObjectFile *OwningObj = Rel.getObject();
+    assert(SymbolIter != OwningObj->symbol_end() &&
+           "relocation Symbol expected");
+    const SymbolRef &Symbol = *SymbolIter;
+    const uint64_t SymbolAddress = cantFail(Symbol.getAddress());
+    const ELFObjectFileBase *ELFObj = dyn_cast<ELFObjectFileBase>(OwningObj);
+    if (!ELFObj)
+      llvm_unreachable("Only ELFObjectFileBase is supported");
+    const int64_t Addend = BinaryContext::getRelocationAddend(ELFObj, Rel);
+    BinaryFunction *BF = BC.getBinaryFunctionAtAddress(SymbolAddress + Addend);
+    if (!BF)
+      continue;
+    BF->setUnsafeICF();
+  }
+}
+
+Error IdenticalCodeFolding::markFunctionsUnsafeToFold(BinaryContext &BC) {
+  Error ErrorStatus = Error::success();
+  ErrorOr<BinarySection &> SecRelData = BC.getUniqueSectionByName(".rela.data");
+  if (!BC.HasRelocations)
+    ErrorStatus = joinErrors(
+        std::move(ErrorStatus),
+        createFatalBOLTError(Twine("BOLT-ERROR: Binary built without "
+                                   "relocations. Safe ICF is not supported")));
+  if (ErrorStatus)
+    return ErrorStatus;
+  if (SecRelData) {
+    SectionRef SecRefRelData = SecRelData->getSectionRef();
+    processDataRelocations(BC, SecRefRelData);
+  }
+
+  ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+    if (BF.getState() == BinaryFunction::State::CFG) {
+      for (const BinaryBasicBlock *BB : BF.getLayout().blocks())
+        for (const MCInst &Inst : *BB)
+          BC.processInstructionForFuncReferences(Inst);
+    }
+  };
+  ParallelUtilities::PredicateTy SkipFunc =
+      [&](const BinaryFunction &BF) -> bool { return (bool)ErrorStatus; };
+  ParallelUtilities::runOnEachFunction(
+      BC, ParallelUtilities::SchedulingPolicy::SP_TRIVIAL, WorkFun, SkipFunc,
+      "markUnsafe", /*ForceSequential*/ false, 2);
+
+  LLVM_DEBUG({
+    for (auto &BFIter : BC.getBinaryFunctions()) {
+      if (BFIter.second.isSafeToICF())
+        continue;
+      dbgs() << "BOLT-DEBUG: skipping function " << BFIter.second.getOneName()
+             << '\n';
+    }
+  });
+  return ErrorStatus;
+}
+
 Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
   const size_t OriginalFunctionCount = BC.getBinaryFunctions().size();
   uint64_t NumFunctionsFolded = 0;
@@ -385,7 +447,7 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
                                            "ICF breakdown", opts::TimeICF);
     for (auto &BFI : BC.getBinaryFunctions()) {
       BinaryFunction &BF = BFI.second;
-      if (!this->shouldOptimize(BF))
+      if (!shouldOptimize(BF))
         continue;
       CongruentBuckets[&BF].emplace(&BF);
     }
@@ -475,7 +537,11 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
 
     LLVM_DEBUG(SinglePass.stopTimer());
   };
-
+  if (BC.getICFLevel() == BinaryContext::ICFLevel::Safe) {
+    if (Error Err = markFunctionsUnsafeToFold(BC)) {
+      return Err;
+    }
+  }
   hashFunctions();
   createCongruentBuckets();
 

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -65,9 +65,6 @@ static cl::opt<bool>
                          cl::desc("eliminate unreachable code"), cl::init(true),
                          cl::cat(BoltOptCategory));
 
-cl::opt<bool> ICF("icf", cl::desc("fold functions with identical code"),
-                  cl::cat(BoltOptCategory));
-
 static cl::opt<bool> JTFootprintReductionFlag(
     "jt-footprint-reduction",
     cl::desc("make jump tables size smaller at the cost of using more "
@@ -398,7 +395,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
                          opts::StripRepRet);
 
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+                       BC.getICFLevel() != BinaryContext::ICFLevel::None);
 
   Manager.registerPass(
       std::make_unique<SpecializeMemcpy1>(NeverPrint, opts::SpecializeMemcpy1),
@@ -423,7 +420,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(std::make_unique<Inliner>(PrintInline));
 
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+                       BC.getICFLevel() != BinaryContext::ICFLevel::None);
 
   Manager.registerPass(std::make_unique<PLTCall>(PrintPLT));
 

bolt/lib/Rewrite/BoltDiff.cpp

@@ -29,6 +29,7 @@ namespace opts {
 extern cl::OptionCategory BoltDiffCategory;
 extern cl::opt<bool> NeverPrint;
 extern cl::opt<bool> ICF;
+extern cl::opt<bool> SafeICF;
 
 static cl::opt<bool> IgnoreLTOSuffix(
     "ignore-lto-suffix",
@@ -697,7 +698,7 @@ void RewriteInstance::compare(RewriteInstance &RI2) {
   }
 
   // Pre-pass ICF
-  if (opts::ICF) {
+  if (BC->getICFLevel() != BinaryContext::ICFLevel::None) {
     IdenticalCodeFolding ICF(opts::NeverPrint);
     outs() << "BOLT-DIFF: Starting ICF pass for binary 1";
     BC->logBOLTErrorsAndQuitOnFatal(ICF.runOnFunctions(*BC));