[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.

1 files changed, 105 insertions, 2 deletions

diff --git a/bolt/lib/Passes/IdenticalCodeFolding.cpp b/bolt/lib/Passes/IdenticalCodeFolding.cpp
index 38e080c..8923562 100644
--- a/bolt/lib/Passes/IdenticalCodeFolding.cpp
+++ b/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();