From 5309356908c823b049fb7fe9c267e99a4bfa5d5d Mon Sep 17 00:00:00 2001 From: Ben Vanik Date: Thu, 6 Feb 2014 21:53:31 -0800 Subject: [PATCH] Fixing tail calls in the jit. --- .../x64/lowering/lowering_sequences.cc | 31 +++++++++++++++++-- src/alloy/backend/x64/lowering/tracers.cc | 6 ++-- src/alloy/backend/x64/x64_emitter.cc | 4 +++ src/alloy/backend/x64/x64_function.cc | 2 +- src/alloy/backend/x64/x64_thunk_emitter.cc | 2 ++ src/alloy/backend/x64/x64_thunk_emitter.h | 8 ++++- 6 files changed, 47 insertions(+), 6 deletions(-) diff --git a/src/alloy/backend/x64/lowering/lowering_sequences.cc b/src/alloy/backend/x64/lowering/lowering_sequences.cc index bb59ca222..38007aed5 100644 --- a/src/alloy/backend/x64/lowering/lowering_sequences.cc +++ b/src/alloy/backend/x64/lowering/lowering_sequences.cc @@ -245,17 +245,30 @@ void IssueCall(X64Emitter& e, FunctionInfo* symbol_info, uint32_t flags) { // Actually jump/call to rax. if (flags & CALL_TAIL) { + // Pass the callers return address over. + e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_RET_ADDR]); + e.add(e.rsp, (uint32_t)e.stack_size()); e.jmp(e.rax); } else { + // Return address is from the previous SET_RETURN_ADDRESS. + e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR]); + e.call(e.rax); } } void IssueCallIndirect(X64Emitter& e, Value* target, uint32_t flags) { - // Resolve address to the function to call and store in rax. - // TODO(benvanik): caching/etc. For now this makes debugging easier. Reg64 r; e.BeginOp(target, r, 0); + + // Check if return. + if (flags & CALL_POSSIBLE_RETURN) { + e.cmp(r.cvt32(), e.dword[e.rsp + StackLayout::GUEST_RET_ADDR]); + e.je("epilog", CodeGenerator::T_NEAR); + } + + // Resolve address to the function to call and store in rax. + // TODO(benvanik): caching/etc. For now this makes debugging easier. if (r != e.rdx) { e.mov(e.rdx, r); } @@ -264,9 +277,15 @@ void IssueCallIndirect(X64Emitter& e, Value* target, uint32_t flags) { // Actually jump/call to rax. if (flags & CALL_TAIL) { + // Pass the callers return address over. + e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_RET_ADDR]); + e.add(e.rsp, (uint32_t)e.stack_size()); e.jmp(e.rax); } else { + // Return address is from the previous SET_RETURN_ADDRESS. + e.mov(e.rdx, e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR]); + e.call(e.rax); } } @@ -434,6 +453,14 @@ table->AddSequence(OPCODE_RETURN_TRUE, [](X64Emitter& e, Instr*& i) { return true; }); +table->AddSequence(OPCODE_SET_RETURN_ADDRESS, [](X64Emitter& e, Instr*& i) { + XEASSERT(i->src1.value->IsConstant()); + e.mov(e.qword[e.rsp + StackLayout::GUEST_CALL_RET_ADDR], + i->src1.value->AsUint64()); + i = e.Advance(i); + return true; +}); + // -------------------------------------------------------------------------- // Branches // -------------------------------------------------------------------------- diff --git a/src/alloy/backend/x64/lowering/tracers.cc b/src/alloy/backend/x64/lowering/tracers.cc index 0d7975847..b4c0ae74c 100644 --- a/src/alloy/backend/x64/lowering/tracers.cc +++ b/src/alloy/backend/x64/lowering/tracers.cc @@ -29,10 +29,12 @@ namespace lowering { #define DFLUSH() #define DPRINT +#define TARGET_THREAD 1 + #define IFLUSH() fflush(stdout) -#define IPRINT if (thread_state->thread_id() == 1) printf +#define IPRINT if (thread_state->thread_id() == TARGET_THREAD) printf #define DFLUSH() fflush(stdout) -#define DPRINT DFLUSH(); if (thread_state->thread_id() == 1) printf +#define DPRINT DFLUSH(); if (thread_state->thread_id() == TARGET_THREAD) printf void TraceString(void* raw_context, const char* str) { diff --git a/src/alloy/backend/x64/x64_emitter.cc b/src/alloy/backend/x64/x64_emitter.cc index c8ddaddb4..02a1aa132 100644 --- a/src/alloy/backend/x64/x64_emitter.cc +++ b/src/alloy/backend/x64/x64_emitter.cc @@ -158,6 +158,10 @@ int X64Emitter::Emit(HIRBuilder* builder, size_t& out_stack_size) { if (emit_prolog) { sub(rsp, (uint32_t)stack_size); mov(qword[rsp + StackLayout::GUEST_RCX_HOME], rcx); + mov(qword[rsp + StackLayout::GUEST_RET_ADDR], rdx); + mov(qword[rsp + StackLayout::GUEST_CALL_RET_ADDR], 0); + // ReloadRDX: + mov(rdx, qword[rcx + 8]); // membase } auto lowering_table = backend_->lowering_table(); diff --git a/src/alloy/backend/x64/x64_function.cc b/src/alloy/backend/x64/x64_function.cc index 0f6b4d12b..71452ac14 100644 --- a/src/alloy/backend/x64/x64_function.cc +++ b/src/alloy/backend/x64/x64_function.cc @@ -48,6 +48,6 @@ int X64Function::CallImpl(ThreadState* thread_state, uint64_t return_address) { thunk( machine_code_, thread_state->raw_context(), - thread_state->memory()->membase()); + (void*)return_address); return 0; } diff --git a/src/alloy/backend/x64/x64_thunk_emitter.cc b/src/alloy/backend/x64/x64_thunk_emitter.cc index 7fc6fab60..0e1922581 100644 --- a/src/alloy/backend/x64/x64_thunk_emitter.cc +++ b/src/alloy/backend/x64/x64_thunk_emitter.cc @@ -34,6 +34,7 @@ HostToGuestThunk X64ThunkEmitter::EmitHostToGuestThunk() { const size_t stack_size = StackLayout::THUNK_STACK_SIZE; // rsp + 0 = return address + mov(qword[rsp + 8 * 3], r8); mov(qword[rsp + 8 * 2], rdx); mov(qword[rsp + 8 * 1], rcx); sub(rsp, stack_size); @@ -88,6 +89,7 @@ HostToGuestThunk X64ThunkEmitter::EmitHostToGuestThunk() { add(rsp, stack_size); mov(rcx, qword[rsp + 8 * 1]); mov(rdx, qword[rsp + 8 * 2]); + mov(r8, qword[rsp + 8 * 3]); ret(); void* fn = Emplace(stack_size); diff --git a/src/alloy/backend/x64/x64_thunk_emitter.h b/src/alloy/backend/x64/x64_thunk_emitter.h index a9f27650b..ae9c7b967 100644 --- a/src/alloy/backend/x64/x64_thunk_emitter.h +++ b/src/alloy/backend/x64/x64_thunk_emitter.h @@ -104,6 +104,10 @@ namespace x64 { * +------------------+ * | rcx / context | rsp + 64 * +------------------+ + * | guest ret addr | rsp + 72 + * +------------------+ + * | call ret addr | rsp + 80 + * +------------------+ * ... locals ... * +------------------+ * | (return address) | @@ -115,8 +119,10 @@ class StackLayout { public: const static size_t THUNK_STACK_SIZE = 120; - const static size_t GUEST_STACK_SIZE = 72; + const static size_t GUEST_STACK_SIZE = 88; const static size_t GUEST_RCX_HOME = 64; + const static size_t GUEST_RET_ADDR = 72; + const static size_t GUEST_CALL_RET_ADDR = 80; };