From dd11305a2b57b0da0d1529164ebcbb3a4b1a0343 Mon Sep 17 00:00:00 2001 From: Maksim Panchenko Date: Tue, 9 Sep 2025 17:40:20 -0700 Subject: [PATCH 1/2] [BOLT][WIP] Support restartable sequences in tcmalloc --- bolt/include/bolt/Rewrite/MetadataRewriters.h | 6 +- bolt/lib/Rewrite/CMakeLists.txt | 1 + bolt/lib/Rewrite/RSeqRewriter.cpp | 71 +++++++++++++++++++ bolt/lib/Rewrite/RewriteInstance.cpp | 2 + bolt/test/X86/rseq.s | 32 +++++++++ 5 files changed, 110 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 b71bd6cad2505..429c2d7712979 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 &); } // namespace bolt diff --git a/bolt/lib/Rewrite/CMakeLists.txt b/bolt/lib/Rewrite/CMakeLists.txt index 775036063dd5a..12914bfe64301 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 NO_EXPORT diff --git a/bolt/lib/Rewrite/RSeqRewriter.cpp b/bolt/lib/Rewrite/RSeqRewriter.cpp new file mode 100644 index 0000000000000..7cbe9ac17da52 --- /dev/null +++ b/bolt/lib/Rewrite/RSeqRewriter.cpp @@ -0,0 +1,71 @@ +//===- 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 postCFGInitializer() override { + for (const BinarySection &Section : BC.allocatableSections()) { + if (Section.getName() != "__rseq_cs") + continue; + + auto handleRelocation = [&](const Relocation &Rel) { + 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.outs() << "BOLT-INFO: restartable sequence detected in " << *BF + << ". Function will not be optimized\n"; + BF->setSimple(false); + } else { + BC.errs() << "BOLT-WARNING: no function found matching dynamic " + "relocation in __rseq_cs\n"; + } + }; + + for (const Relocation &Rel : Section.dynamicRelocations()) + handleRelocation(Rel); + + for (const Relocation &Rel : Section.relocations()) + handleRelocation(Rel); + } + + 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 a6e4dbc9c192f..6e07d2ea8d975 100644 --- a/bolt/lib/Rewrite/RewriteInstance.cpp +++ b/bolt/lib/Rewrite/RewriteInstance.cpp @@ -3338,6 +3338,8 @@ void RewriteInstance::initializeMetadataManager() { MetadataManager.registerRewriter(createPseudoProbeRewriter(*BC)); + MetadataManager.registerRewriter(createRSeqRewriter(*BC)); + MetadataManager.registerRewriter(createSDTRewriter(*BC)); } diff --git a/bolt/test/X86/rseq.s b/bolt/test/X86/rseq.s new file mode 100644 index 0000000000000..7ebc57e1128f4 --- /dev/null +++ b/bolt/test/X86/rseq.s @@ -0,0 +1,32 @@ +## Check that llvm-bolt avoids optimization of functions referenced from +## __rseq_cs section, i.e. containing critical sections 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 \ +# RUN: | FileCheck %s --check-prefix=CHECK-NO-PIE +# RUN: %clang %cflags %s -o %t.pie -nostdlib -pie -Wl,-q +# RUN: llvm-bolt %t.pie -o %t.pie.bolt 2>&1 \ +# RUN: | FileCheck %s --check-prefix=CHECK-PIE + +# CHECK-NO-PIE: Binary Function "_start" +# CHECK-NO-PIE: IsSimple +# CHECK-NO-PIE-SAME: 0 + +# CHECK-PIE: restartable sequence detected in _start + +.global _start + .type _start, %function +_start: + pushq %rbp + mov %rsp, %rbp +.L1: + pop %rbp + retq +.size _start, .-_start + +.reloc 0, R_X86_64_NONE + +.section __rseq_cs, "aw" +.balign 32 + .quad .L1 From a972175f3114f083e6ac181088be53b81d250070 Mon Sep 17 00:00:00 2001 From: Maksim Panchenko Date: Thu, 11 Sep 2025 13:28:27 -0700 Subject: [PATCH 2/2] fixup! [BOLT][WIP] Support restartable sequences in tcmalloc --- bolt/lib/Rewrite/RSeqRewriter.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bolt/lib/Rewrite/RSeqRewriter.cpp b/bolt/lib/Rewrite/RSeqRewriter.cpp index 7cbe9ac17da52..3a3230818943c 100644 --- a/bolt/lib/Rewrite/RSeqRewriter.cpp +++ b/bolt/lib/Rewrite/RSeqRewriter.cpp @@ -45,7 +45,7 @@ class RSeqRewriter final : public MetadataRewriter { if (BF) { BC.outs() << "BOLT-INFO: restartable sequence detected in " << *BF << ". Function will not be optimized\n"; - BF->setSimple(false); + BF->setIgnored(); } else { BC.errs() << "BOLT-WARNING: no function found matching dynamic " "relocation in __rseq_cs\n";