[ARM] Remove redundant BTI instructions for table jumps

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.

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

diff --git a/llvm/lib/Target/ARM/ARMBranchTargets.cpp b/llvm/lib/Target/ARM/ARMBranchTargets.cpp
index 8ba3e62..17d0bdd 100644 (file)
--- a/llvm/lib/Target/ARM/ARMBranchTargets.cpp
+++ b/llvm/lib/Target/ARM/ARMBranchTargets.cpp
@@ -65,31 +65,19 @@ bool ARMBranchTargets::runOnMachineFunction(MachineFunction &MF) {
   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;
     }

diff --git a/llvm/lib/Target/ARM/ARMConstantIslandPass.cpp b/llvm/lib/Target/ARM/ARMConstantIslandPass.cpp
index 0fa7258..0c910eb 100644 (file)
--- a/llvm/lib/Target/ARM/ARMConstantIslandPass.cpp
+++ b/llvm/lib/Target/ARM/ARMConstantIslandPass.cpp
@@ -628,6 +628,11 @@ void ARMConstantIslands::doInitialJumpTablePlacement(
     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:

diff --git a/llvm/test/CodeGen/Thumb2/bti-indirect-branches.ll b/llvm/test/CodeGen/Thumb2/bti-indirect-branches.ll
index e238c67..59e3465 100644 (file)
--- a/llvm/test/CodeGen/Thumb2/bti-indirect-branches.ll
+++ b/llvm/test/CodeGen/Thumb2/bti-indirect-branches.ll
@@ -19,20 +19,16 @@ define internal i32 @table_switch(i32 %x) {
 ; 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:

diff --git a/llvm/test/CodeGen/Thumb2/bti-jump-table.mir b/llvm/test/CodeGen/Thumb2/bti-jump-table.mir
index 22e7a77..71a8484 100644 (file)
--- a/llvm/test/CodeGen/Thumb2/bti-jump-table.mir
+++ b/llvm/test/CodeGen/Thumb2/bti-jump-table.mir
@@ -62,7 +62,6 @@ body:             |
   ; 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:
@@ -70,14 +69,11 @@ body:             |
   ; 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:

diff --git a/llvm/test/MC/ARM/remove-redundant-bti.ll b/llvm/test/MC/ARM/remove-redundant-bti.ll
new file mode 100644 (file)
index 0000000..e241193
--- /dev/null
+++ b/llvm/test/MC/ARM/remove-redundant-bti.ll
@@ -0,0 +1,83 @@
+; 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
+}