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

Anatoly Trosinenko via llvm-commits llvm-commits at lists.llvm.org
Thu Oct 5 07:34:34 PDT 2023


https://github.com/atrosinenko created https://github.com/llvm/llvm-project/pull/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.

>From 98361034e6038233ef0bb8cfe72d3bd9ac575709 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Tue, 3 Oct 2023 19:30:58 +0300
Subject: [PATCH] [AArch64][BTI] Prevent Machine Scheduler from moving branch
 targets

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.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  |  38 +++-
 llvm/lib/Target/AArch64/AArch64InstrInfo.h    |   3 +
 .../machine-outliner-noreturn-no-stack.mir    |   4 +-
 .../machine-outliner-noreturn-save-lr.mir     |   8 +-
 .../AArch64/misched-branch-targets.mir        | 195 ++++++++++++++++++
 5 files changed, 237 insertions(+), 11 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/misched-branch-targets.mir

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