[BOLT] Extend Inliner to work on functions with Pointer Authentication

[bgergely0]

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.

[bgergely0]

• [BOLT] Extend Inliner to work on functions with Pointer Authentication #162458 👈 (View in Graphite)
• [BOLT][NFC] Add MCPlusBuilder unittests for PAuth helpers #162251
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Gergely Bálint (bgergely0)

Changes

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.

---

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

4 Files Affected:

• (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+6)
• (modified) bolt/lib/Passes/Inliner.cpp (+18)
• (modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+29)
• (added) bolt/test/AArch64/inline-armv8.3-returns.s (+45)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 2772de73081d1..ba28016a234fb 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 df4f42128605e..5fb1d7458a4bc 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<MCPhysReg> 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

[bgergely0]

force push: synced this to main, to have the change from the previously merged PR in the stack.

[peterwaller-arm]

Thanks for the patch. A couple of problems spotted - I'm thinking it would be good to test more of the variants you intend to function per the opcodes in createMatchingAuth.

[bgergely0]

force push: rebased to main.

[github-actions]

🐧 Linux x64 Test Results

• 582 tests passed
• 36 tests skipped

[bgergely0]

I've added unittests for all types of insts:

• pauth-lr versions using a register
• pauth-lr versions using a label
• normal Armv8.3-A versions

Plus the fix for the tailcall situation also has a unittest.

[peterwaller-arm]

Just one more tweak on the comment but otherwise LGTM.