[ARM] Remove redundant BTI instructions for table jumps

[JiruiWu]

Summary

A BTI instruction was previously inserted at the beginning of each block
that has its address stored in a jump table. Jump tables only emit
indirect jumps in ARM or Thumb1 modes. However, PACBTI is not supported
in these modes. As a result, BTI instructions emitted by jump tables are
redundant. Removing redundant BTI instructions improves the code size
and prevents potential gadgets.

[olista01]

llvm/lib/Target/ARM/ARMBranchTargets.cpp

67–68

This comment can be removed too.

llvm/test/MC/AArch64/remove-redundant-bti.s

1 ↗

(On Diff #499105)

The filename for this should end in .ll, because it contains IR, not assembly.

30 ↗

(On Diff #499105)

These tests should also check the tbb/tbh/mov instructions which branch to these labels, to make sure that the two test cases are testing different was of implementing jump tables.

[JiruiWu]

Addressed review comments.

[olista01]

LGTM

llvm/lib/Target/ARM/ARMBranchTargets.cpp

 bool ARMBranchTargets::runOnMachineFunction(MachineFunction &MF) {
   if (!MF.getInfo<ARMFunctionInfo>()->branchTargetEnforcement())
     return false;
 
   LLVM_DEBUG(dbgs() << "********** ARM Branch Targets  **********\n"
                     << "********** Function: " << MF.getName() << '\n');
   const ARMInstrInfo &TII =
       *static_cast<const ARMInstrInfo *>(MF.getSubtarget().getInstrInfo());
 
-  // LLVM does not consider basic blocks which are the targets of jump tables
-  // to be address-taken (the address can't escape anywhere else), but they are
-  // used for indirect branches, so need BTI instructions.
-  SmallPtrSet<const MachineBasicBlock *, 8> JumpTableTargets;
-  if (const MachineJumpTableInfo *JTI = MF.getJumpTableInfo())
-    for (const MachineJumpTableEntry &JTE : JTI->getJumpTables())
-      for (const MachineBasicBlock *MBB : JTE.MBBs)
-        JumpTableTargets.insert(MBB);
-
   bool MadeChange = false;
   for (MachineBasicBlock &MBB : MF) {
-    bool NeedBTI = false;
     bool IsFirstBB = &MBB == &MF.front();
 
     // Every function can potentially be called indirectly (even if it has
     // static linkage, due to linker-generated veneers).
-    if (IsFirstBB)
-      NeedBTI = true;
-
     // If the block itself is address-taken, or is an exception landing pad, it
     // could be indirectly branched to.
-    if (MBB.hasAddressTaken() || MBB.isEHPad() || JumpTableTargets.count(&MBB))
-      NeedBTI = true;
+    // Jump tables only emit indirect jumps (JUMPTABLE_ADDRS) in ARM or Thumb1
+    // modes. These modes do not support PACBTI. As a result, BTI instructions
+    // are not added in the destination blocks.
 
-    if (NeedBTI) {
+    if (IsFirstBB || MBB.hasAddressTaken() || MBB.isEHPad()) {
       addBTI(TII, MBB, IsFirstBB);
       MadeChange = true;
     }
   }
 
   return MadeChange;
 }
 

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

for (MachineBasicBlock &MBB : *MF) {
     switch (MI->getOpcode()) {
     default:
       continue;
     case ARM::BR_JTadd:
     case ARM::BR_JTr:
     case ARM::tBR_JTr:
     case ARM::BR_JTm_i12:
     case ARM::BR_JTm_rs:
+      // These instructions are emitted only in ARM or Thumb1 modes which do not
+      // support PACBTI. Hence we don't add BTI instructions in the destination
+      // blocks.
+      assert(!MF->getInfo<ARMFunctionInfo>()->branchTargetEnforcement() &&
+             "Branch protection must not be enabled for Arm or Thumb1 modes");
       JTOpcode = ARM::JUMPTABLE_ADDRS;
       break;
     case ARM::t2BR_JT:
       JTOpcode = ARM::JUMPTABLE_INSTS;
       break;
     case ARM::tTBB_JT:
     case ARM::t2TBB_JT:
       JTOpcode = ARM::JUMPTABLE_TBB;

llvm/test/CodeGen/Thumb2/bti-indirect-branches.ll

 ; CHECK-NEXT:  @ %bb.2:
 ; CHECK-NEXT:  .LJTI0_0:
 ; CHECK-NEXT:    .byte (.LBB0_5-(.LCPI0_0+4))/2
 ; CHECK-NEXT:    .byte (.LBB0_3-(.LCPI0_0+4))/2
 ; CHECK-NEXT:    .byte (.LBB0_6-(.LCPI0_0+4))/2
 ; CHECK-NEXT:    .byte (.LBB0_7-(.LCPI0_0+4))/2
 ; CHECK-NEXT:    .p2align 1
 ; CHECK-NEXT:  .LBB0_3: @ %bb2
-; CHECK-NEXT:    bti
 ; CHECK-NEXT:    movs r0, #2
 ; CHECK-NEXT:    bx lr
 ; CHECK-NEXT:  .LBB0_4: @ %sw.epilog
 ; CHECK-NEXT:    movs r0, #0
 ; CHECK-NEXT:  .LBB0_5: @ %return
-; CHECK-NEXT:    bti
 ; CHECK-NEXT:    bx lr
 ; CHECK-NEXT:  .LBB0_6: @ %bb3
-; CHECK-NEXT:    bti
 ; CHECK-NEXT:    movs r0, #3
 ; CHECK-NEXT:    bx lr
 ; CHECK-NEXT:  .LBB0_7: @ %bb4
-; CHECK-NEXT:    bti
 ; CHECK-NEXT:    movs r0, #4
 ; CHECK-NEXT:    bx lr
 entry:
   switch i32 %x, label %sw.epilog [
     i32 1, label %bb1
     i32 2, label %bb2
     i32 3, label %bb3
     i32 4, label %bb4

llvm/test/CodeGen/Thumb2/bti-jump-table.mir

body:             |
   ; CHECK: bb.1.entry:
   ; CHECK:   successors: %bb.4(0x20000000), %bb.2(0x20000000), %bb.5(0x20000000), %bb.6(0x20000000)
   ; CHECK:   liveins: $r1
   ; CHECK:   renamable $r0, dead $cpsr = tMOVi8 1, 14 /* CC::al */, $noreg
   ; CHECK:   renamable $r2 = t2LEApcrelJT %jump-table.0, 14 /* CC::al */, $noreg
   ; CHECK:   renamable $r2 = t2ADDrs killed renamable $r2, renamable $r1, 18, 14 /* CC::al */, $noreg, $noreg
   ; CHECK:   t2BR_JT killed renamable $r2, killed renamable $r1, %jump-table.0
   ; CHECK: bb.2.bb2:
-  ; CHECK:   t2BTI
   ; CHECK:   renamable $r0, dead $cpsr = tMOVi8 2, 14 /* CC::al */, $noreg
   ; CHECK:   tBX_RET 14 /* CC::al */, $noreg, implicit $r0
   ; CHECK: bb.3.sw.epilog:
   ; CHECK:   successors: %bb.4(0x80000000)
   ; CHECK:   renamable $r0, dead $cpsr = tMOVi8 0, 14 /* CC::al */, $noreg
   ; CHECK: bb.4.return:
   ; CHECK:   liveins: $r0
-  ; CHECK:   t2BTI
   ; CHECK:   tBX_RET 14 /* CC::al */, $noreg, implicit $r0
   ; CHECK: bb.5.bb3:
-  ; CHECK:   t2BTI
   ; CHECK:   renamable $r0, dead $cpsr = tMOVi8 3, 14 /* CC::al */, $noreg
   ; CHECK:   tBX_RET 14 /* CC::al */, $noreg, implicit $r0
   ; CHECK: bb.6.bb4:
-  ; CHECK:   t2BTI
   ; CHECK:   renamable $r0, dead $cpsr = tMOVi8 4, 14 /* CC::al */, $noreg
   ; CHECK:   tBX_RET 14 /* CC::al */, $noreg, implicit $r0
   bb.0.entry:
     successors: %bb.5, %bb.1
     liveins: $r0
 
     renamable $r1, dead $cpsr = tSUBi3 killed renamable $r0, 1, 14 /* CC::al */, $noreg
     tCMPi8 renamable $r1, 3, 14 /* CC::al */, $noreg, implicit-def $cpsr

llvm/test/MC/ARM/remove-redundant-bti.ll

+; RUN: llc -O2 -mtriple=thumbv7m-eabi < %s | FileCheck %s
+; RUN: llc -O2 -mtriple=thumbv8m.main-eabi < %s | FileCheck %s
+
+define dso_local i32 @test_for_tbb_tbh_cases(i32 noundef %x) local_unnamed_addr #0 {
+entry:
+  switch i32 %x, label %sw.default [
+    i32 0, label %sw.bb
+    i32 1, label %return
+    i32 2, label %sw.bb2
+    i32 3, label %sw.bb3
+  ]
+
+sw.bb:                                            ; preds = %entry
+  %call = tail call i32 @g(i32 noundef 0) #2
+  %add = add nsw i32 %call, 1
+  br label %return
+
+sw.bb2:                                           ; preds = %entry
+  br label %return
+
+sw.bb3:                                           ; preds = %entry
+  br label %return
+
+sw.default:                                       ; preds = %entry
+  br label %return
+
+return:                                           ; preds = %entry, %sw.default, %sw.bb3, %sw.bb2, %sw.bb
+  %retval.0 = phi i32 [ 5, %sw.default ], [ 4, %sw.bb3 ], [ 3, %sw.bb2 ], [ %add, %sw.bb ], [ 2, %entry ]
+  ret i32 %retval.0
+
+; CHECK: tbb  [pc, r1]
+; CHECK-LABEL: .LBB0_3:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB0_4:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB0_5:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB0_6:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB0_7:
+; CHECK-NOT: bti
+}
+
+declare dso_local i32 @g(i32 noundef) local_unnamed_addr #1
+
+define dso_local i32 @test_for_direct_jump_cases(i32 noundef %x) local_unnamed_addr #0 {
+entry:
+  switch i32 %x, label %sw.default [
+    i32 0, label %sw.bb
+    i32 2, label %cleanup
+    i32 8, label %sw.bb2
+    i32 5, label %sw.bb3
+  ]
+
+sw.bb:                                            ; preds = %entry
+  tail call void asm sideeffect ".space 140000", ""()
+  br label %cleanup
+
+sw.bb2:                                           ; preds = %entry
+  br label %cleanup
+
+sw.bb3:                                           ; preds = %entry
+  br label %cleanup
+
+sw.default:                                       ; preds = %entry
+  br label %cleanup
+
+cleanup:                                          ; preds = %entry, %sw.default, %sw.bb3, %sw.bb2, %sw.bb
+  %retval.0 = phi i32 [ 5, %sw.default ], [ 4, %sw.bb3 ], [ 3, %sw.bb2 ], [ 1, %sw.bb ], [ %x, %entry ]
+  ret i32 %retval.0
+
+; CHECK: mov  pc, r1
+; CHECK-LABEL: .LBB1_3:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB1_4:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB1_5:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB1_6:
+; CHECK-NOT: bti
+; CHECK-LABEL: .LBB1_7:
+; CHECK-NOT: bti
+}