safe-icf

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: https://phabricator.intern.facebook.com/D65799965

Author: Alexander Yermolovich

bolt/include/bolt/Core/BinaryContext.h

@@ -282,6 +282,12 @@ class BinaryContext {
                       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.
@@ -1350,6 +1356,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

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 setUnsetToICF() { IsSafeToICF = false; }
+
   /// Returns the raw binary encoding of this function.
   ErrorOr<ArrayRef<uint8_t>> getData() const;
 

bolt/include/bolt/Passes/IdenticalCodeFolding.h

@@ -31,11 +31,21 @@ class IdenticalCodeFolding : public BinaryFunctionPass {
   }
 
 public:
-  explicit IdenticalCodeFolding(const cl::opt<bool> &PrintPass)
-      : BinaryFunctionPass(PrintPass) {}
+  explicit IdenticalCodeFolding(const cl::opt<bool> &PrintPass, bool IsSafeICF)
+      : BinaryFunctionPass(PrintPass), IsSafeICF(IsSafeICF) {}
 
   const char *getName() const override { return "identical-code-folding"; }
   Error runOnFunctions(BinaryContext &BC) override;
+
+private:
+  /// Create a skip list of functions that should not be folded.
+  Error createFoldSkipList(BinaryContext &BC);
+  /// Processes relocations in the .data section to identify function
+  /// references.
+  void processDataRelocations(BinaryContext &BC,
+                              const SectionRef &SecRefRelData);
+
+  bool IsSafeICF;
 };
 
 } // namespace bolt

bolt/lib/Core/BinaryContext.cpp

@@ -75,8 +75,63 @@ 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<bool> ICF("icf", cl::desc("fold functions with identical code"),
+                  cl::cat(BoltOptCategory));
+
+cl::opt<bool> SafeICF("safe-icf",
+                      cl::desc("Enable safe identical code folding"),
+                      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 {
 
@@ -156,6 +211,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 +2010,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 (!opts::SafeICF || skipInstruction(Inst, *this))
+    return;
+  for (const MCOperand &Op : MCPlus::primeOperands(Inst)) {
+    if (Op.isExpr()) {
+      const MCExpr &Expr = *Op.getExpr();
+      if (Expr.getKind() == MCExpr::SymbolRef) {
+        const MCSymbol &Symbol = cast<MCSymbolRefExpr>(Expr).getSymbol();
+        if (BinaryFunction *BF = getFunctionForSymbol(&Symbol))
+          BF->setUnsetToICF();
+      }
+    }
+  }
+}
+
 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,75 @@ 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)
+      assert(false && "Only ELFObjectFileBase is supported");
+    const int64_t Addend = BinaryContext::getRelocationAddend(ELFObj, Rel);
+    BinaryFunction *BF = BC.getBinaryFunctionAtAddress(SymbolAddress + Addend);
+    if (!BF)
+      continue;
+    BF->setUnsetToICF();
+  }
+}
+
+Error IdenticalCodeFolding::createFoldSkipList(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({
+    std::vector<StringRef> Vect;
+    std::mutex PrintMutex;
+    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+      if (BF.isSafeToICF())
+        return;
+      std::lock_guard<std::mutex> Lock(PrintMutex);
+      Vect.push_back(BF.getOneName());
+    };
+    ParallelUtilities::PredicateTy SkipFunc =
+        [&](const BinaryFunction &BF) -> bool { return false; };
+    ParallelUtilities::runOnEachFunction(
+        BC, ParallelUtilities::SchedulingPolicy::SP_TRIVIAL, WorkFun, SkipFunc,
+        "markUnsafe", /*ForceSequential*/ false, 2);
+    llvm::sort(Vect);
+    for (const auto &FuncName : Vect)
+      dbgs() << "BOLT-DEBUG: skipping function " << FuncName << '\n';
+  });
+  return ErrorStatus;
+}
+
 Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
   const size_t OriginalFunctionCount = BC.getBinaryFunctions().size();
   uint64_t NumFunctionsFolded = 0;
@@ -350,6 +422,9 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
   std::atomic<uint64_t> NumFoldedLastIteration{0};
   CongruentBucketsMap CongruentBuckets;
 
+  auto SkipFuncShared = [&](const BinaryFunction &BF) {
+    return !shouldOptimize(BF) || !BF.isSafeToICF();
+  };
   // Hash all the functions
   auto hashFunctions = [&]() {
     NamedRegionTimer HashFunctionsTimer("hashing", "hashing", "ICF breakdown",
@@ -369,7 +444,7 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
     };
 
     ParallelUtilities::PredicateTy SkipFunc = [&](const BinaryFunction &BF) {
-      return !shouldOptimize(BF);
+      return SkipFuncShared(BF);
     };
 
     ParallelUtilities::runOnEachFunction(
@@ -385,7 +460,7 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
                                            "ICF breakdown", opts::TimeICF);
     for (auto &BFI : BC.getBinaryFunctions()) {
       BinaryFunction &BF = BFI.second;
-      if (!this->shouldOptimize(BF))
+      if (SkipFuncShared(BF))
         continue;
       CongruentBuckets[&BF].emplace(&BF);
     }
@@ -475,7 +550,11 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
 
     LLVM_DEBUG(SinglePass.stopTimer());
   };
-
+  if (IsSafeICF) {
+    if (Error Err = createFoldSkipList(BC)) {
+      return Err;
+    }
+  }
   hashFunctions();
   createCongruentBuckets();
 

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -54,6 +54,8 @@ extern cl::opt<bool> PrintDynoStats;
 extern cl::opt<bool> DumpDotAll;
 extern cl::opt<std::string> AsmDump;
 extern cl::opt<bolt::PLTCall::OptType> PLT;
+extern cl::opt<bool> ICF;
+extern cl::opt<bool> SafeICF;
 
 static cl::opt<bool>
 DynoStatsAll("dyno-stats-all",
@@ -65,9 +67,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 "
@@ -397,8 +396,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
     Manager.registerPass(std::make_unique<StripRepRet>(NeverPrint),
                          opts::StripRepRet);
 
-  Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+  Manager.registerPass(
+      std::make_unique<IdenticalCodeFolding>(PrintICF, opts::SafeICF),
+      opts::ICF || opts::SafeICF);
 
   Manager.registerPass(
       std::make_unique<SpecializeMemcpy1>(NeverPrint, opts::SpecializeMemcpy1),
@@ -422,8 +422,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
 
   Manager.registerPass(std::make_unique<Inliner>(PrintInline));
 
-  Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+  Manager.registerPass(
+      std::make_unique<IdenticalCodeFolding>(PrintICF, opts::SafeICF),
+      opts::ICF || opts::SafeICF);
 
   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",
@@ -698,7 +699,7 @@ void RewriteInstance::compare(RewriteInstance &RI2) {
 
   // Pre-pass ICF
   if (opts::ICF) {
-    IdenticalCodeFolding ICF(opts::NeverPrint);
+    IdenticalCodeFolding ICF(opts::NeverPrint, opts::SafeICF);
     outs() << "BOLT-DIFF: Starting ICF pass for binary 1";
     BC->logBOLTErrorsAndQuitOnFatal(ICF.runOnFunctions(*BC));
     outs() << "BOLT-DIFF: Starting ICF pass for binary 2";