[BOLT] Support restartable sequences in tcmalloc

[maksfb]

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.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/167195.diff

5 Files Affected:

• (modified) bolt/include/bolt/Rewrite/MetadataRewriters.h (+4-2)
• (modified) bolt/lib/Rewrite/CMakeLists.txt (+1)
• (added) bolt/lib/Rewrite/RSeqRewriter.cpp (+72)
• (modified) bolt/lib/Rewrite/RewriteInstance.cpp (+2)
• (added) bolt/test/X86/rseq.s (+38)

diff --git a/bolt/include/bolt/Rewrite/MetadataRewriters.h b/bolt/include/bolt/Rewrite/MetadataRewriters.h
index 2c09c879b9128..6b74b0e776997 100644
--- a/bolt/include/bolt/Rewrite/MetadataRewriters.h
+++ b/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 &);
diff --git a/bolt/lib/Rewrite/CMakeLists.txt b/bolt/lib/Rewrite/CMakeLists.txt
index 5b15edcacb482..bc1b2ed3c2e3c 100644
--- a/bolt/lib/Rewrite/CMakeLists.txt
+++ b/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
 
diff --git a/bolt/lib/Rewrite/RSeqRewriter.cpp b/bolt/lib/Rewrite/RSeqRewriter.cpp
new file mode 100644
index 0000000000000..46bce66d13ddf
--- /dev/null
+++ b/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);
+}
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 5769577aa3f74..8d6731e7540a8 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/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));
diff --git a/bolt/test/X86/rseq.s b/bolt/test/X86/rseq.s
new file mode 100644
index 0000000000000..ef81bca02c8b7
--- /dev/null
+++ b/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

[aaupov]

Does it need to be defined in this and not outer scope?

[maksfb]

Thanks for the review.

Doesn't have to be, but I selected the minimal scope for handleRelocation().