[llvm] [AArch64][BTI] Add BTI at EH entries: BTI j for Itanium EH, BTI c for… (PR #155308)

[Shashi Shankar via llvm-commits]

https://github.com/shashforge created https://github.com/llvm/llvm-project/pull/155308

[AArch64][BTI] Add BTI at EH entries

Mark EH landing pads as indirect-branch targets so we emit `bti j` on Itanium EH; treat WinEH funclet entries as call-like to emit `bti c`.

Tests:
  * bti-ehpad.ll: asserts BTI j (or HINT #36) at the landing pad.
  * wineh-bti-funclet.ll: enables +bti and asserts BTI-c (or PAC) at entry;
    use POSIX whitespace classes in FileCheck for portability.

Fixes: #149267


>From ea8c72f931c3817367607b27b94d9c6ab44accc5 Mon Sep 17 00:00:00 2001
From: Shashi Shankar <shashishankar1687 at gmail.com>
Date: Mon, 25 Aug 2025 23:13:19 +0200
Subject: [PATCH] [AArch64][BTI] Add BTI at EH entries: BTI j for Itanium EH,
 BTI c for WinEH

Signed-off-by: Shashi Shankar <shashishankar1687 at gmail.com>
---
 .../Target/AArch64/AArch64BranchTargets.cpp   |  9 ++++-
 llvm/test/CodeGen/AArch64/bti-ehpad.ll        | 29 ++++++++++++++++
 .../test/CodeGen/AArch64/wineh-bti-funclet.ll | 34 +++++++++++++++++++
 3 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/bti-ehpad.ll
 create mode 100644 llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll

diff --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index 3436dc9ef4521..42315ec7eb503 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -97,12 +97,19 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
          (F.hasAddressTaken() || !F.hasLocalLinkage())))
       CouldCall = true;
 
+
     // If the block itself is address-taken, it could be indirectly branched
     // to, but not called.
     if (MBB.isMachineBlockAddressTaken() || MBB.isIRBlockAddressTaken() ||
-        JumpTableTargets.count(&MBB))
+        JumpTableTargets.count(&MBB) || MBB.isEHPad())
       CouldJump = true;
 
+    if (MBB.isEHPad()) {
+      if (HasWinCFI && (MBB.isEHFuncletEntry() || MBB.isCleanupFuncletEntry()))
+        CouldCall = true;
+      else
+        CouldJump = true;
+    }
     if (CouldCall || CouldJump) {
       addBTI(MBB, CouldCall, CouldJump, HasWinCFI);
       MadeChange = true;
diff --git a/llvm/test/CodeGen/AArch64/bti-ehpad.ll b/llvm/test/CodeGen/AArch64/bti-ehpad.ll
new file mode 100644
index 0000000000000..1845cca46e72e
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/bti-ehpad.ll
@@ -0,0 +1,29 @@
+; REQUIRES: aarch64-registered-target
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu %s -o - | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+declare i32 @__gxx_personality_v0(...)
+declare void @may_throw()
+
+define void @test() #0 personality ptr @__gxx_personality_v0 {
+entry:
+  invoke void @may_throw()
+          to label %ret unwind label %lpad
+
+lpad:
+  landingpad { ptr, i32 } cleanup
+  ret void
+
+ret:
+  ret void
+}
+
+; Request BTI in codegen.
+attributes #0 = { "branch-target-enforcement"="true" "target-features"="+bti" }
+
+; CHECK-LABEL: test:
+; The assembler prints the label line like: ".LBB0_2: // %lpad"
+; Match the EH pad comment, then the very next line must be BTI j (or HINT #36).
+; CHECK:      {{//[[:space:]]*%lpad}}
+; CHECK-NEXT: {{(bti[[:space:]]+j|hint[[:space:]]+#36)}}
diff --git a/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll b/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll
new file mode 100644
index 0000000000000..31e193b5425d4
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/wineh-bti-funclet.ll
@@ -0,0 +1,34 @@
+; REQUIRES: aarch64-registered-target
+; RUN: llc -mtriple=aarch64-windows -mattr=+bti -o - %s | FileCheck %s
+
+target triple = "aarch64-unknown-windows-msvc"
+
+declare i32 @__CxxFrameHandler3(...)
+declare void @may_throw()
+
+define dso_local void @"?w@@YAXXZ"() #0 personality ptr @__CxxFrameHandler3 {
+entry:
+  invoke void @may_throw()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:
+  %cs = catchswitch within none [label %catch] unwind to caller
+
+catch:
+  %cp = catchpad within %cs [ptr null, i32 0, ptr null]
+  ; Put some code in the funclet so it materializes with text.
+  call void @may_throw() ["funclet"(token %cp)]
+  catchret from %cp to label %try.cont
+
+try.cont:
+  ret void
+}
+
+; Ask for BTI enforcement at the function level (the RUN line already enables +bti).
+attributes #0 = { "branch-target-enforcement"="true" }
+
+; Function entry should accept call-like entries:
+;  - Either 'hint #34' / 'bti c'
+;  - Or a PAC prologue (paciasp/pacibsp), which also satisfies BTI-c
+; CHECK-LABEL: "?w@@YAXXZ":
+; CHECK: {{(hint[[:space:]]+#34|bti[[:space:]]+c|paciasp|pacibsp)}}