[BOLT] Add support for safe-icf

bolt/docs/CommandLineArgumentReference.md

@@ -498,9 +498,12 @@
   Automatically put hot code on 2MB page(s) (hugify) at runtime. No manual call
   to hugify is needed in the binary (which is what --hot-text relies on).
 
-- `--icf`
+- `--icf=<value>`
 
   Fold functions with identical code
+  - `all`:  Enable identical code folding
+  - `none`: Disable identical code folding (default)
+  - `safe`: Enable safe identical code folding
 
 - `--icp`
 

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 function is referenced by non-control flow instruction.
+  bool HasAddressTaken{false};
+
   /// 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;
   }
 
+  /// Return true if function is referenced in non-control flow instructions..
+  bool hasAddressTaken() const { return HasAddressTaken; }
+
+  /// Set whether function is referenced in non-control flow instructions.
+  void setHasAddressTaken(bool Hat) { HasAddressTaken = Hat; }
+
   /// Returns the raw binary encoding of this function.
   ErrorOr<ArrayRef<uint8_t>> getData() const;
 
@@ -2094,6 +2103,10 @@ class BinaryFunction {
   // adjustments.
   void handleAArch64IndirectCall(MCInst &Instruction, const uint64_t Offset);
 
+  /// Processes code section to identify function references.
+  void processInstructionForFuncReferences(BinaryContext &BC,
+                                           const MCInst &Inst);
+
   /// Scan function for references to other functions. In relocation mode,
   /// add relocations for external references. In non-relocation mode, detect
   /// and mark new entry points.

bolt/include/bolt/Core/Relocation.h

@@ -16,6 +16,8 @@
 
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCStreamer.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/Object/ObjectFile.h"
 #include "llvm/TargetParser/Triple.h"
 
 namespace llvm {
@@ -178,6 +180,11 @@ inline raw_ostream &operator<<(raw_ostream &OS, const Relocation &Rel) {
   return OS;
 }
 
+uint32_t getRelocationSymbol(const object::ELFObjectFileBase *Obj,
+                             const object::RelocationRef &Rel);
+
+int64_t getRelocationAddend(const object::ELFObjectFileBase *Obj,
+                            const object::RelocationRef &Rel);
 } // namespace bolt
 } // namespace llvm
 

bolt/include/bolt/Passes/IdenticalCodeFolding.h

@@ -27,15 +27,42 @@ class IdenticalCodeFolding : public BinaryFunctionPass {
       return false;
     if (BF.hasSDTMarker())
       return false;
+    if (BF.hasAddressTaken())
+      return false;
     return BinaryFunctionPass::shouldOptimize(BF);
   }
 
 public:
+  enum class ICFLevel {
+    None, // No ICF. (Default)
+    Safe, // Safe ICF for all sections.
+    All,  // Aggressive ICF for code.
+  };
   explicit IdenticalCodeFolding(const cl::opt<bool> &PrintPass)
-      : BinaryFunctionPass(PrintPass) {}
+      : BinaryFunctionPass(PrintPass) {
+    VtableBitVector.resize((((uint64_t)1) << 32) / 8);
+  }
 
   const char *getName() const override { return "identical-code-folding"; }
   Error runOnFunctions(BinaryContext &BC) override;
+
+private:
+  /// Bit vector of memory addresses of vtables.
+  llvm::BitVector VtableBitVector;
+  bool isInVTable(uint64_t Address) const {
+    return VtableBitVector.test(Address / 8);
+  }
+  /// Scans symbol table and creates a bit vector of memory addresses of
+  /// vtables.
+  void processSymbolTable(const BinaryContext &BC);
+  /// Analyze .text section and relocations and mark functions that are not
+  /// safe to fold.
+  Error markFunctionsUnsafeToFold(BinaryContext &BC);
+  /// Process relocations in the .data section to identify function
+  /// references.
+  Error processDataRelocations(BinaryContext &BC,
+                               const SectionRef &SecRefRelData,
+                               const bool HasAddressTaken);
 };
 
 } // namespace bolt

bolt/lib/Core/BinaryFunction.cpp

@@ -1513,6 +1513,20 @@ MCSymbol *BinaryFunction::registerBranch(uint64_t Src, uint64_t Dst) {
   return Target;
 }
 
+void BinaryFunction::processInstructionForFuncReferences(BinaryContext &BC,
+                                                         const MCInst &Inst) {
+  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 = BC.getFunctionForSymbol(&Symbol))
+        BF->setHasAddressTaken(true);
+    }
+  }
+}
+
 bool BinaryFunction::scanExternalRefs() {
   bool Success = true;
   bool DisassemblyFailed = false;
@@ -1633,6 +1647,8 @@ bool BinaryFunction::scanExternalRefs() {
                              [](const MCOperand &Op) { return Op.isExpr(); })) {
       // Skip assembly if the instruction may not have any symbolic operands.
       continue;
+    } else {
+      processInstructionForFuncReferences(BC, Instruction);
     }
 
     // Emit the instruction using temp emitter and generate relocations.

bolt/lib/Core/Relocation.cpp

@@ -1115,3 +1115,65 @@ void Relocation::print(raw_ostream &OS) const {
     OS << ", 0x" << Twine::utohexstr(Addend);
   OS << ", 0x" << Twine::utohexstr(Value);
 }
+
+namespace {
+template <typename ELFT>
+int64_t getRelocationAddend(const llvm::object::ELFObjectFile<ELFT> *Obj,
+                            const llvm::object::RelocationRef &RelRef) {
+  using ELFShdrTy = typename ELFT::Shdr;
+  using Elf_Rela = typename ELFT::Rela;
+  int64_t Addend = 0;
+  const llvm::object::ELFFile<ELFT> &EF = Obj->getELFFile();
+  llvm::object::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 llvm::object::ELFObjectFile<ELFT> *Obj,
+                             const llvm::object::RelocationRef &RelRef) {
+  using ELFShdrTy = typename ELFT::Shdr;
+  uint32_t Symbol = 0;
+  const llvm::object::ELFFile<ELFT> &EF = Obj->getELFFile();
+  llvm::object::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;
+}
+} // namespace
+
+namespace llvm {
+namespace bolt {
+uint32_t getRelocationSymbol(const llvm::object::ELFObjectFileBase *Obj,
+                             const llvm::object::RelocationRef &Rel) {
+  return ::getRelocationSymbol(cast<llvm::object::ELF64LEObjectFile>(Obj), Rel);
+}
+
+int64_t getRelocationAddend(const llvm::object::ELFObjectFileBase *Obj,
+                            const llvm::object::RelocationRef &Rel) {
+  return ::getRelocationAddend(cast<llvm::object::ELF64LEObjectFile>(Obj), Rel);
+}
+} // namespace bolt
+} // namespace llvm

bolt/lib/Passes/IdenticalCodeFolding.cpp

@@ -13,6 +13,7 @@
 #include "bolt/Passes/IdenticalCodeFolding.h"
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ThreadPool.h"
@@ -42,6 +43,19 @@ TimeICF("time-icf",
   cl::ReallyHidden,
   cl::ZeroOrMore,
   cl::cat(BoltOptCategory));
+
+cl::opt<bolt::IdenticalCodeFolding::ICFLevel> ICF(
+    "icf", cl::desc("fold functions with identical code"),
+    cl::init(bolt::IdenticalCodeFolding::ICFLevel::None),
+    cl::values(clEnumValN(bolt::IdenticalCodeFolding::ICFLevel::All, "all",
+                          "Enable identical code folding"),
+               clEnumValN(bolt::IdenticalCodeFolding::ICFLevel::All, "",
+                          "Enable identical code folding"),
+               clEnumValN(bolt::IdenticalCodeFolding::ICFLevel::None, "none",
+                          "Disable identical code folding (default)"),
+               clEnumValN(bolt::IdenticalCodeFolding::ICFLevel::Safe, "safe",
+                          "Enable safe identical code folding")),
+    cl::ZeroOrMore, cl::ValueOptional, cl::cat(BoltOptCategory));
 } // namespace opts
 
 /// Compare two jump tables in 2 functions. The function relies on consistent
@@ -340,6 +354,113 @@ typedef std::unordered_map<BinaryFunction *, std::vector<BinaryFunction *>,
 
 namespace llvm {
 namespace bolt {
+/// Scans symbol table and creates a bit vector of memory addresses of vtables.
+void IdenticalCodeFolding::processSymbolTable(const BinaryContext &BC) {
+  for (auto &[Address, Data] : BC.getBinaryData()) {
+    // Filter out all symbols that are not vtables.
+    if (!Data->getName().starts_with("_ZTV"))
+      continue;
+    for (uint64_t I = Address / 8, End = I + (Data->getSize() / 8); I < End;
+         ++I)
+      VtableBitVector.set(I);
+  }
+}
+Error IdenticalCodeFolding::processDataRelocations(
+    BinaryContext &BC, const SectionRef &SecRefRelData,
+    const bool HasAddressTaken) {
+  for (const RelocationRef &Rel : SecRefRelData.relocations()) {
+    if (isInVTable(Rel.getOffset()))
+      continue;
+    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)
+      return createFatalBOLTError(
+          Twine("BOLT-ERROR: Only ELFObjectFileBase is supported"));
+    const int64_t Addend = getRelocationAddend(ELFObj, Rel);
+    BinaryFunction *BF = BC.getBinaryFunctionAtAddress(SymbolAddress + Addend);
+    if (!BF)
+      continue;
+    BF->setHasAddressTaken(HasAddressTaken);
+  }
+  return Error::success();
+}
+
+Error IdenticalCodeFolding::markFunctionsUnsafeToFold(BinaryContext &BC) {
+  if (!BC.isX86())
+    BC.outs() << "BOLT-WARNING: safe ICF is only supported for x86\n";
+  processSymbolTable(BC);
+  for (const auto &Sec : BC.sections()) {
+    if (!Sec.hasSectionRef() || !Sec.isRela())
+      continue;
+    const SectionRef &SecRef = Sec.getSectionRef();
+    section_iterator RelocatedSecIter = cantFail(SecRef.getRelocatedSection());
+    assert(RelocatedSecIter != SecRef.getObject()->section_end() &&
+           "Relocation section exists without corresponding relocated section");
+    const SectionRef &RelocatedSecRef = *RelocatedSecIter;
+    const StringRef RelocatedSectionName = cantFail(RelocatedSecRef.getName());
+    const bool SkipRelocs =
+        StringSwitch<bool>(RelocatedSectionName)
+            .Cases(".plt", ".got.plt", ".eh_frame", ".gcc_except_table",
+                   ".fini_array", ".init_array", true)
+            .Default(false);
+    if (!RelocatedSecRef.isData() || SkipRelocs)
+      continue;
+
+    Error ErrorStatus =
+        processDataRelocations(BC, SecRef, /* HasAddressTaken */ true);
+    if (ErrorStatus)
+      return ErrorStatus;
+  }
+  ErrorOr<BinarySection &> SecRelDataIA =
+      BC.getUniqueSectionByName(".rela.init_array");
+  if (SecRelDataIA) {
+    const SectionRef SecRefRelData = SecRelDataIA->getSectionRef();
+    Error ErrorStatus =
+        processDataRelocations(BC, SecRefRelData, /* !HasAddressTaken */ false);
+    if (ErrorStatus)
+      return ErrorStatus;
+  }
+  ErrorOr<BinarySection &> SecRelDataFIA =
+      BC.getUniqueSectionByName(".rela.fini_array");
+  if (SecRelDataFIA) {
+    const SectionRef SecRefRelData = SecRelDataFIA->getSectionRef();
+    Error ErrorStatus =
+        processDataRelocations(BC, SecRefRelData, /* !HasAddressTaken */ false);
+    if (ErrorStatus)
+      return ErrorStatus;
+  }
+
+  ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+    for (const BinaryBasicBlock *BB : BF.getLayout().blocks()) {
+      for (const MCInst &Inst : *BB) {
+        if (!(BC.MIB->isCall(Inst) || BC.MIB->isBranch(Inst)))
+          BF.processInstructionForFuncReferences(BC, Inst);
+      }
+    }
+  };
+  ParallelUtilities::PredicateTy SkipFunc =
+      [&](const BinaryFunction &BF) -> bool {
+    return BF.getState() != BinaryFunction::State::CFG;
+  };
+  ParallelUtilities::runOnEachFunction(
+      BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
+      SkipFunc, "markUnsafe", /*ForceSequential*/ false, 2);
+
+  LLVM_DEBUG({
+    for (auto &BFIter : BC.getBinaryFunctions()) {
+      if (!BFIter.second.hasAddressTaken())
+        continue;
+      dbgs() << "BOLT-DEBUG: skipping function " << BFIter.second.getOneName()
+             << '\n';
+    }
+  });
+  return Error::success();
+}
 
 Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
   const size_t OriginalFunctionCount = BC.getBinaryFunctions().size();
@@ -385,7 +506,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 +596,9 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
 
     LLVM_DEBUG(SinglePass.stopTimer());
   };
-
+  if (opts::ICF == ICFLevel::Safe)
+    if (Error Err = markFunctionsUnsafeToFold(BC))
+      return Err;
   hashFunctions();
   createCongruentBuckets();