[llvm] f1b2dd2 - [AArch64][BTI] Prevent Machine Scheduler from moving branch targets (#68313)

via llvm-commits llvm-commits at lists.llvm.org
Fri Oct 6 08:05:07 PDT 2023


Author: Anatoly Trosinenko
Date: 2023-10-06T18:05:03+03:00
New Revision: f1b2dd2a111f038420b3f69d4ce0b3b3f245c873

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

LOG: [AArch64][BTI] Prevent Machine Scheduler from moving branch targets (#68313)

Moving instructions that are recognized as branch targets by BTI can
result in runtime crash.

In outliner tests, replaced "BRK 1" with "HINT 0" (a.k.a. NOP) as a
generic outlinable instruction.

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

Modified: 
    llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
    llvm/lib/Target/AArch64/AArch64InstrInfo.h
    llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir
    llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 830bf77b4e5d05c..5a234bceb25ed0a 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/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;
 }

diff  --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.h b/llvm/lib/Target/AArch64/AArch64InstrInfo.h
index f1a4928939bcd9e..f5874d7856f8d24 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.h
+++ b/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);
 

diff  --git a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir b/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir
index b13887a02ebd232..f926040f3932d71 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-no-stack.mir
+++ b/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

diff  --git a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir b/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir
index 157cbbb51b5b782..c688f4370b0e9dd 100644
--- a/llvm/test/CodeGen/AArch64/machine-outliner-noreturn-save-lr.mir
+++ b/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
 ...

diff  --git a/llvm/test/CodeGen/AArch64/misched-branch-targets.mir b/llvm/test/CodeGen/AArch64/misched-branch-targets.mir
new file mode 100644
index 000000000000000..f32c1e964f97356
--- /dev/null
+++ b/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
+
+...


        


More information about the llvm-commits mailing list