From 9ad894f5778423bcff5eb872c04ffb11543860b6 Mon Sep 17 00:00:00 2001 From: Maksim Panchenko Date: Tue, 9 Sep 2025 17:40:20 -0700 Subject: [PATCH] [BOLT] Support restartable sequences in tcmalloc 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 | 6 +- bolt/lib/Rewrite/CMakeLists.txt | 1 + bolt/lib/Rewrite/RSeqRewriter.cpp | 72 +++++++++++++++++++ bolt/lib/Rewrite/RewriteInstance.cpp | 2 + bolt/test/X86/rseq.s | 38 ++++++++++ 5 files changed, 117 insertions(+), 2 deletions(-) create mode 100644 bolt/lib/Rewrite/RSeqRewriter.cpp create mode 100644 bolt/test/X86/rseq.s 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 createLinuxKernelRewriter(BinaryContext &); - std::unique_ptr createBuildIDRewriter(BinaryContext &); +std::unique_ptr createLinuxKernelRewriter(BinaryContext &); + std::unique_ptr createPseudoProbeRewriter(BinaryContext &); +std::unique_ptr createRSeqRewriter(BinaryContext &); + std::unique_ptr createSDTRewriter(BinaryContext &); std::unique_ptr 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 +llvm::bolt::createRSeqRewriter(BinaryContext &BC) { + return std::make_unique("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