From ad23391cd09a731cd89f51640dde2d8261854f67 Mon Sep 17 00:00:00 2001 From: Gergely Balint Date: Wed, 8 Oct 2025 11:17:31 +0000 Subject: [PATCH] [BOLT] Extend Inliner to work on functions with Pointer Autentication The inliner uses DirectSP to check if a function has instructions that modify the SP. Exceptions are stack Push and Pop instructions. We can also allow pointer signing and authentication instructions. The inliner removes the Return instructions from the inlined functions. If it is a fused pointer-authentication-and-return (e.g. RETAA), we have to generate a new authentication instruction in place of the Return. --- bolt/include/bolt/Core/MCPlusBuilder.h | 6 +++ bolt/lib/Passes/Inliner.cpp | 18 ++++++++ .../Target/AArch64/AArch64MCPlusBuilder.cpp | 29 ++++++++++++ bolt/test/AArch64/inline-armv8.3-returns.s | 45 +++++++++++++++++++ 4 files changed, 98 insertions(+) create mode 100644 bolt/test/AArch64/inline-armv8.3-returns.s diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h index d666c10885ad5..5b2c9bcd5ba49 100644 --- a/bolt/include/bolt/Core/MCPlusBuilder.h +++ b/bolt/include/bolt/Core/MCPlusBuilder.h @@ -632,6 +632,12 @@ class MCPlusBuilder { return false; } + /// Generate the matching pointer authentication instruction from a fused + /// pauth-and-return instruction. + virtual void createMatchingAuth(const MCInst &AuthAndRet, MCInst &Auth) { + llvm_unreachable("not implemented"); + } + /// Returns the register used as a return address. Returns std::nullopt if /// not applicable, such as reading the return address from a system register /// or from the stack. diff --git a/bolt/lib/Passes/Inliner.cpp b/bolt/lib/Passes/Inliner.cpp index 9b28c7efde5bf..913ff3d554a5b 100644 --- a/bolt/lib/Passes/Inliner.cpp +++ b/bolt/lib/Passes/Inliner.cpp @@ -195,6 +195,13 @@ InliningInfo getInliningInfo(const BinaryFunction &BF) { if (BC.MIB->isPush(Inst) || BC.MIB->isPop(Inst)) continue; + // Pointer signing and authenticatin instructions are used around + // Push and Pop. These are also straightforward to handle. + if (BC.isAArch64() && + (BC.MIB->isPSignOnLR(Inst) || BC.MIB->isPAuthOnLR(Inst) || + BC.MIB->isPAuthAndRet(Inst))) + continue; + DirectSP |= BC.MIB->hasDefOfPhysReg(Inst, SPReg) || BC.MIB->hasUseOfPhysReg(Inst, SPReg); } @@ -338,6 +345,17 @@ Inliner::inlineCall(BinaryBasicBlock &CallerBB, BC.Ctx.get()); } + // Handling fused authentication and return instructions (Armv8.3-A): + // if the Return here is RETA(A|B), we have to keep the authentication + // part. + // RETAA -> AUTIASP + RET + // RETAB -> AUTIBSP + RET + if (BC.isAArch64() && BC.MIB->isPAuthAndRet(Inst)) { + MCInst Auth; + BC.MIB->createMatchingAuth(Inst, Auth); + InsertII = + std::next(InlinedBB->insertInstruction(InsertII, std::move(Auth))); + } if (CSIsTailCall || (!MIB.isCall(Inst) && !MIB.isReturn(Inst))) { InsertII = std::next(InlinedBB->insertInstruction(InsertII, std::move(Inst))); diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp index 6954cb295e86a..ff480942643a7 100644 --- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp +++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp @@ -266,6 +266,35 @@ class AArch64MCPlusBuilder : public MCPlusBuilder { Inst.getOpcode() == AArch64::RETABSPPCr; } + void createMatchingAuth(const MCInst &AuthAndRet, MCInst &Auth) override { + assert(isPAuthAndRet(AuthAndRet) && + "Not a fused pauth-and-return instruction"); + + Auth.clear(); + switch (AuthAndRet.getOpcode()) { + case AArch64::RETAA: + Auth.setOpcode(AArch64::AUTIASP); + break; + case AArch64::RETAB: + Auth.setOpcode(AArch64::AUTIBSP); + break; + case AArch64::RETAASPPCi: + Auth.setOpcode(AArch64::AUTIASPPCi); + break; + case AArch64::RETABSPPCi: + Auth.setOpcode(AArch64::AUTIBSPPCi); + break; + case AArch64::RETAASPPCr: + Auth.setOpcode(AArch64::AUTIASPPCr); + break; + case AArch64::RETABSPPCr: + Auth.setOpcode(AArch64::AUTIBSPPCr); + break; + default: + llvm_unreachable("Unhandled fused pauth-and-return instruction"); + } + } + std::optional getSignedReg(const MCInst &Inst) const override { switch (Inst.getOpcode()) { case AArch64::PACIA: diff --git a/bolt/test/AArch64/inline-armv8.3-returns.s b/bolt/test/AArch64/inline-armv8.3-returns.s new file mode 100644 index 0000000000000..055b589476caf --- /dev/null +++ b/bolt/test/AArch64/inline-armv8.3-returns.s @@ -0,0 +1,45 @@ +# This test checks that inlining functions with fused pointer-auth-and-return +# instructions is properly handled by BOLT. + +# REQUIRES: system-linux + +# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown -mattr=+v8.3a %s -o %t.o +# RUN: %clang %cflags -O0 %t.o -o %t.exe -Wl,-q +# RUN: llvm-bolt --inline-all --print-inline --print-only=_Z3barP1A \ +# RUN: %t.exe -o %t.bolt | FileCheck %s + +# CHECK: BOLT-INFO: inlined 0 calls at 1 call sites in 2 iteration(s). Change in binary size: 8 bytes. +# CHECK: Binary Function "_Z3barP1A" after inlining { +# CHECK-NOT: bl _Z3fooP1A +# CHECK: ldr x8, [x0] +# CHECK-NEXT: ldr w0, [x8] +# CHECK-NEXT: autiasp + + .text + .globl _Z3fooP1A + .type _Z3fooP1A,@function +_Z3fooP1A: + paciasp + ldr x8, [x0] + ldr w0, [x8] + retaa + .size _Z3fooP1A, .-_Z3fooP1A + + .globl _Z3barP1A + .type _Z3barP1A,@function +_Z3barP1A: + stp x29, x30, [sp, #-16]! + mov x29, sp + bl _Z3fooP1A + mul w0, w0, w0 + ldp x29, x30, [sp], #16 + ret + .size _Z3barP1A, .-_Z3barP1A + + .globl main + .p2align 2 + .type main,@function +main: + mov w0, wzr + ret + .size main, .-main