[BOLT] Support restartable sequences in tcmalloc

bolt/include/bolt/Rewrite/MetadataRewriters.h

@@ -19,12 +19,14 @@ class BinaryContext;
 
 // The list of rewriter build functions.
 
-std::unique_ptr<MetadataRewriter> createLinuxKernelRewriter(BinaryContext &);
-
 std::unique_ptr<MetadataRewriter> createBuildIDRewriter(BinaryContext &);
 
+std::unique_ptr<MetadataRewriter> createLinuxKernelRewriter(BinaryContext &);
+
 std::unique_ptr<MetadataRewriter> createPseudoProbeRewriter(BinaryContext &);
 
+std::unique_ptr<MetadataRewriter> createRSeqRewriter(BinaryContext &);
+
 std::unique_ptr<MetadataRewriter> createSDTRewriter(BinaryContext &);
 
 std::unique_ptr<MetadataRewriter> createGNUPropertyRewriter(BinaryContext &);

bolt/lib/Rewrite/CMakeLists.txt

@@ -24,6 +24,7 @@ add_llvm_library(LLVMBOLTRewrite
   BuildIDRewriter.cpp
   PseudoProbeRewriter.cpp
   RewriteInstance.cpp
+  RSeqRewriter.cpp
   SDTRewriter.cpp
   GNUPropertyRewriter.cpp
 

bolt/lib/Rewrite/RSeqRewriter.cpp

@@ -0,0 +1,72 @@
+//===- bolt/Rewrite/RSeqRewriter.cpp --------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Basic support for restartable sequences used by tcmalloc. Prevent critical
+// section overrides by ignoring optimizations in containing functions.
+//
+// References:
+//   * https://google.github.io/tcmalloc/rseq.html
+//   * tcmalloc/internal/percpu_rseq_x86_64.S
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Rewrite/MetadataRewriter.h"
+#include "bolt/Rewrite/MetadataRewriters.h"
+#include "llvm/Support/Errc.h"
+
+using namespace llvm;
+using namespace bolt;
+
+namespace {
+
+class RSeqRewriter final : public MetadataRewriter {
+public:
+  RSeqRewriter(StringRef Name, BinaryContext &BC)
+      : MetadataRewriter(Name, BC) {}
+
+  Error preCFGInitializer() override {
+    for (const BinarySection &Section : BC.allocatableSections()) {
+      if (Section.getName() != "__rseq_cs")
+        continue;
+
+      auto handleRelocation = [&](const Relocation &Rel, bool IsDynamic) {
+        BinaryFunction *BF = nullptr;
+        if (Rel.Symbol)
+          BF = BC.getFunctionForSymbol(Rel.Symbol);
+        else if (Relocation::isRelative(Rel.Type))
+          BF = BC.getBinaryFunctionContainingAddress(Rel.Addend);
+
+        if (!BF) {
+          BC.errs() << "BOLT-WARNING: no function found matching "
+                    << (IsDynamic ? "dynamic " : "")
+                    << "relocation in __rseq_cs\n";
+        } else if (!BF->isIgnored()) {
+          BC.outs() << "BOLT-INFO: restartable sequence reference detected in "
+                    << *BF << ". Function will not be optimized\n";
+          BF->setIgnored();
+        }
+      };
+
+      for (const Relocation &Rel : Section.dynamicRelocations())
+        handleRelocation(Rel, /*IsDynamic*/ true);
+
+      for (const Relocation &Rel : Section.relocations())
+        handleRelocation(Rel, /*IsDynamic*/ false);
+    }
+
+    return Error::success();
+  }
+};
+
+} // namespace
+
+std::unique_ptr<MetadataRewriter>
+llvm::bolt::createRSeqRewriter(BinaryContext &BC) {
+  return std::make_unique<RSeqRewriter>("rseq-cs-rewriter", BC);
+}

bolt/lib/Rewrite/RewriteInstance.cpp

@@ -3346,6 +3346,8 @@ void RewriteInstance::initializeMetadataManager() {
 
   MetadataManager.registerRewriter(createPseudoProbeRewriter(*BC));
 
+  MetadataManager.registerRewriter(createRSeqRewriter(*BC));
+
   MetadataManager.registerRewriter(createSDTRewriter(*BC));
 
   MetadataManager.registerRewriter(createGNUPropertyRewriter(*BC));

bolt/test/X86/rseq.s

@@ -0,0 +1,38 @@
+## Check that llvm-bolt avoids optimization of functions referenced from
+## __rseq_cs section, i.e. containing critical sections and abort handlers used
+## by restartable sequences in tcmalloc.
+
+# RUN: %clang %cflags %s -o %t -nostdlib -no-pie -Wl,-q
+# RUN: llvm-bolt %t -o %t.bolt --print-cfg 2>&1 | FileCheck %s
+# RUN: %clang %cflags %s -o %t.pie -nostdlib -pie -Wl,-q
+# RUN: llvm-bolt %t.pie -o %t.pie.bolt 2>&1 | FileCheck %s
+
+# CHECK: restartable sequence reference detected in _start
+# CHECK: restartable sequence reference detected in __rseq_abort
+
+## Force relocations against .text
+  .text
+.reloc 0, R_X86_64_NONE
+
+  .global _start
+  .type _start, %function
+_start:
+  pushq %rbp
+  mov %rsp, %rbp
+.L1:
+  pop %rbp
+.L2:
+  retq
+  .size _start, .-_start
+
+  .section __rseq_abort, "ax"
+## Signature for rseq abort IP. Unmarked in the symbol table.
+  .byte 0x0f, 0x1f, 0x05
+  .long 0x42424242
+.L3:
+  jmp .L2
+
+.section __rseq_cs, "aw"
+.balign 32
+  .quad .L1
+  .quad .L3