Skip to content

Commit 603d90a

Browse files
momchil-velikovmomchil-velikov
committed
Fix a crash when the probing instruction to replace is last in block
1 parent 94f8691 commit 603d90a

File tree

3 files changed

+174
-53
lines changed

3 files changed

+174
-53
lines changed

‎llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

+25-44
Original file line numberDiff line numberDiff line change
@@ -4276,7 +4276,7 @@ AArch64FrameLowering::inlineStackProbeLoopExactMultiple(
42764276
return ExitMBB->begin();
42774277
}
42784278

4279-
MachineBasicBlock::iterator AArch64FrameLowering::inlineStackProbeFixed(
4279+
void AArch64FrameLowering::inlineStackProbeFixed(
42804280
MachineBasicBlock::iterator MBBI, Register ScratchReg, int64_t FrameSize,
42814281
StackOffset CFAOffset) const {
42824282
MachineBasicBlock *MBB = MBBI->getParent();
@@ -4353,54 +4353,35 @@ MachineBasicBlock::iterator AArch64FrameLowering::inlineStackProbeFixed(
43534353
.setMIFlags(MachineInstr::FrameSetup);
43544354
}
43554355
}
4356-
4357-
MachineBasicBlock::iterator Next = std::next(MBBI);
4358-
return Next;
4359-
}
4360-
4361-
MachineBasicBlock::iterator AArch64FrameLowering::inlineStackProbeFixed(
4362-
MachineBasicBlock::iterator MBBI) const {
4363-
4364-
Register ScratchReg = MBBI->getOperand(0).getReg();
4365-
int64_t FrameSize = MBBI->getOperand(1).getImm();
4366-
StackOffset CFAOffset = StackOffset::get(MBBI->getOperand(2).getImm(),
4367-
MBBI->getOperand(3).getImm());
4368-
4369-
MachineBasicBlock::iterator NextInst =
4370-
inlineStackProbeFixed(MBBI, ScratchReg, FrameSize, CFAOffset);
4371-
4372-
MBBI->eraseFromParent();
4373-
return NextInst;
4374-
}
4375-
4376-
MachineBasicBlock::iterator AArch64FrameLowering::inlineStackProbeVar(
4377-
MachineBasicBlock::iterator MBBI) const {
4378-
MachineBasicBlock &MBB = *MBBI->getParent();
4379-
MachineFunction &MF = *MBB.getParent();
4380-
const AArch64InstrInfo *TII =
4381-
MF.getSubtarget<AArch64Subtarget>().getInstrInfo();
4382-
4383-
DebugLoc DL = MBB.findDebugLoc(MBBI);
4384-
Register TargetReg = MBBI->getOperand(0).getReg();
4385-
4386-
MachineBasicBlock::iterator NextInst =
4387-
TII->probedStackAlloc(MBBI, TargetReg, true);
4388-
4389-
MBBI->eraseFromParent();
4390-
return NextInst;
43914356
}
43924357

43934358
void AArch64FrameLowering::inlineStackProbe(MachineFunction &MF,
43944359
MachineBasicBlock &MBB) const {
4395-
for (auto MBBI = MBB.begin(), E = MBB.end(); MBBI != E;) {
4396-
if (MBBI->getOpcode() == AArch64::PROBED_STACKALLOC) {
4397-
MBBI = inlineStackProbeFixed(MBBI);
4398-
E = MBBI->getParent()->end();
4399-
} else if (MBBI->getOpcode() == AArch64::PROBED_STACKALLOC_VAR) {
4400-
MBBI = inlineStackProbeVar(MBBI);
4401-
E = MBBI->getParent()->end();
4360+
// Get the instructions that need to be replaced. We emit at most two of
4361+
// these. Remember them in order to avoid complications coming from the need
4362+
// to traverse the block while potentially creating more blocks.
4363+
SmallVector<MachineInstr *, 4> ToReplace;
4364+
for (MachineInstr &MI : MBB)
4365+
if (MI.getOpcode() == AArch64::PROBED_STACKALLOC ||
4366+
MI.getOpcode() == AArch64::PROBED_STACKALLOC_VAR)
4367+
ToReplace.push_back(&MI);
4368+
4369+
for (MachineInstr *MI : ToReplace) {
4370+
if (MI->getOpcode() == AArch64::PROBED_STACKALLOC) {
4371+
Register ScratchReg = MI->getOperand(0).getReg();
4372+
int64_t FrameSize = MI->getOperand(1).getImm();
4373+
StackOffset CFAOffset = StackOffset::get(MI->getOperand(2).getImm(),
4374+
MI->getOperand(3).getImm());
4375+
inlineStackProbeFixed(MI->getIterator(), ScratchReg, FrameSize,
4376+
CFAOffset);
44024377
} else {
4403-
++MBBI;
4378+
assert(MI->getOpcode() == AArch64::PROBED_STACKALLOC_VAR &&
4379+
"Stack probe pseudo-instruction expected");
4380+
const AArch64InstrInfo *TII =
4381+
MI->getMF()->getSubtarget<AArch64Subtarget>().getInstrInfo();
4382+
Register TargetReg = MI->getOperand(0).getReg();
4383+
(void)TII->probedStackAlloc(MI->getIterator(), TargetReg, true);
44044384
}
4385+
MI->eraseFromParent();
44054386
}
44064387
}

‎llvm/lib/Target/AArch64/AArch64FrameLowering.h

+3-9
Original file line numberDiff line numberDiff line change
@@ -164,15 +164,9 @@ class AArch64FrameLowering : public TargetFrameLowering {
164164
void inlineStackProbe(MachineFunction &MF,
165165
MachineBasicBlock &PrologueMBB) const override;
166166

167-
MachineBasicBlock::iterator
168-
inlineStackProbeFixed(MachineBasicBlock::iterator MBBI, Register ScratchReg,
169-
int64_t FrameSize, StackOffset CFAOffset) const;
170-
171-
MachineBasicBlock::iterator
172-
inlineStackProbeFixed(MachineBasicBlock::iterator MBBI) const;
173-
174-
MachineBasicBlock::iterator
175-
inlineStackProbeVar(MachineBasicBlock::iterator MBBI) const;
167+
void inlineStackProbeFixed(MachineBasicBlock::iterator MBBI,
168+
Register ScratchReg, int64_t FrameSize,
169+
StackOffset CFAOffset) const;
176170

177171
MachineBasicBlock::iterator
178172
inlineStackProbeLoopExactMultiple(MachineBasicBlock::iterator MBBI,

‎llvm/test/CodeGen/AArch64/stack-probing-last-in-block.mir

+146
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
2+
# RUN: llc -run-pass=prologepilog %s -o - | FileCheck %s
3+
# Regression test for a crash when the probing instruction
4+
# to replace is last in the block.
5+
--- |
6+
source_filename = "tt.ll"
7+
target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
8+
target triple = "aarch64-linux"
9+
10+
declare i1 @g(ptr)
11+
12+
define void @f(ptr %out) #0 {
13+
entry:
14+
%p = alloca i32, i32 50000, align 4
15+
br label %loop
16+
17+
loop: ; preds = %loop, %entry
18+
%c = call i1 @g(ptr %p)
19+
br i1 %c, label %loop, label %exit
20+
21+
exit: ; preds = %loop
22+
ret void
23+
}
24+
25+
attributes #0 = { uwtable "frame-pointer"="none" "probe-stack"="inline-asm" "target-features"="+sve" }
26+
27+
...
28+
---
29+
name: f
30+
alignment: 4
31+
exposesReturnsTwice: false
32+
legalized: false
33+
regBankSelected: false
34+
selected: false
35+
failedISel: false
36+
tracksRegLiveness: true
37+
hasWinCFI: false
38+
callsEHReturn: false
39+
callsUnwindInit: false
40+
hasEHCatchret: false
41+
hasEHScopes: false
42+
hasEHFunclets: false
43+
isOutlined: false
44+
debugInstrRef: false
45+
failsVerification: false
46+
tracksDebugUserValues: true
47+
registers: []
48+
liveins: []
49+
frameInfo:
50+
isFrameAddressTaken: false
51+
isReturnAddressTaken: false
52+
hasStackMap: false
53+
hasPatchPoint: false
54+
stackSize: 0
55+
offsetAdjustment: 0
56+
maxAlignment: 4
57+
adjustsStack: true
58+
hasCalls: true
59+
stackProtector: ''
60+
functionContext: ''
61+
maxCallFrameSize: 0
62+
cvBytesOfCalleeSavedRegisters: 0
63+
hasOpaqueSPAdjustment: false
64+
hasVAStart: false
65+
hasMustTailInVarArgFunc: false
66+
hasTailCall: false
67+
localFrameSize: 200000
68+
savePoint: ''
69+
restorePoint: ''
70+
fixedStack: []
71+
stack:
72+
- { id: 0, name: p, type: default, offset: 0, size: 200000, alignment: 4,
73+
stack-id: default, callee-saved-register: '', callee-saved-restored: true,
74+
local-offset: -200000, debug-info-variable: '', debug-info-expression: '',
75+
debug-info-location: '' }
76+
entry_values: []
77+
callSites: []
78+
debugValueSubstitutions: []
79+
constants: []
80+
machineFunctionInfo: {}
81+
body: |
82+
; CHECK-LABEL: name: f
83+
; CHECK: bb.0.entry:
84+
; CHECK-NEXT: successors: %bb.3(0x80000000)
85+
; CHECK-NEXT: liveins: $lr, $fp
86+
; CHECK-NEXT: {{ $}}
87+
; CHECK-NEXT: early-clobber $sp = frame-setup STPXpre killed $fp, killed $lr, $sp, -2 :: (store (s64) into %stack.2), (store (s64) into %stack.1)
88+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 16
89+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w30, -8
90+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION offset $w29, -16
91+
; CHECK-NEXT: $x9 = frame-setup SUBXri $sp, 48, 12
92+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa $w9, 196624
93+
; CHECK-NEXT: {{ $}}
94+
; CHECK-NEXT: bb.3.entry:
95+
; CHECK-NEXT: successors: %bb.4(0x40000000), %bb.3(0x40000000)
96+
; CHECK-NEXT: liveins: $x9
97+
; CHECK-NEXT: {{ $}}
98+
; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 1, 12
99+
; CHECK-NEXT: frame-setup STRXui $xzr, $sp, 0
100+
; CHECK-NEXT: $xzr = frame-setup SUBSXrx64 $sp, $x9, 24, implicit-def $nzcv
101+
; CHECK-NEXT: frame-setup Bcc 1, %bb.3, implicit $nzcv
102+
; CHECK-NEXT: {{ $}}
103+
; CHECK-NEXT: bb.4.entry:
104+
; CHECK-NEXT: successors: %bb.1(0x80000000)
105+
; CHECK-NEXT: {{ $}}
106+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_register $wsp
107+
; CHECK-NEXT: $sp = frame-setup SUBXri $sp, 3392, 0
108+
; CHECK-NEXT: frame-setup CFI_INSTRUCTION def_cfa_offset 200016
109+
; CHECK-NEXT: frame-setup STRXui $xzr, $sp, 0
110+
; CHECK-NEXT: {{ $}}
111+
; CHECK-NEXT: bb.1.loop:
112+
; CHECK-NEXT: successors: %bb.1(0x7c000000), %bb.2(0x04000000)
113+
; CHECK-NEXT: {{ $}}
114+
; CHECK-NEXT: $x0 = ADDXri $sp, 0, 0
115+
; CHECK-NEXT: BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit-def $sp, implicit-def $w0
116+
; CHECK-NEXT: TBNZW killed renamable $w0, 0, %bb.1
117+
; CHECK-NEXT: B %bb.2
118+
; CHECK-NEXT: {{ $}}
119+
; CHECK-NEXT: bb.2.exit:
120+
; CHECK-NEXT: $sp = frame-destroy ADDXri $sp, 48, 12
121+
; CHECK-NEXT: frame-destroy CFI_INSTRUCTION def_cfa_offset 3408
122+
; CHECK-NEXT: $sp = frame-destroy ADDXri $sp, 3392, 0
123+
; CHECK-NEXT: frame-destroy CFI_INSTRUCTION def_cfa_offset 16
124+
; CHECK-NEXT: early-clobber $sp, $fp, $lr = frame-destroy LDPXpost $sp, 2 :: (load (s64) from %stack.2), (load (s64) from %stack.1)
125+
; CHECK-NEXT: frame-destroy CFI_INSTRUCTION def_cfa_offset 0
126+
; CHECK-NEXT: frame-destroy CFI_INSTRUCTION restore $w30
127+
; CHECK-NEXT: frame-destroy CFI_INSTRUCTION restore $w29
128+
; CHECK-NEXT: RET_ReallyLR
129+
bb.0.entry:
130+
successors: %bb.1(0x80000000)
131+
132+
133+
bb.1.loop:
134+
successors: %bb.1(0x7c000000), %bb.2(0x04000000)
135+
136+
ADJCALLSTACKDOWN 0, 0, implicit-def dead $sp, implicit $sp
137+
$x0 = ADDXri %stack.0.p, 0, 0
138+
BL @g, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit $x0, implicit-def $sp, implicit-def $w0
139+
ADJCALLSTACKUP 0, 0, implicit-def dead $sp, implicit $sp
140+
TBNZW killed renamable $w0, 0, %bb.1
141+
B %bb.2
142+
143+
bb.2.exit:
144+
RET_ReallyLR
145+
146+
...

0 commit comments

Comments
 (0)