[llvm] [AArch64] Fix stack probing clobbering flags (PR #81879)

Momchil Velikov via llvm-commits llvm-commits at lists.llvm.org
Tue Feb 20 05:27:41 PST 2024


https://github.com/momchil-velikov updated https://github.com/llvm/llvm-project/pull/81879

>From 965cec4603efbaa1296e6ed2236f65df4f929f43 Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov at arm.com>
Date: Thu, 15 Feb 2024 17:11:05 +0000
Subject: [PATCH 1/2] [AArch64] Fix stack probing clobbering flags

Certain stack probing sequences might clobber flags, then we can't use a
block as a prologue if the flags register is a live-in on entry to that
block.
---
 .../Target/AArch64/AArch64FrameLowering.cpp   |   5 +
 .../AArch64/stack-probing-shrink-wrap.mir     | 107 ++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/stack-probing-shrink-wrap.mir

diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index 0e9adde564b3e5..c62e750a9b07b0 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1061,6 +1061,11 @@ bool AArch64FrameLowering::canUseAsPrologue(
       return false;
   }
 
+  // Certain stack probing sequences might clobber flags, then we can't use
+  // the block as a prologue if the flags register is a live-in.
+  if (TLI->hasInlineStackProbe(*MF) && MBB.isLiveIn(AArch64::NZCV))
+    return false;
+
   // Don't need a scratch register if we're not going to re-align the stack or
   // emit stack probes.
   if (!RegInfo->hasStackRealignment(*MF) && !TLI->hasInlineStackProbe(*MF))
diff --git a/llvm/test/CodeGen/AArch64/stack-probing-shrink-wrap.mir b/llvm/test/CodeGen/AArch64/stack-probing-shrink-wrap.mir
new file mode 100644
index 00000000000000..83aa90d389a4a2
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/stack-probing-shrink-wrap.mir
@@ -0,0 +1,107 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# RUN: llc %s --start-before=shrink-wrap --stop-after=prologepilog -o - | FileCheck %s
+--- |
+  target triple = "aarch64-linux"
+
+  define void @f(i32 %n) #0 {
+  entry:
+    %a = alloca i8, i32 150000, align 8
+    %c0 = icmp sle i32 %n, 1
+    br i1 %c0, label %if.then1, label %exit
+
+  if.then1:                                         ; preds = %entry
+    %0 = icmp sle i32 %n, 1
+    %v = select i1 %0, i32 0, i32 1
+    call void @g(ptr %a, i32 %v)
+    br label %exit
+
+  exit:                                             ; preds = %if.then1, %entry
+    ret void
+  }
+
+  declare void @g(...)
+
+  attributes #0 = { nounwind "probe-stack"="inline-asm" "stack-probe-size"="4096" }
+
+...
+---
+name:            f
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$w0', virtual-reg: '' }
+frameInfo:
+  localFrameSize:  150000
+stack:
+  - { id: 0, name: a, type: default, offset: 0, size: 150000, alignment: 8,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -150000, debug-info-variable: '', debug-info-expression: '',
+      debug-info-location: '' }
+entry_values:    []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  ; CHECK-LABEL: name: f
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.3(0x80000000)
+  ; CHECK-NEXT:   liveins: $w0, $lr
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   early-clobber $sp = frame-setup STPXpre killed $fp, killed $lr, $sp, -2 :: (store (s64) into %stack.2), (store (s64) into %stack.1)
+  ; CHECK-NEXT:   $x9 = frame-setup SUBXri $sp, 36, 12
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3.entry:
+  ; CHECK-NEXT:   successors: %bb.4(0x40000000), %bb.3(0x40000000)
+  ; CHECK-NEXT:   liveins: $w0, $x9
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $sp = frame-setup SUBXri $sp, 1, 12
+  ; CHECK-NEXT:   frame-setup STRXui $xzr, $sp, 0
+  ; CHECK-NEXT:   $xzr = frame-setup SUBSXrx64 $sp, $x9, 24, implicit-def $nzcv
+  ; CHECK-NEXT:   frame-setup Bcc 1, %bb.3, implicit $nzcv
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.4.entry:
+  ; CHECK-NEXT:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK-NEXT:   liveins: $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $sp = frame-setup SUBXri $sp, 2544, 0
+  ; CHECK-NEXT:   frame-setup STRXui $xzr, $sp, 0
+  ; CHECK-NEXT:   dead $wzr = SUBSWri killed renamable $w0, 1, 0, implicit-def $nzcv
+  ; CHECK-NEXT:   Bcc 12, %bb.2, implicit $nzcv
+  ; CHECK-NEXT:   B %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1.if.then1:
+  ; CHECK-NEXT:   successors: %bb.2(0x80000000)
+  ; CHECK-NEXT:   liveins: $nzcv
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   renamable $w1 = CSINCWr $wzr, $wzr, 13, implicit killed $nzcv
+  ; CHECK-NEXT:   $x0 = ADDXri $sp, 0, 0
+  ; CHECK-NEXT:   BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit $w1, implicit-def $sp
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2.exit:
+  ; CHECK-NEXT:   $sp = frame-destroy ADDXri $sp, 36, 12
+  ; CHECK-NEXT:   $sp = frame-destroy ADDXri $sp, 2544, 0
+  ; CHECK-NEXT:   early-clobber $sp, $fp, $lr = frame-destroy LDPXpost $sp, 2 :: (load (s64) from %stack.2), (load (s64) from %stack.1)
+  ; CHECK-NEXT:   RET_ReallyLR
+  bb.0.entry:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $w0
+
+    dead $wzr = SUBSWri killed renamable $w0, 1, 0, implicit-def $nzcv
+    Bcc 12, %bb.2, implicit $nzcv
+    B %bb.1
+
+  bb.1.if.then1:
+    successors: %bb.2(0x80000000)
+    liveins: $nzcv
+
+    renamable $w1 = CSINCWr $wzr, $wzr, 13, implicit killed $nzcv
+    ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
+    $x0 = ADDXri %stack.0.a, 0, 0
+    BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit $w1, implicit-def $sp
+    ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
+
+  bb.2.exit:
+    RET_ReallyLR
+
+...

>From 21dae48d45f88ec8fe225cfcf6b08489a1d2c58c Mon Sep 17 00:00:00 2001
From: Momchil Velikov <momchil.velikov at arm.com>
Date: Tue, 20 Feb 2024 13:19:00 +0000
Subject: [PATCH 2/2] [fixup] Use `hasStackProbing`

The difference now is that the condition would trigger
for Windows, where stack probing does clobber flags.
---
 llvm/lib/Target/AArch64/AArch64FrameLowering.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
index c62e750a9b07b0..31b1661e73e363 100644
--- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp
@@ -1063,7 +1063,8 @@ bool AArch64FrameLowering::canUseAsPrologue(
 
   // Certain stack probing sequences might clobber flags, then we can't use
   // the block as a prologue if the flags register is a live-in.
-  if (TLI->hasInlineStackProbe(*MF) && MBB.isLiveIn(AArch64::NZCV))
+  if (MF->getInfo<AArch64FunctionInfo>()->hasStackProbing() &&
+      MBB.isLiveIn(AArch64::NZCV))
     return false;
 
   // Don't need a scratch register if we're not going to re-align the stack or



More information about the llvm-commits mailing list