Skip to content

Commit ca22613

Browse files
bors[bot]ptersilie
andauthored
48: Add function prologue info to stackmaps. r=vext01 a=ptersilie Co-authored-by: Lukas Diekmann <lukas.diekmann@gmail.com>
2 parents bd1d951 + d06f72f commit ca22613

File tree

5 files changed

+110
-61
lines changed

5 files changed

+110
-61
lines changed

llvm/include/llvm/CodeGen/StackMaps.h

+24-12
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ class StackMapOpers {
3838
enum { IDPos, NBytesPos };
3939

4040
private:
41-
const MachineInstr* MI;
41+
const MachineInstr *MI;
4242

4343
public:
4444
explicit StackMapOpers(const MachineInstr *MI);
@@ -204,6 +204,12 @@ class StackMaps {
204204
: Reg(Reg), DwarfRegNum(DwarfRegNum), Size(Size) {}
205205
};
206206

207+
// Stores callee-saved register info.
208+
struct CSR {
209+
int Reg = 0;
210+
int Offset = 0;
211+
};
212+
207213
// OpTypes are used to encode information about the following logical
208214
// operand (which may consist of several MachineOperands) for the
209215
// OpParser.
@@ -237,13 +243,19 @@ class StackMaps {
237243
using LiveVarsVec = SmallVector<LocationVec, 8>;
238244
using LiveOutVec = SmallVector<LiveOutReg, 8>;
239245
using ConstantPool = MapVector<uint64_t, uint64_t>;
246+
using CSRVec = SmallVector<CSR, 8>;
240247

241248
struct FunctionInfo {
242249
uint64_t StackSize = 0;
243250
uint64_t RecordCount = 1;
251+
bool HasFramePointer;
252+
CSRVec SpilledRegisters;
244253

245254
FunctionInfo() = default;
246-
explicit FunctionInfo(uint64_t StackSize) : StackSize(StackSize) {}
255+
explicit FunctionInfo(uint64_t StackSize, bool HasFramePointer,
256+
CSRVec SpilledRegisters)
257+
: StackSize(StackSize), HasFramePointer(HasFramePointer),
258+
SpilledRegisters(SpilledRegisters) {}
247259
};
248260

249261
struct CallsiteInfo {
@@ -265,16 +277,13 @@ class StackMaps {
265277
/// Generate a stackmap record for a stackmap instruction.
266278
///
267279
/// MI must be a raw STACKMAP, not a PATCHPOINT.
268-
void recordStackMap(const MCSymbol &L,
269-
const MachineInstr &MI);
280+
void recordStackMap(const MCSymbol &L, const MachineInstr &MI);
270281

271282
/// Generate a stackmap record for a patchpoint instruction.
272-
void recordPatchPoint(const MCSymbol &L,
273-
const MachineInstr &MI);
283+
void recordPatchPoint(const MCSymbol &L, const MachineInstr &MI);
274284

275285
/// Generate a stackmap record for a statepoint instruction.
276-
void recordStatepoint(const MCSymbol &L,
277-
const MachineInstr &MI);
286+
void recordStatepoint(const MCSymbol &L, const MachineInstr &MI);
278287

279288
/// If there is any stack map data, create a stack map section and serialize
280289
/// the map info into it. This clears the stack map data structures
@@ -294,6 +303,7 @@ class StackMaps {
294303
CallsiteInfoList CSInfos;
295304
ConstantPool ConstPool;
296305
FnInfoMap FnInfos;
306+
bool HasFramePointer;
297307

298308
MachineInstr::const_mop_iterator
299309
parseOperand(MachineInstr::const_mop_iterator MOI,
@@ -321,10 +331,9 @@ class StackMaps {
321331
/// STACKMAP, and PATCHPOINT the label is expected to immediately *preceed*
322332
/// lowering of the MI to MCInsts. For STATEPOINT, it expected to
323333
/// immediately *follow*. It's not clear this difference was intentional,
324-
/// but it exists today.
325-
void recordStackMapOpers(const MCSymbol &L,
326-
const MachineInstr &MI, uint64_t ID,
327-
MachineInstr::const_mop_iterator MOI,
334+
/// but it exists today.
335+
void recordStackMapOpers(const MCSymbol &L, const MachineInstr &MI,
336+
uint64_t ID, MachineInstr::const_mop_iterator MOI,
328337
MachineInstr::const_mop_iterator MOE,
329338
bool recordResult = false);
330339

@@ -340,6 +349,9 @@ class StackMaps {
340349
/// Emit the callsite info for each stackmap/patchpoint intrinsic call.
341350
void emitCallsiteEntries(MCStreamer &OS);
342351

352+
/// Emit information about the function prologue.
353+
void emitCSRInfo(MCStreamer &OS);
354+
343355
void print(raw_ostream &OS);
344356
void debug() { print(dbgs()); }
345357
};

llvm/include/llvm/Transforms/Yk/ControlPoint.h

+4
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@
1010
// The name of the new control point replacing the user's dummy control point.
1111
#define YK_NEW_CONTROL_POINT "__ykrt_control_point"
1212

13+
// The name of the function which reconstructs the stackframe and jumps to the
14+
// right instruction in AOT from where to continue.
15+
#define YK_RECONSTRUCT_FRAMES "__ykrt_reconstruct_frames"
16+
1317
namespace llvm {
1418
ModulePass *createYkControlPointPass();
1519
} // namespace llvm

llvm/lib/CodeGen/StackMaps.cpp

+54-15
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "llvm/CodeGen/MachineFunction.h"
1616
#include "llvm/CodeGen/MachineInstr.h"
1717
#include "llvm/CodeGen/MachineOperand.h"
18+
#include "llvm/CodeGen/TargetFrameLowering.h"
1819
#include "llvm/CodeGen/TargetOpcodes.h"
1920
#include "llvm/CodeGen/TargetRegisterInfo.h"
2021
#include "llvm/CodeGen/TargetSubtargetInfo.h"
@@ -54,10 +55,8 @@ static uint64_t getConstMetaVal(const MachineInstr &MI, unsigned Idx) {
5455
return MO.getImm();
5556
}
5657

57-
StackMapOpers::StackMapOpers(const MachineInstr *MI)
58-
: MI(MI) {
59-
assert(getVarIdx() <= MI->getNumOperands() &&
60-
"invalid stackmap definition");
58+
StackMapOpers::StackMapOpers(const MachineInstr *MI) : MI(MI) {
59+
assert(getVarIdx() <= MI->getNumOperands() && "invalid stackmap definition");
6160
}
6261

6362
PatchPointOpers::PatchPointOpers(const MachineInstr *MI)
@@ -81,11 +80,10 @@ unsigned PatchPointOpers::getNextScratchIdx(unsigned StartIdx) const {
8180

8281
// Find the next scratch register (implicit def and early clobber)
8382
unsigned ScratchIdx = StartIdx, e = MI->getNumOperands();
84-
while (ScratchIdx < e &&
85-
!(MI->getOperand(ScratchIdx).isReg() &&
86-
MI->getOperand(ScratchIdx).isDef() &&
87-
MI->getOperand(ScratchIdx).isImplicit() &&
88-
MI->getOperand(ScratchIdx).isEarlyClobber()))
83+
while (ScratchIdx < e && !(MI->getOperand(ScratchIdx).isReg() &&
84+
MI->getOperand(ScratchIdx).isDef() &&
85+
MI->getOperand(ScratchIdx).isImplicit() &&
86+
MI->getOperand(ScratchIdx).isEarlyClobber()))
8987
++ScratchIdx;
9088

9189
assert(ScratchIdx != e && "No scratch register available");
@@ -274,8 +272,8 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
274272
if (SubRegIdx)
275273
Offset = TRI->getSubRegIdxOffset(SubRegIdx);
276274

277-
Locs.emplace_back(Location::Register, TRI->getSpillSize(*RC),
278-
DwarfRegNum, Offset);
275+
Locs.emplace_back(Location::Register, TRI->getSpillSize(*RC), DwarfRegNum,
276+
Offset);
279277
return ++MOI;
280278
}
281279

@@ -570,6 +568,7 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
570568

571569
// Record the stack size of the current function and update callsite count.
572570
const MachineFrameInfo &MFI = AP.MF->getFrameInfo();
571+
const TargetRegisterInfo *TRI = AP.MF->getSubtarget().getRegisterInfo();
573572
const TargetRegisterInfo *RegInfo = AP.MF->getSubtarget().getRegisterInfo();
574573
bool HasDynamicFrameSize =
575574
MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(*(AP.MF));
@@ -578,17 +577,30 @@ void StackMaps::recordStackMapOpers(const MCSymbol &MILabel,
578577
auto CurrentIt = FnInfos.find(AP.CurrentFnSym);
579578
if (CurrentIt != FnInfos.end())
580579
CurrentIt->second.RecordCount++;
581-
else
582-
FnInfos.insert(std::make_pair(AP.CurrentFnSym, FunctionInfo(FrameSize)));
580+
else {
581+
// Collect callee-saved-register spills.
582+
CSRVec CSRInfo;
583+
std::vector<CalleeSavedInfo> CSI = MFI.getCalleeSavedInfo();
584+
for (auto cs : CSI) {
585+
int dwreg = getDwarfRegNum(cs.getReg(), TRI);
586+
CSR Entry = {dwreg, cs.getFrameIdx()};
587+
CSRInfo.push_back(Entry);
588+
}
589+
// Check whether this function pushed a frame pointer.
590+
const TargetFrameLowering *TFL = AP.MF->getSubtarget().getFrameLowering();
591+
bool HasFramePointer = TFL->hasFP(*(AP.MF));
592+
FnInfos.insert(std::make_pair(
593+
AP.CurrentFnSym, FunctionInfo(FrameSize, HasFramePointer, CSRInfo)));
594+
}
583595
}
584596

585597
void StackMaps::recordStackMap(const MCSymbol &L, const MachineInstr &MI) {
586598
assert(MI.getOpcode() == TargetOpcode::STACKMAP && "expected stackmap");
587599

588600
StackMapOpers opers(&MI);
589601
const int64_t ID = MI.getOperand(PatchPointOpers::IDPos).getImm();
590-
recordStackMapOpers(L, MI, ID, std::next(MI.operands_begin(),
591-
opers.getVarIdx()),
602+
recordStackMapOpers(L, MI, ID,
603+
std::next(MI.operands_begin(), opers.getVarIdx()),
592604
MI.operands_end());
593605
}
594606

@@ -784,6 +796,32 @@ void StackMaps::emitCallsiteEntries(MCStreamer &OS) {
784796
}
785797
}
786798

799+
// Emit information about the function prologue: pushed frame pointer and
800+
// callee-saved registers.
801+
// CSRInfo[NumFunctions] {
802+
// uint8 : HasFramePtr
803+
// uint8 : Padding
804+
// uint32 : NumSpills
805+
//
806+
// Spills[NumSpills] {
807+
// uint16 : DwarfRegister
808+
// uint16 : Padding
809+
// int32 : Offset
810+
// }
811+
// }
812+
void StackMaps::emitCSRInfo(MCStreamer &OS) {
813+
for (auto const &FR : FnInfos) {
814+
OS.emitInt8(FR.second.HasFramePointer);
815+
OS.emitInt8(0); // Padding
816+
OS.emitInt32(FR.second.SpilledRegisters.size());
817+
for (auto &Ent : FR.second.SpilledRegisters) {
818+
OS.emitInt16(Ent.Reg);
819+
OS.emitInt16(0); // Padding
820+
OS.emitInt32(Ent.Offset);
821+
}
822+
}
823+
}
824+
787825
/// Serialize the stackmap data.
788826
void StackMaps::serializeToStackMapSection() {
789827
(void)WSMP;
@@ -812,6 +850,7 @@ void StackMaps::serializeToStackMapSection() {
812850
emitFunctionFrameRecords(OS);
813851
emitConstantPoolEntries(OS);
814852
emitCallsiteEntries(OS);
853+
emitCSRInfo(OS);
815854
OS.addBlankLine();
816855

817856
// Clean up.

llvm/lib/Transforms/Yk/ControlPoint.cpp

+27-33
Original file line numberDiff line numberDiff line change
@@ -162,25 +162,11 @@ class YkControlPoint : public ModulePass {
162162
->getType();
163163

164164
// Create the new control point, which is of the form:
165-
// bool new_control_point(YkMT*, YkLocation*, CtrlPointVars*,
166-
// ReturnValue*)
167-
// If the return type of the control point's caller is void (i.e. if a
168-
// function f calls yk_control_point and f's return type is void), create
169-
// an Int1 pointer as a dummy. We have to pass something as the yk_stopgap
170-
// signature expects a pointer, even if its never used.
171-
Type *ReturnTy = Caller->getReturnType();
172-
Type *ReturnPtrTy;
173-
if (ReturnTy->isVoidTy()) {
174-
// Create dummy pointer which we pass in but which is never written to.
175-
ReturnPtrTy = Type::getInt1Ty(Context);
176-
} else {
177-
ReturnPtrTy = ReturnTy;
178-
}
179-
FunctionType *FType =
180-
FunctionType::get(Type::getInt1Ty(Context),
181-
{YkMTTy, YkLocTy, CtrlPointVarsTy->getPointerTo(),
182-
ReturnPtrTy->getPointerTo()},
183-
false);
165+
// void* new_control_point(YkMT*, YkLocation*, CtrlPointVars*, void*)
166+
PointerType *VoidPtr = PointerType::get(Context, 0);
167+
FunctionType *FType = FunctionType::get(
168+
VoidPtr, {YkMTTy, YkLocTy, CtrlPointVarsTy->getPointerTo(), VoidPtr},
169+
false);
184170
Function *NF = Function::Create(FType, GlobalVariable::ExternalLinkage,
185171
YK_NEW_CONTROL_POINT, M);
186172

@@ -190,10 +176,6 @@ class YkControlPoint : public ModulePass {
190176
IRBuilder<> Builder(Caller->getEntryBlock().getFirstNonPHI());
191177
Value *InputStruct = Builder.CreateAlloca(CtrlPointVarsTy, 0, "");
192178

193-
// Also at the top, generate storage for the interpreted return of the
194-
// control points caller.
195-
Value *ReturnPtr = Builder.CreateAlloca(ReturnPtrTy, 0, "");
196-
197179
Builder.SetInsertPoint(OldCtrlPointCall);
198180
unsigned LvIdx = 0;
199181
for (Value *LV : LiveVals) {
@@ -205,11 +187,17 @@ class YkControlPoint : public ModulePass {
205187
LvIdx++;
206188
}
207189

190+
// Create a call to the llvm.frameaddress intrinsic, which we pass into the
191+
// control point. This is later required for stack reconstruction.
192+
Function *FrameAddress =
193+
Intrinsic::getDeclaration(&M, Intrinsic::frameaddress, {VoidPtr});
194+
Value *FAddr = Builder.CreateCall(FrameAddress, {Builder.getInt32(0)});
195+
208196
// Insert call to the new control point.
209197
Instruction *NewCtrlPointCallInst = Builder.CreateCall(
210198
NF, {OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_MT_IDX),
211199
OldCtrlPointCall->getArgOperand(YK_CONTROL_POINT_ARG_LOC_IDX),
212-
InputStruct, ReturnPtr});
200+
InputStruct, FAddr});
213201

214202
// Once the control point returns we need to extract the (potentially
215203
// mutated) values from the returned YkCtrlPointStruct and reassign them to
@@ -238,14 +226,14 @@ class YkControlPoint : public ModulePass {
238226
// Create the new exit block.
239227
BasicBlock *ExitBB = BasicBlock::Create(Context, "", Caller);
240228
Builder.SetInsertPoint(ExitBB);
241-
// YKFIXME: We need to return the value of interpreted return and the return
242-
// type must be that of the control point's caller.
243-
if (ReturnTy->isVoidTy()) {
244-
Builder.CreateRetVoid();
245-
} else {
246-
Value *ReturnValue = Builder.CreateLoad(ReturnTy, ReturnPtr);
247-
Builder.CreateRet(ReturnValue);
248-
}
229+
230+
// Create call to frame reconstructor.
231+
FunctionType *YKFRType =
232+
FunctionType::get(Type::getVoidTy(Context), {VoidPtr}, false);
233+
Function *YKFR = Function::Create(YKFRType, GlobalVariable::ExternalLinkage,
234+
YK_RECONSTRUCT_FRAMES, M);
235+
Builder.CreateCall(YKFR, {NewCtrlPointCallInst});
236+
Builder.CreateUnreachable();
249237

250238
// To do so we need to first split up the current block and then
251239
// insert a conditional branch that either continues or returns.
@@ -256,7 +244,13 @@ class YkControlPoint : public ModulePass {
256244
Instruction &OldBr = BB->back();
257245
OldBr.eraseFromParent();
258246
Builder.SetInsertPoint(BB);
259-
Builder.CreateCondBr(NewCtrlPointCallInst, ExitBB, ContBB);
247+
// The return value of the control point tells us whether a guard has failed
248+
// or not (i.e. anything other than a nullptr means a guard has failed). If
249+
// it has jump to `ExitBB` which calls the code that copies over the new
250+
// stack from the given pointer.
251+
Value *HasGuardFailed = Builder.CreateICmpEQ(
252+
NewCtrlPointCallInst, ConstantPointerNull::get(VoidPtr));
253+
Builder.CreateCondBr(HasGuardFailed, ContBB, ExitBB);
260254

261255
#ifndef NDEBUG
262256
// Our pass runs after LLVM normally does its verify pass. In debug builds

llvm/test/CodeGen/X86/statepoint-stackmap-size.ll

+1-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
; spaces) starting with a `.` follow (e.g. ` .byte`).
55
;
66
; CHECK: .section .llvm_stackmaps,{{.*$}}
7-
; CHECK-NEXT:{{(.+$[[:space:]]){51}[[:space:]]}}
7+
; CHECK-NEXT:{{(.+$[[:space:]]){54}[[:space:]]}}
88
; CHECK-NOT:{{.|[[:space:]]}}
99

1010
target triple = "x86_64-pc-linux-gnu"

0 commit comments

Comments
 (0)