[BOLT] Add support for safe-icf (#116275)

Identical Code Folding (ICF) folds functions that are identical into one
function, and updates symbol addresses to the new address. This reduces
the size of a binary, but can lead to problems. For example when
function pointers are compared. This can be done either explicitly in
the code or generated IR by optimization passes like Indirect Call
Promotion (ICP). After ICF what used to be two different addresses
become the same address. This can lead to a different code path being
taken.

This is where safe ICF comes in. Linker (LLD) does it using address
significant section generated by clang. If symbol is in it, or an object
doesn't have this section symbols are not folded.

BOLT does not have the information regarding which objects do not have
this section, so can't re-use this mechanism.

This implementation scans code section and conservatively marks
functions symbols as unsafe. It treats symbols as unsafe if they are
used in non-control flow instruction. It also scans through the data
relocation sections and does the same for relocations that reference a
function symbol. The latter handles the case when function pointer is
stored in a local or global variable, etc. If a relocation address
points within a vtable these symbols are skipped.

Author: ayermolo

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());
@@ -822,6 +825,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;
 
@@ -2135,6 +2146,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

@@ -1504,6 +1504,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;
@@ -1624,6 +1638,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 "
@@ -403,7 +403,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),
@@ -428,7 +428,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));