[AArch64][BTI] Prevent Machine Scheduler from moving branch targets

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -1133,6 +1133,11 @@ bool AArch64InstrInfo::isSchedulingBoundary(const MachineInstr &MI,
                                             const MachineFunction &MF) const {
   if (TargetInstrInfo::isSchedulingBoundary(MI, MBB, MF))
     return true;
+
+  // Do not move an instruction that can be recognized as a branch target.
+  if (hasBTISemantics(MI))
+    return true;
+
   switch (MI.getOpcode()) {
   case AArch64::HINT:
     // CSDB hints are scheduling barriers.
@@ -4081,6 +4086,32 @@ bool AArch64InstrInfo::isQForm(const MachineInstr &MI) {
   return llvm::any_of(MI.operands(), IsQFPR);
 }
 
+bool AArch64InstrInfo::hasBTISemantics(const MachineInstr &MI) {
+  switch (MI.getOpcode()) {
+  case AArch64::BRK:
+  case AArch64::HLT:
+  case AArch64::PACIASP:
+  case AArch64::PACIBSP:
+    // Implicit BTI behavior.
+    return true;
+  case AArch64::PAUTH_PROLOGUE:
+    // PAUTH_PROLOGUE expands to PACI(A|B)SP.
+    return true;
+  case AArch64::HINT: {
+    unsigned Imm = MI.getOperand(0).getImm();
+    // Explicit BTI instruction.
+    if (Imm == 32 || Imm == 34 || Imm == 36 || Imm == 38)
+      return true;
+    // PACI(A|B)SP instructions.
+    if (Imm == 25 || Imm == 27)
+      return true;
+    return false;
+  }
+  default:
+    return false;
+  }
+}
+
 bool AArch64InstrInfo::isFpOrNEON(const MachineInstr &MI) {
   auto IsFPR = [&](const MachineOperand &Op) {
     if (!Op.isReg())
@@ -8829,11 +8860,8 @@ AArch64InstrInfo::getOutliningTypeImpl(MachineBasicBlock::iterator &MIT,
 
   // Don't outline BTI instructions, because that will prevent the outlining
   // site from being indirectly callable.
-  if (MI.getOpcode() == AArch64::HINT) {
-    int64_t Imm = MI.getOperand(0).getImm();
-    if (Imm == 32 || Imm == 34 || Imm == 36 || Imm == 38)
-      return outliner::InstrType::Illegal;
-  }
+  if (hasBTISemantics(MI))
+    return outliner::InstrType::Illegal;
 
   return outliner::InstrType::Legal;
 }

llvm/lib/Target/AArch64/AArch64InstrInfo.h

@@ -120,6 +120,9 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
   /// Returns whether the instruction is in Q form (128 bit operands)
   static bool isQForm(const MachineInstr &MI);
 
+  /// Returns whether the instruction can be compatible with non-zero BTYPE.
+  static bool hasBTISemantics(const MachineInstr &MI);
+
   /// Returns the index for the immediate for a given instruction.
   static unsigned getLoadStoreImmIdx(unsigned Opc);
 

llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir

@@ -34,7 +34,7 @@ body:             |
     $x0 = ORRXrs $xzr, $x1, 0
     $x1 = ORRXrs $xzr, $x2, 0
     BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8,  implicit-def $x9, implicit-def $x10,  implicit-def $x11, implicit-def $x12, implicit-def $x13,  implicit-def $x14, implicit-def $x15, implicit-def $x18
-    BRK 1
+    HINT 0
 ...
 ---
 name:            stack_2
@@ -56,7 +56,7 @@ body:             |
     $x0 = ORRXrs $xzr, $x1, 0
     $x1 = ORRXrs $xzr, $x2, 0
     BL @baz, implicit-def dead $lr, implicit $sp, implicit $x8, implicit $x0, implicit $x1, implicit $x3, implicit $x4, implicit-def $sp, implicit-def $x5, implicit-def $x6, implicit-def $x7, implicit-def $x8,  implicit-def $x9, implicit-def $x10,  implicit-def $x11, implicit-def $x12, implicit-def $x13,  implicit-def $x14, implicit-def $x15, implicit-def $x18
-    BRK 1
+    HINT 0
 ...
 ---
 name:            baz

llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir

@@ -31,7 +31,7 @@ body:             |
     ; CHECK: $lr = ORRXrs $xzr, $x0, 0
     $w3 = ORRWri $wzr, 1
     $w4 = ORRWri $wzr, 1
-    BRK 1
+    HINT 0
     $w5 = ORRWri $wzr, 1
     $w6 = ORRWri $wzr, 1
 ...
@@ -53,7 +53,7 @@ body:             |
     ; CHECK: $lr = ORRXrs $xzr, $x0, 0
     $w3 = ORRWri $wzr, 1
     $w4 = ORRWri $wzr, 1
-    BRK 1
+    HINT 0
     $w5 = ORRWri $wzr, 1
     $w6 = ORRWri $wzr, 1
 ...
@@ -75,7 +75,7 @@ body:             |
     ; CHECK: $lr = ORRXrs $xzr, $x0, 0
     $w3 = ORRWri $wzr, 1
     $w4 = ORRWri $wzr, 1
-    BRK 1
+    HINT 0
     $w5 = ORRWri $wzr, 1
     $w6 = ORRWri $wzr, 1
 ...
@@ -97,7 +97,7 @@ body:             |
     ; CHECK: $lr = ORRXrs $xzr, $x0, 0
     $w3 = ORRWri $wzr, 1
     $w4 = ORRWri $wzr, 1
-    BRK 1
+    HINT 0
     $w5 = ORRWri $wzr, 1
     $w6 = ORRWri $wzr, 1
 ...

llvm/test/CodeGen/AArch64/misched-branch-targets.mir

@@ -0,0 +1,195 @@
+# RUN: llc -o - -run-pass=machine-scheduler -misched=shuffle %s | FileCheck %s
+# RUN: llc -o - -run-pass=postmisched %s | FileCheck %s
+
+# Check that instructions that are recognized as branch targets by BTI
+# are not reordered by machine instruction schedulers.
+
+--- |
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64-unknown-linux-gnu"
+
+  define i32 @f_pac_pseudo(i32 %a, i32 %b, i32 %c) #0 "sign-return-address"="all" {
+  entry:
+    ret i32 0
+  }
+
+  define i32 @f_pac(i32 %a, i32 %b, i32 %c) #0 "sign-return-address"="all" {
+  entry:
+    ret i32 0
+  }
+
+  define i32 @f_bti(i32 %a, i32 %b, i32 %c) #0 {
+  entry:
+    ret i32 0
+  }
+
+  define i32 @f_brk(i32 %a, i32 %b, i32 %c) #0 {
+  entry:
+    ret i32 0
+  }
+
+  define i32 @f_hlt(i32 %a, i32 %b, i32 %c) #0 {
+  entry:
+    ret i32 0
+  }
+
+  define i32 @f_nop(i32 %a, i32 %b, i32 %c) #0 {
+  entry:
+    ret i32 0
+  }
+
+  attributes #0 = { nounwind memory(none) "target-features"="+v8.2a" }
+
+...
+---
+name:            f_pac_pseudo
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit $lr, implicit $sp
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# PAUTH_EPILOGUE instruction is omitted for simplicity as it is technically possible
+# to move it, so it may end up at a less obvious position in a basic block.
+
+# CHECK-LABEL: name: f_pac_pseudo
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           frame-setup PAUTH_PROLOGUE implicit-def $lr, implicit {{.*}}$lr, implicit $sp
+# CHECK-NEXT:      $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-NEXT:      $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...
+---
+name:            f_pac
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    frame-setup PACIASP implicit-def $lr, implicit $lr, implicit $sp
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# AUTIASP is omitted, see above.
+
+# CHECK-LABEL: name: f_pac
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           frame-setup PACIASP implicit-def $lr, implicit {{.*}}$lr, implicit $sp
+# CHECK-NEXT:      $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-NEXT:      $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...
+---
+name:            f_bti
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    HINT 34
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# CHECK-LABEL: name: f_bti
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           HINT 34
+# CHECK-NEXT:      $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-NEXT:      $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...
+---
+name:            f_brk
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    BRK 1
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# CHECK-LABEL: name: f_brk
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           BRK 1
+# CHECK-NEXT:      $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-NEXT:      $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...
+---
+name:            f_hlt
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    HLT 1
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# CHECK-LABEL: name: f_hlt
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           HLT 1
+# CHECK-NEXT:      $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-NEXT:      $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...
+---
+name:            f_nop
+alignment:       4
+tracksRegLiveness: true
+body:             |
+  bb.0.entry:
+    liveins: $w0, $w1, $w2, $lr
+
+    HINT 0
+    $w8 = ADDWrs $w0, $w1, 0
+    $w0 = MADDWrrr $w8, $w2, $wzr
+    RET undef $lr, implicit $w0
+
+# Check that BTI-related instructions are left intact not because *anything*
+# is left intact.
+
+# CHECK-LABEL: name: f_nop
+# CHECK:       body:             |
+# CHECK-NEXT:    bb.0.entry:
+# CHECK-NEXT:      liveins: $w0, $w1, $w2, $lr
+#
+# CHECK:           $w8 = ADDWrs {{.*}}$w0, {{.*}}$w1, 0
+# CHECK-DAG:       $w0 = MADDWrrr {{.*}}$w8, {{.*}}$w2, $wzr
+# CHECK-DAG:       HINT 0
+# CHECK-NEXT:      RET undef $lr, implicit {{.*}}$w0
+
+...