Skip to content

Commit c3d9f31

Browse files
authored
[AArch64][SVE] Avoid using the SVE stack size before it is determined (#147561)
This patch adds asserts to `getStackSizeSVE()` and `getSVECalleeSavedStackSize()` to check they are only called after the SVE stack size has been determined. This patch then fixes issues in three places: * `AArch64FrameLowering::homogeneousPrologEpilog()` - This function is called before callee saves or SVE stack sizes have been determined - This check has been updated to use `isLikelyToHaveSVEStack()` - `isLikelyToHaveSVEStack()` conservatively returns if a function is likely to have an SVE stack * `produceCompactUnwindFrame()` - This function checked the SVE CS stack size before callee-saves had been determined - This has been replaced with a more conservative `isSVECC()` check * `AArch64FrameLowering::resolveFrameOffsetReference()` - This was hit by some post-PEI MIR tests - This case was fixed by adding "stackSizeSVE" to the YAML for `AArch64FunctionInfo`
1 parent d3d4066 commit c3d9f31

File tree

5 files changed

+43
-4
lines changed

5 files changed

+43
-4
lines changed

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,26 @@ static Register findScratchNonCalleeSaveRegister(MachineBasicBlock *MBB,
335335
bool HasCall = false);
336336
static bool requiresSaveVG(const MachineFunction &MF);
337337

338+
// Conservatively, returns true if the function is likely to have an SVE vectors
339+
// on the stack. This function is safe to be called before callee-saves or
340+
// object offsets have been determined.
341+
static bool isLikelyToHaveSVEStack(MachineFunction &MF) {
342+
auto *AFI = MF.getInfo<AArch64FunctionInfo>();
343+
if (AFI->isSVECC())
344+
return true;
345+
346+
if (AFI->hasCalculatedStackSizeSVE())
347+
return bool(getSVEStackSize(MF));
348+
349+
const MachineFrameInfo &MFI = MF.getFrameInfo();
350+
for (int FI = MFI.getObjectIndexBegin(); FI < MFI.getObjectIndexEnd(); FI++) {
351+
if (MFI.getStackID(FI) == TargetStackID::ScalableVector)
352+
return true;
353+
}
354+
355+
return false;
356+
}
357+
338358
/// Returns true if a homogeneous prolog or epilog code can be emitted
339359
/// for the size optimization. If possible, a frame helper call is injected.
340360
/// When Exit block is given, this check is for epilog.
@@ -350,8 +370,9 @@ bool AArch64FrameLowering::homogeneousPrologEpilog(
350370
// TODO: Window is supported yet.
351371
if (needsWinCFI(MF))
352372
return false;
373+
353374
// TODO: SVE is not supported yet.
354-
if (getSVEStackSize(MF))
375+
if (isLikelyToHaveSVEStack(MF))
355376
return false;
356377

357378
// Bail on stack adjustment needed on return for simplicity.
@@ -3039,7 +3060,7 @@ static bool produceCompactUnwindFrame(MachineFunction &MF) {
30393060
!(Subtarget.getTargetLowering()->supportSwiftError() &&
30403061
Attrs.hasAttrSomewhere(Attribute::SwiftError)) &&
30413062
MF.getFunction().getCallingConv() != CallingConv::SwiftTail &&
3042-
!requiresSaveVG(MF) && AFI->getSVECalleeSavedStackSize() == 0;
3063+
!requiresSaveVG(MF) && !AFI->isSVECC();
30433064
}
30443065

30453066
static bool invalidateWindowsRegisterPairing(unsigned Reg1, unsigned Reg2,

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ using namespace llvm;
2525

2626
yaml::AArch64FunctionInfo::AArch64FunctionInfo(
2727
const llvm::AArch64FunctionInfo &MFI)
28-
: HasRedZone(MFI.hasRedZone()) {}
28+
: HasRedZone(MFI.hasRedZone()),
29+
StackSizeSVE(MFI.hasCalculatedStackSizeSVE()
30+
? std::optional<uint64_t>(MFI.getStackSizeSVE())
31+
: std::nullopt) {}
2932

3033
void yaml::AArch64FunctionInfo::mappingImpl(yaml::IO &YamlIO) {
3134
MappingTraits<AArch64FunctionInfo>::mapping(YamlIO, *this);
@@ -35,6 +38,8 @@ void AArch64FunctionInfo::initializeBaseYamlFields(
3538
const yaml::AArch64FunctionInfo &YamlMFI) {
3639
if (YamlMFI.HasRedZone)
3740
HasRedZone = YamlMFI.HasRedZone;
41+
if (YamlMFI.StackSizeSVE)
42+
setStackSizeSVE(*YamlMFI.StackSizeSVE);
3843
}
3944

4045
static std::pair<bool, bool> GetSignReturnAddress(const Function &F) {

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
8282
unsigned CalleeSavedStackSize = 0;
8383
unsigned SVECalleeSavedStackSize = 0;
8484
bool HasCalleeSavedStackSize = false;
85+
bool HasSVECalleeSavedStackSize = false;
8586

8687
/// Number of TLS accesses using the special (combinable)
8788
/// _TLS_MODULE_BASE_ symbol.
@@ -306,7 +307,10 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
306307
StackSizeSVE = S;
307308
}
308309

309-
uint64_t getStackSizeSVE() const { return StackSizeSVE; }
310+
uint64_t getStackSizeSVE() const {
311+
assert(hasCalculatedStackSizeSVE());
312+
return StackSizeSVE;
313+
}
310314

311315
bool hasStackFrame() const { return HasStackFrame; }
312316
void setHasStackFrame(bool s) { HasStackFrame = s; }
@@ -400,8 +404,11 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
400404
// Saves the CalleeSavedStackSize for SVE vectors in 'scalable bytes'
401405
void setSVECalleeSavedStackSize(unsigned Size) {
402406
SVECalleeSavedStackSize = Size;
407+
HasSVECalleeSavedStackSize = true;
403408
}
404409
unsigned getSVECalleeSavedStackSize() const {
410+
assert(HasSVECalleeSavedStackSize &&
411+
"SVECalleeSavedStackSize has not been calculated");
405412
return SVECalleeSavedStackSize;
406413
}
407414

@@ -592,6 +599,7 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
592599
namespace yaml {
593600
struct AArch64FunctionInfo final : public yaml::MachineFunctionInfo {
594601
std::optional<bool> HasRedZone;
602+
std::optional<uint64_t> StackSizeSVE;
595603

596604
AArch64FunctionInfo() = default;
597605
AArch64FunctionInfo(const llvm::AArch64FunctionInfo &MFI);
@@ -603,6 +611,7 @@ struct AArch64FunctionInfo final : public yaml::MachineFunctionInfo {
603611
template <> struct MappingTraits<AArch64FunctionInfo> {
604612
static void mapping(IO &YamlIO, AArch64FunctionInfo &MFI) {
605613
YamlIO.mapOptional("hasRedZone", MFI.HasRedZone);
614+
YamlIO.mapOptional("stackSizeSVE", MFI.StackSizeSVE);
606615
}
607616
};
608617

llvm/test/DebugInfo/AArch64/asan-stack-vars.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,8 @@ frameInfo:
365365
stackProtector: '%stack.0.StackGuardSlot'
366366
maxCallFrameSize: 0
367367
localFrameSize: 144
368+
machineFunctionInfo:
369+
stackSizeSVE: 0
368370
stack:
369371
- { id: 0, name: StackGuardSlot, offset: -40, size: 8, alignment: 8,
370372
stack-id: default, local-offset: -8 }

llvm/test/DebugInfo/AArch64/compiler-gen-bbs-livedebugvalues.mir

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,8 @@ frameInfo:
6868
adjustsStack: true
6969
hasCalls: true
7070
maxCallFrameSize: 0
71+
machineFunctionInfo:
72+
stackSizeSVE: 0
7173
stack:
7274
- { id: 0, type: spill-slot, offset: -20, size: 4, alignment: 4, stack-id: default }
7375
- { id: 1, type: spill-slot, offset: -8, size: 8, alignment: 8, stack-id: default,

0 commit comments

Comments
 (0)