[Bolt] Teach bolt about no-return functions

Author: elhewaty

bolt/include/bolt/Core/BinaryContext.h

@@ -272,6 +272,10 @@ class BinaryContext {
   /// DWARF line info for CUs.
   std::map<unsigned, DwarfLineTable> DwarfLineTablesCUMap;
 
+  /// Container to cache results about functions that their paths lead to
+  /// a no-return function.
+  std::unordered_map<BinaryFunction *, bool> CallsNoReturnFunction;
+
   /// Internal helper for removing section name from a lookup table.
   void deregisterSectionName(const BinarySection &Section);
 
@@ -282,6 +286,16 @@ class BinaryContext {
                       std::unique_ptr<DWARFContext> DwCtx,
                       JournalingStreams Logger);
 
+  void setHasPathToNoReturn(BinaryFunction *Func, bool value = true) {
+    CallsNoReturnFunction[Func] = value;
+  }
+
+  bool cachedInNoReturnMap(BinaryFunction *Func) {
+    return CallsNoReturnFunction.find(Func) != CallsNoReturnFunction.end();
+  }
+
+  bool hasPathToNoReturn(BinaryFunction *Func);
+
   /// 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.

bolt/include/bolt/Passes/DiscoverNoReturnPass.h

@@ -0,0 +1,36 @@
+//===- bolt/Passes/Discover.h -----------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_PASSES_DISCOVERNORETURN_H
+#define BOLT_PASSES_DISCOVERNORETURN_H
+
+#include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Passes/BinaryPasses.h"
+#include <unordered_map>
+
+namespace llvm {
+
+namespace bolt {
+
+class DiscoverNoReturnPass : public BinaryFunctionPass {
+public:
+  explicit DiscoverNoReturnPass() : BinaryFunctionPass(false) {}
+
+  const char *getName() const override { return "discover-no-return"; }
+
+  Error runOnFunctions(BinaryContext &BC) override;
+
+private:
+  std::unordered_map<BinaryFunction *, bool> Visited;
+  bool traverseFromFunction(BinaryFunction *Func, BinaryContext &BC);
+};
+} // namespace bolt
+} // namespace llvm
+
+#endif

bolt/lib/Core/BinaryContext.cpp

@@ -308,6 +308,20 @@ Expected<std::unique_ptr<BinaryContext>> BinaryContext::createBinaryContext(
   return std::move(BC);
 }
 
+bool BinaryContext::hasPathToNoReturn(BinaryFunction *Func) {
+  // Dummy way to mark no-return functions.
+  // FIXME: Find a better way.
+  if (std::string FuncName = Func->getPrintName();
+      FuncName == "__cxa_throw@PLT" || FuncName != "_Unwind_Resume@PLT" ||
+      FuncName == "__cxa_rethrow@PLT" || FuncName != "exit@PLT" ||
+      FuncName == "abort@PLT" || FuncName == "setjmp@PLT" ||
+      FuncName == "longjmp@PLT")
+    return true;
+
+  auto itr = CallsNoReturnFunction.find(Func);
+  return itr != CallsNoReturnFunction.end() && itr->second;
+}
+
 bool BinaryContext::forceSymbolRelocations(StringRef SymbolName) const {
   if (opts::HotText &&
       (SymbolName == "__hot_start" || SymbolName == "__hot_end"))

bolt/lib/Passes/CMakeLists.txt

@@ -8,6 +8,7 @@ add_llvm_library(LLVMBOLTPasses
   CacheMetrics.cpp
   DataflowAnalysis.cpp
   DataflowInfoManager.cpp
+  DiscoverNoReturnPass.cpp
   FrameAnalysis.cpp
   FrameOptimizer.cpp
   FixRelaxationPass.cpp

bolt/lib/Passes/DiscoverNoReturnPass.cpp

@@ -0,0 +1,84 @@
+//===- bolt/Passes/ReorderSection.cpp - Reordering of section data --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements DiscoverNoReturnPass class.
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Passes/DiscoverNoReturnPass.h"
+
+using namespace llvm;
+
+namespace opts {
+extern cl::OptionCategory BoltCategory;
+
+static cl::opt<bool> DiscoverNoReturnAnalysis(
+    "discover-no-return",
+    cl::desc("analyze the binary and mark no-return functions"), cl::init(true),
+    cl::cat(BoltCategory), cl::ReallyHidden);
+} // namespace opts
+
+namespace llvm {
+namespace bolt {
+
+Error DiscoverNoReturnPass::runOnFunctions(BinaryContext &BC) {
+  bool Changed;
+  do {
+    Changed = false;
+    for (auto &BFI : BC.getBinaryFunctions()) {
+      auto &Func = BFI.second;
+      bool PrevStat = BC.hasPathToNoReturn(&Func);
+      bool CurStat = traverseFromFunction(&Func, BC);
+      Changed |= (PrevStat != CurStat);
+    }
+  } while (Changed);
+
+  return Error::success();
+}
+
+bool DiscoverNoReturnPass::traverseFromFunction(BinaryFunction *Func,
+                                                BinaryContext &BC) {
+  // The Function cached before, so return its value
+  if (BC.cachedInNoReturnMap(Func))
+    return BC.hasPathToNoReturn(Func);
+
+  Visited[Func] = true;
+  bool Result = true;
+  bool hasCalls = 0;
+  bool hasReturns = 0;
+  for (auto &BB : *Func) {
+    if (!BB.getNumCalls())
+      continue;
+    for (auto &Inst : BB) {
+      if (BC.MIB->isCall(Inst)) {
+        hasCalls = true;
+        const MCSymbol *TargetSymbol = BC.MIB->getTargetSymbol(Inst);
+        BinaryFunction *TargetFunction = BC.SymbolToFunctionMap[TargetSymbol];
+        if (!Visited.count(TargetFunction))
+          Result &= traverseFromFunction(TargetFunction, BC);
+      }
+      hasReturns |= BC.MIB->isReturn(Inst);
+    }
+  }
+
+  // This functions is represented as a leaf in the call graph and doesn't
+  // have a no-return attribute.
+  if (!hasCalls && hasReturns)
+    Result = false;
+
+  // If the function doens't have a return instruction then it's a
+  // no-return function.
+  if (!hasReturns)
+    Result = true;
+
+  BC.setHasPathToNoReturn(Func, Result);
+  return Result;
+}
+
+} // end namespace bolt
+} // end namespace llvm

bolt/lib/Rewrite/BinaryPassManager.cpp

@@ -13,6 +13,7 @@
 #include "bolt/Passes/AsmDump.h"
 #include "bolt/Passes/CMOVConversion.h"
 #include "bolt/Passes/ContinuityStats.h"
+#include "bolt/Passes/DiscoverNoReturnPass.h"
 #include "bolt/Passes/FixRISCVCallsPass.h"
 #include "bolt/Passes/FixRelaxationPass.h"
 #include "bolt/Passes/FrameOptimizer.h"
@@ -443,6 +444,7 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
 
   Manager.registerPass(std::make_unique<CMOVConversion>(),
                        opts::CMOVConversionFlag);
+  Manager.registerPass(std::make_unique<DiscoverNoReturnPass>());
 
   // This pass syncs local branches with CFG. If any of the following
   // passes breaks the sync - they either need to re-run the pass or