[BOLT] Support restartable sequences in tcmalloc (#167195)

Add `RSeqRewriter` to detect code references from `__rseq_cs` section
and ignore function referenced from that section. Code references are
detected via relocations (static or dynamic).

Note that the abort handler is preceded by a 4-byte signature byte
sequence and we cannot relocate the handler without that the signature,
otherwise the application may crash. Thus we are ignoring the function,
i.e. making sure it's not separated from its signature.

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