[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 a 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,14 @@ class BinaryFunction {
     return nullptr;
   }
 
+  /// Return true if function is referenced in a non-control flow instruction.
+  /// This flag is set when the code and relocation analyses are being
+  /// performed, which occurs when safe ICF (Identical Code Folding) is enabled.
+  bool hasAddressTaken() const { return HasAddressTaken; }
+
+  /// Set whether function is referenced in a non-control flow instruction.
+  void setHasAddressTaken(bool AddressTaken) { HasAddressTaken = AddressTaken; }
+
   /// Returns the raw binary encoding of this function.
   ErrorOr<ArrayRef<uint8_t>> getData() const;
 
@@ -2094,6 +2105,9 @@ class BinaryFunction {
   // adjustments.
   void handleAArch64IndirectCall(MCInst &Instruction, const uint64_t Offset);
 
+  /// Analyze instruction to identify a function reference.
+  void analyzeInstructionForFuncReference(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

@@ -11,6 +11,7 @@
 
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Passes/BinaryPasses.h"
+#include "llvm/ADT/SparseBitVector.h"
 
 namespace llvm {
 namespace bolt {
@@ -20,22 +21,72 @@ namespace bolt {
 ///
 class IdenticalCodeFolding : public BinaryFunctionPass {
 protected:
-  bool shouldOptimize(const BinaryFunction &BF) const override {
-    if (BF.hasUnknownControlFlow())
-      return false;
-    if (BF.isFolded())
-      return false;
-    if (BF.hasSDTMarker())
-      return false;
-    return BinaryFunctionPass::shouldOptimize(BF);
-  }
+  /// Return true if the function is safe to fold.
+  bool shouldOptimize(const BinaryFunction &BF) const override;
 
 public:
+  enum class ICFLevel {
+    None, /// No ICF. (Default)
+    Safe, /// Safe ICF.
+    All,  /// Aggressive ICF.
+  };
   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::SparseBitVector<> VTableBitVector;
+
+  /// Return true if the memory address is in a vtable.
+  bool isAddressInVTable(uint64_t Address) const {
+    return VTableBitVector.test(Address / 8);
+  }
+
+  /// Mark memory address of a vtable as used.
+  void setAddressUsedInVTable(uint64_t Address) {
+    VTableBitVector.set(Address / 8);
+  }
+
+  /// Scan symbol table and mark memory addresses of
+  /// vtables.
+  void initVTableReferences(const BinaryContext &BC);
+
+  /// Analyze code 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 mark them as unsafe to fold. It filters out
+  /// symbol references that are in vtables.
+  void analyzeDataRelocations(BinaryContext &BC);
+
+  /// Process functions that have been disassembled and mark functions that are
+  /// used in non-control flow instructions as unsafe to fold.
+  void analyzeFunctions(BinaryContext &BC);
+};
+
+class DeprecatedICFNumericOptionParser
+    : public cl::parser<IdenticalCodeFolding::ICFLevel> {
+public:
+  explicit DeprecatedICFNumericOptionParser(cl::Option &O)
+      : cl::parser<IdenticalCodeFolding::ICFLevel>(O) {}
+
+  bool parse(cl::Option &O, StringRef ArgName, StringRef Arg,
+             IdenticalCodeFolding::ICFLevel &Value) {
+    if (Arg == "0" || Arg == "1") {
+      Value = (Arg == "0") ? IdenticalCodeFolding::ICFLevel::None
+                           : IdenticalCodeFolding::ICFLevel::All;
+      errs() << formatv("BOLT-WARNING: specifying numeric value \"{0}\" "
+                        "for option -{1} is deprecated\n",
+                        Arg, ArgName);
+      return false;
+    }
+    return cl::parser<IdenticalCodeFolding::ICFLevel>::parse(O, ArgName, Arg,
+                                                             Value);
+  }
 };
 
 } // 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(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();
+    // Set HasAddressTaken for a function regardless of the ICF level.
+    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(Instruction);
     }
 
     // Emit the instruction using temp emitter and generate relocations.

bolt/lib/Passes/IdenticalCodeFolding.cpp

@@ -15,6 +15,7 @@
 #include "bolt/Core/ParallelUtilities.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/Timer.h"
 #include <atomic>
@@ -42,8 +43,41 @@ TimeICF("time-icf",
   cl::ReallyHidden,
   cl::ZeroOrMore,
   cl::cat(BoltOptCategory));
+
+cl::opt<bolt::IdenticalCodeFolding::ICFLevel, false,
+        DeprecatedICFNumericOptionParser>
+    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::None, "0",
+                              "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
 
+bool IdenticalCodeFolding::shouldOptimize(const BinaryFunction &BF) const {
+  if (BF.hasUnknownControlFlow())
+    return false;
+  if (BF.isFolded())
+    return false;
+  if (BF.hasSDTMarker())
+    return false;
+  if (BF.isPseudo())
+    return false;
+  if (opts::ICF == ICFLevel::Safe && BF.hasAddressTaken())
+    return false;
+  return BinaryFunctionPass::shouldOptimize(BF);
+}
+
 /// Compare two jump tables in 2 functions. The function relies on consistent
 /// ordering of basic blocks in both binary functions (e.g. DFS).
 static bool equalJumpTables(const JumpTable &JumpTableA,
@@ -340,6 +374,74 @@ typedef std::unordered_map<BinaryFunction *, std::vector<BinaryFunction *>,
 
 namespace llvm {
 namespace bolt {
+void IdenticalCodeFolding::initVTableReferences(const BinaryContext &BC) {
+  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 dynamic relocations there are two cases:
+    // 1: No symbol and only addend.
+    // 2: There is a symbol, but it does not references a function in a binary.
+    for (const auto &Rel : Sec.dynamicRelocations()) {
+      const uint64_t RelAddr = Rel.Offset + Sec.getAddress();
+      if (isAddressInVTable(RelAddr))
+        continue;
+      if (BinaryFunction *BF = BC.getBinaryFunctionAtAddress(Rel.Addend))
+        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(Inst);
+  };
+  ParallelUtilities::PredicateTy SkipFunc =
+      [&](const BinaryFunction &BF) -> bool { return !BF.hasCFG(); };
+  ParallelUtilities::runOnEachFunction(
+      BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
+      SkipFunc, "markUnsafe");
+
+  LLVM_DEBUG({
+    for (const 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 +487,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 +577,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,9 @@ 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, false,
+               llvm::bolt::DeprecatedICFNumericOptionParser>
+    ICF;
 
 static cl::opt<bool>
 DynoStatsAll("dyno-stats-all",
@@ -65,9 +68,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 +398,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 +423,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,9 @@ 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, false,
+               llvm::bolt::DeprecatedICFNumericOptionParser>
+    ICF;
 
 static cl::opt<bool> IgnoreLTOSuffix(
     "ignore-lto-suffix",
@@ -697,7 +699,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));