[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};
 
+  /// 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);
 
+  /// Analyze instruction to identify a function references.
+  void analyzeInstructionForFuncReference(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/Passes/IdenticalCodeFolding.h

@@ -27,15 +27,53 @@ 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) {}
 
   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;
+  /// Returns true if the memory address is in the bit vector of vtables.
+  bool isAddressInVTable(uint64_t Address) const {
+    return VtableBitVector.test(Address / 8);
+  }
+  /// Initialize bit vector of memory addresses of vtables.
+  void initVtable() {
+    constexpr uint64_t PotentialRelocationAddressRange =
+        (((uint64_t)1) << 32) / 8;
+    VtableBitVector.resize(PotentialRelocationAddressRange);
+  }
+  /// Mark memory address of vtable as used.
+  void setAddressUsedInVTable(uint64_t Address) {
+    VtableBitVector.set(Address / 8);
+  }
+  /// Scans symbol table and creates a bit vector of memory addresses of
+  /// vtables.
+  void initVTableReferences(const BinaryContext &BC);
+  /// Analyze .text section and relocations and mark functions that are not
+  /// safe to fold.
+  void markFunctionsUnsafeToFold(BinaryContext &BC);
+  /// Process static and dynamic relocations in the data sections to identify
+  /// function references, and marks them as unsafe to fold. It filters out
+  /// symbol references that are in vtables.
+  void analyzeDataRelocations(BinaryContext &BC);
+  /// Process functions that have CFG created and mark functions unsafe to fold
+  /// that are used in non-control flow instructions.
+  void analyzeFunctions(BinaryContext &BC);
 };
 
 } // namespace bolt

bolt/lib/Core/BinaryFunction.cpp

@@ -1513,6 +1513,20 @@ MCSymbol *BinaryFunction::registerBranch(uint64_t Src, uint64_t Dst) {
   return Target;
 }
 
+void BinaryFunction::analyzeInstructionForFuncReference(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)
+      continue;
+    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 {
+      analyzeInstructionForFuncReference(BC, Instruction);
     }
 
     // Emit the instruction using temp emitter and generate relocations.

bolt/lib/Passes/IdenticalCodeFolding.cpp

@@ -42,6 +42,21 @@ 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, "1",
+                          "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 +355,77 @@ typedef std::unordered_map<BinaryFunction *, std::vector<BinaryFunction *>,
 
 namespace llvm {
 namespace bolt {
+void IdenticalCodeFolding::initVTableReferences(const BinaryContext &BC) {
+  initVtable();
+  for (const 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, End = I + Data->getSize(); I < End; I += 8)
+      setAddressUsedInVTable(I);
+  }
+}
+void IdenticalCodeFolding::analyzeDataRelocations(BinaryContext &BC) {
+  initVTableReferences(BC);
+  // For static relocations there should be a symbol for function references.
+  for (const BinarySection &Sec : BC.sections()) {
+    if (!Sec.hasSectionRef() || !Sec.isData())
+      continue;
+    for (const auto &Rel : Sec.relocations()) {
+      const uint64_t RelAddr = Rel.Offset + Sec.getAddress();
+      if (isAddressInVTable(RelAddr))
+        continue;
+      if (BinaryFunction *BF = BC.getFunctionForSymbol(Rel.Symbol))
+        BF->setHasAddressTaken(true);
+    }
+    // For dyanmic relocations there are two cases:
+    // 1: No symbol and only addend.
+    // 2: There is symbol, but it references undefined symbol, or things like
+    // type information. As the result only using addend to lookup BF is a valid
+    // case.
+    for (const auto &Rel : Sec.dynamicRelocations()) {
+      const uint64_t RelAddr = Rel.Offset + Sec.getAddress();
+      if (isAddressInVTable(RelAddr))
+        continue;
+      if (BinaryFunction *BF =
+              BC.getBinaryFunctionContainingAddress(Rel.Addend,
+                                                    /*CheckPastEnd*/ false,
+                                                    /*UseMaxSize*/ true))
+        BF->setHasAddressTaken(true);
+    }
+  }
+}
+void IdenticalCodeFolding::analyzeFunctions(BinaryContext &BC) {
+  ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+    for (const BinaryBasicBlock &BB : BF)
+      for (const MCInst &Inst : BB)
+        if (!(BC.MIB->isCall(Inst) || BC.MIB->isBranch(Inst)))
+          BF.analyzeInstructionForFuncReference(BC, Inst);
+  };
+  ParallelUtilities::PredicateTy SkipFunc =
+      [&](const BinaryFunction &BF) -> bool { return !BF.hasCFG(); };
+  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 with reference taken "
+             << BFIter.second.getOneName() << '\n';
+    }
+  });
+}
+void IdenticalCodeFolding::markFunctionsUnsafeToFold(BinaryContext &BC) {
+  NamedRegionTimer MarkFunctionsUnsafeToFoldTimer(
+      "markFunctionsUnsafeToFold", "markFunctionsUnsafeToFold", "ICF breakdown",
+      "ICF breakdown", opts::TimeICF);
+  if (!BC.isX86())
+    BC.outs() << "BOLT-WARNING: safe ICF is only supported for x86\n";
+  analyzeDataRelocations(BC);
+  analyzeFunctions(BC);
+}
 
 Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
   const size_t OriginalFunctionCount = BC.getBinaryFunctions().size();
@@ -385,7 +471,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 +561,8 @@ Error IdenticalCodeFolding::runOnFunctions(BinaryContext &BC) {
 
     LLVM_DEBUG(SinglePass.stopTimer());
   };
-
+  if (opts::ICF == ICFLevel::Safe)
+    markFunctionsUnsafeToFold(BC);
   hashFunctions();
   createCongruentBuckets();
 

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -54,6 +54,7 @@ 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<bolt::IdenticalCodeFolding::ICFLevel> ICF;
 
 static cl::opt<bool>
 DynoStatsAll("dyno-stats-all",
@@ -65,9 +66,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 +396,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
                          opts::StripRepRet);
 
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+                       opts::ICF != IdenticalCodeFolding::ICFLevel::None);
 
   Manager.registerPass(
       std::make_unique<SpecializeMemcpy1>(NeverPrint, opts::SpecializeMemcpy1),
@@ -423,7 +421,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
   Manager.registerPass(std::make_unique<Inliner>(PrintInline));
 
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
-                       opts::ICF);
+                       opts::ICF != IdenticalCodeFolding::ICFLevel::None);
 
   Manager.registerPass(std::make_unique<PLTCall>(PrintPLT));
 

bolt/lib/Rewrite/BoltDiff.cpp

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

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -19,6 +19,7 @@
 #include "bolt/Core/Relocation.h"
 #include "bolt/Passes/BinaryPasses.h"
 #include "bolt/Passes/CacheMetrics.h"
+#include "bolt/Passes/IdenticalCodeFolding.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Profile/BoltAddressTranslation.h"
 #include "bolt/Profile/DataAggregator.h"
@@ -86,6 +87,7 @@ extern cl::opt<bolt::ReorderFunctions::ReorderType> ReorderFunctions;
 extern cl::opt<bool> TerminalTrap;
 extern cl::opt<bool> TimeBuild;
 extern cl::opt<bool> TimeRewrite;
+extern cl::opt<bolt::IdenticalCodeFolding::ICFLevel> ICF;
 
 cl::opt<bool> AllowStripped("allow-stripped",
                             cl::desc("allow processing of stripped binaries"),
@@ -2050,6 +2052,13 @@ void RewriteInstance::adjustCommandLineOptions() {
     exit(1);
   }
 
+  if (!BC->HasRelocations &&
+      opts::ICF == IdenticalCodeFolding::ICFLevel::Safe) {
+    BC->errs() << "BOLT-ERROR: binary built without relocations. Safe ICF is "
+                  "not supported\n";
+    exit(1);
+  }
+
   if (opts::Instrument ||
       (opts::ReorderFunctions != ReorderFunctions::RT_NONE &&
        !opts::HotText.getNumOccurrences())) {