[AArch64] Add SEH_Nop opcodes for BTI hints

[mstorsjo]

Summary

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.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64BranchTargets.cpp

 public:
   static char ID;
   AArch64BranchTargets() : MachineFunctionPass(ID) {}
   void getAnalysisUsage(AnalysisUsage &AU) const override;
   bool runOnMachineFunction(MachineFunction &MF) override;
   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
 
 char AArch64BranchTargets::ID = 0;
 
 INITIALIZE_PASS(AArch64BranchTargets, "aarch64-branch-targets",
                 AARCH64_BRANCH_TARGETS_NAME, false, false)
 
bool AArch64BranchTargets::runOnMachineFunction(MachineFunction &MF) {
   // used for indirect branches, so need BTI instructions.
   SmallPtrSet<MachineBasicBlock *, 8> JumpTableTargets;
   if (auto *JTI = MF.getJumpTableInfo())
     for (auto &JTE : JTI->getJumpTables())
       for (auto *MBB : JTE.MBBs)
         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
     // directly in its translation unit, it can still be called indirectly if
     // the linker decides to add a thunk to it for whatever reason (say, for
     // example, if it is finally placed far from its call site and a BL is not
     // long-range enough). PLT entries and tail-calls use BR, but when they are
     // are in guarded pages should all use x16 or x17 to hold the called
     // address, so we don't need to set CouldJump here. BR instructions in
     // non-guarded pages (which might be non-BTI-aware code) are allowed to
     // branch to a "BTI c" using any register.
     if (&MBB == &*MF.begin())
       CouldCall = true;
 
     // If the block itself is address-taken, it could be indirectly branched
     // to, but not called.
     if (MBB.hasAddressTaken() || JumpTableTargets.count(&MBB))
       CouldJump = true;
 
     if (CouldCall || CouldJump) {
-      addBTI(MBB, CouldCall, CouldJump);
+      addBTI(MBB, CouldCall, CouldJump, HasWinCFI);
       MadeChange = true;
     }
   }
 
   return MadeChange;
 }
 
 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");
 
   const AArch64InstrInfo *TII = static_cast<const AArch64InstrInfo *>(
       MBB.getParent()->getSubtarget().getInstrInfo());
 
   unsigned HintNum = 32;
void AArch64BranchTargets::addBTI(MachineBasicBlock &MBB, bool CouldCall,
 
   // SCTLR_EL1.BT[01] is set to 0 by default which means
   // PACI[AB]SP are implicitly BTI C so no BTI C instruction is needed there.
   if (MBBI != MBB.end() && HintNum == 34 &&
       (MBBI->getOpcode() == AArch64::PACIASP ||
        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);
 }

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

+; 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