[llvm] 018ac78 - [AArch64] Add SEH_Nop opcodes for BTI hints

[Martin Storsjö via llvm-commits]

Author: Martin Storsjö
Date: 2022-10-11T14:32:01+03:00
New Revision: 018ac7847bb5a18ff68cf05b1d68f37fb32792d8

URL: https://github.com/llvm/llvm-project/commit/018ac7847bb5a18ff68cf05b1d68f37fb32792d8
DIFF: https://github.com/llvm/llvm-project/commit/018ac7847bb5a18ff68cf05b1d68f37fb32792d8.diff

LOG: [AArch64] Add SEH_Nop opcodes for BTI hints

These are harmless for the unwinder - the unwinder doesn't need to
handle them for being able to unwind correctly.

Only add the opcodes when the branch target is in a SEH prologue;
for jumptables e.g. within a function, we shouldn't add any SEH
opcodes.

Differential Revision: https://reviews.llvm.org/D135277

Added: 
    llvm/test/CodeGen/AArch64/wineh-bti.ll

Modified: 
    llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
index d2acd1de74c8b..e2d10511cccd3 100644
--- a/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
+++ b/llvm/lib/Target/AArch64/AArch64BranchTargets.cpp
@@ -39,7 +39,8 @@ class AArch64BranchTargets : public MachineFunctionPass {
   StringRef getPassName() const override { return AARCH64_BRANCH_TARGETS_NAME; }
 
 private:
-  void addBTI(MachineBasicBlock &MBB, bool CouldCall, bool CouldJump);
+  void addBTI(MachineBasicBlock &MBB, bool CouldCall, bool CouldJump,
+              bool NeedsWinCFI);
 };
 } // end anonymous namespace
 
@@ -75,6 +76,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
         JumpTableTargets.insert(MBB);
 
   bool MadeChange = false;
+  bool HasWinCFI = MF.hasWinCFI();
   for (MachineBasicBlock &MBB : MF) {
     bool CouldCall = false, CouldJump = false;
     // Even in cases where a function has internal linkage and is only called
@@ -95,7 +97,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
       CouldJump = true;
 
     if (CouldCall || CouldJump) {
-      addBTI(MBB, CouldCall, CouldJump);
+      addBTI(MBB, CouldCall, CouldJump, HasWinCFI);
       MadeChange = true;
     }
   }
@@ -104,7 +106,7 @@ bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
 }
 
 void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
-                                  bool CouldJump) {
+                                  bool CouldJump, bool HasWinCFI) {
   LLVM_DEBUG(dbgs() << "Adding BTI " << (CouldJump ? "j" : "")
                     << (CouldCall ? "c" : "") << " to " << MBB.getName()
                     << "\n");
@@ -134,6 +136,10 @@ void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
        MBBI->getOpcode() == AArch64::PACIBSP))
     return;
 
+  if (HasWinCFI && MBBI->getFlag(MachineInstr::FrameSetup)) {
+    BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
+            TII->get(AArch64::SEH_Nop));
+  }
   BuildMI(MBB, MBB.begin(), MBB.findDebugLoc(MBB.begin()),
           TII->get(AArch64::HINT))
       .addImm(HintNum);

diff  --git a/llvm/test/CodeGen/AArch64/wineh-bti.ll b/llvm/test/CodeGen/AArch64/wineh-bti.ll
new file mode 100644
index 0000000000000..125b0cb792394
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/wineh-bti.ll
@@ -0,0 +1,65 @@
+; RUN: llc < %s -mtriple=aarch64-windows | FileCheck %s
+
+define dso_local i32 @func(i32 %in) {
+entry:
+  call void asm sideeffect "", "~{x19}"()
+  switch i32 %in, label %def [
+    i32 0, label %lbl1
+    i32 1, label %lbl2
+    i32 2, label %lbl3
+    i32 4, label %lbl4
+  ]
+
+def:
+  ret i32 0
+
+lbl1:
+  call void asm sideeffect "", ""()
+  ret i32 1
+
+lbl2:
+  ret i32 2
+
+lbl3:
+  ret i32 4
+
+lbl4:
+  ret i32 8
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 8, !"branch-target-enforcement", i32 1}
+!1 = !{i32 8, !"sign-return-address", i32 1}
+
+; CHECK-LABEL: func:
+; CHECK-NEXT: .seh_proc func
+; CHECK-NEXT: // %bb.0:
+; CHECK-NEXT: hint #34
+; CHECK-NEXT: .seh_nop
+; CHECK-NEXT: str x19, [sp, #-16]!
+; CHECK-NEXT: .seh_save_reg_x x19, 16
+; CHECK-NEXT: .seh_endprologue
+
+; CHECK:      .LBB0_2:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #1
+
+; CHECK:      .LBB0_3:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #2
+
+; CHECK:      .LBB0_4:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #4
+
+; CHECK:      .LBB0_5:
+; CHECK-NEXT: hint #36
+; CHECK-NEXT: mov w0, #8
+
+; CHECK:      .seh_startepilogue
+; CHECK-NEXT: ldr x19, [sp], #16
+; CHECK-NEXT: .seh_save_reg_x x19, 16
+; CHECK-NEXT: .seh_endepilogue
+; CHECK-NEXT: ret
+; CHECK-NEXT: .seh_endfunclet