Reuse the epilogue generation functions for tail call emission

Instead of building and copying the new frame over the old one, make use
of the frame shrink/grow pseudo-instructions to move the frame, and then
reuse the existing epilogue generation functions to setup the tail call.

Co-authored-by: Jamey Sharp <jsharp@fastly.com>

Author: elliottt

cranelift/codegen/src/isa/aarch64/lower/isle.rs

@@ -100,6 +100,9 @@ impl Context for IsleContext<'_, '_, MInst, AArch64Backend> {
             self.lower_ctx.sigs(),
             callee_sig,
             &callee,
+            // TODO: this should be Opcode::ReturnCall, once aarch64 has been ported to the new
+            // tail call strategy.
+            Opcode::Call,
             distance,
             caller_conv,
             self.backend.flags().clone(),

cranelift/codegen/src/isa/riscv64/lower/isle.rs

@@ -18,7 +18,7 @@ use crate::machinst::{VCodeConstant, VCodeConstantData};
 use crate::{
     ir::{
         immediates::*, types::*, AtomicRmwOp, BlockCall, ExternalName, Inst, InstructionData,
-        MemFlags, StackSlot, TrapCode, Value, ValueList,
+        MemFlags, Opcode, StackSlot, TrapCode, Value, ValueList,
     },
     isa::riscv64::inst::*,
     machinst::{ArgPair, InstOutput},
@@ -82,6 +82,9 @@ impl generated_code::Context for RV64IsleContext<'_, '_, MInst, Riscv64Backend>
             self.lower_ctx.sigs(),
             callee_sig,
             &callee,
+            // TODO: this should be Opcode::ReturnCall, once riscv64 has been ported to the new
+            // tail call strategy.
+            Opcode::Call,
             distance,
             caller_conv,
             self.backend.flags().clone(),

cranelift/codegen/src/isa/x64/abi.rs

@@ -926,8 +926,9 @@ impl ABIMachineSpec for X64ABIMachineSpec {
 
 impl X64CallSite {
     pub fn emit_return_call(mut self, ctx: &mut Lower<Inst>, args: isle::ValueSlice) {
-        let (new_stack_arg_size, old_stack_arg_size) =
-            self.emit_temporary_tail_call_frame(ctx, args);
+        let new_stack_arg_size =
+            u32::try_from(self.sig(ctx.sigs()).sized_stack_arg_space()).unwrap();
+        let old_stack_arg_size = ctx.abi().stack_args_size(ctx.sigs());
 
         match new_stack_arg_size.cmp(&old_stack_arg_size) {
             core::cmp::Ordering::Equal => {}
@@ -947,43 +948,15 @@ impl X64CallSite {
             }
         }
 
-        // Make a copy of the frame pointer, since we use it when copying down
-        // the new stack frame.
-        let fp = ctx.temp_writable_gpr();
-        let rbp = PReg::from(regs::rbp().to_real_reg().unwrap());
-        ctx.emit(Inst::MovFromPReg { src: rbp, dst: fp });
-
-        // Load the return address, because copying our new stack frame
-        // over our current stack frame might overwrite it, and we'll need to
-        // place it in the correct location after we do that copy.
-        //
-        // But we only need to actually move the return address if the size of
-        // stack arguments changes.
-        let ret_addr = if new_stack_arg_size != old_stack_arg_size {
-            let ret_addr = ctx.temp_writable_gpr();
-            ctx.emit(Inst::Mov64MR {
-                src: SyntheticAmode::Real(Amode::ImmReg {
-                    simm32: 8,
-                    base: *fp.to_reg(),
-                    flags: MemFlags::trusted(),
-                }),
-                dst: ret_addr,
-            });
-            Some(ret_addr.to_reg())
-        } else {
-            None
-        };
+        // Put all arguments in registers and stack slots (within that newly
+        // allocated stack space).
+        self.emit_args(ctx, args);
+        self.emit_stack_ret_arg_for_tail_call(ctx);
 
         // Finally, emit the macro instruction to copy the new stack frame over
         // our current one and do the actual tail call!
-
         let dest = self.dest().clone();
         let info = Box::new(ReturnCallInfo {
-            new_stack_arg_size,
-            old_stack_arg_size,
-            ret_addr,
-            fp: fp.to_reg(),
-            tmp: ctx.temp_writable_gpr(),
             uses: self.take_uses(),
         });
         match dest {

cranelift/codegen/src/isa/x64/inst/emit.rs

@@ -1,6 +1,6 @@
 use crate::ir;
 use crate::ir::immediates::{Ieee32, Ieee64};
-use crate::ir::{KnownSymbol, MemFlags};
+use crate::ir::KnownSymbol;
 use crate::isa::x64::encoding::evex::{EvexInstruction, EvexVectorLength, RegisterOrAmode};
 use crate::isa::x64::encoding::rex::{
     emit_simm, emit_std_enc_enc, emit_std_enc_mem, emit_std_reg_mem, emit_std_reg_reg, int_reg_enc,
@@ -1628,18 +1628,7 @@ pub(crate) fn emit(
             callee,
             info: call_info,
         } => {
-            emit_return_call_common_sequence(
-                allocs,
-                sink,
-                info,
-                state,
-                call_info.new_stack_arg_size,
-                call_info.old_stack_arg_size,
-                call_info.ret_addr,
-                call_info.fp,
-                call_info.tmp,
-                &call_info.uses,
-            );
+            emit_return_call_common_sequence(allocs, sink, info, state, &call_info.uses);
 
             // Finally, jump to the callee!
             //
@@ -1660,18 +1649,7 @@ pub(crate) fn emit(
         } => {
             let callee = callee.with_allocs(allocs);
 
-            emit_return_call_common_sequence(
-                allocs,
-                sink,
-                info,
-                state,
-                call_info.new_stack_arg_size,
-                call_info.old_stack_arg_size,
-                call_info.ret_addr,
-                call_info.fp,
-                call_info.tmp,
-                &call_info.uses,
-            );
+            emit_return_call_common_sequence(allocs, sink, info, state, &call_info.uses);
 
             Inst::JmpUnknown { target: callee }.emit(&[], sink, info, state);
             sink.add_call_site(ir::Opcode::ReturnCallIndirect);
@@ -1737,14 +1715,21 @@ pub(crate) fn emit(
             // TODO: this needs to be accounted for by the stack check, before `GrowFrame` is
             // emitted, so that the increment here is expected.
             //
-            // Decrement SP by `amount`
+            // Decrement SP and FP by `amount`
             Inst::alu_rmi_r(
                 OperandSize::Size64,
                 AluRmiROpcode::Sub,
                 RegMemImm::imm(*amount),
                 Writable::from_reg(regs::rsp()),
             )
             .emit(&[], sink, info, state);
+            Inst::alu_rmi_r(
+                OperandSize::Size64,
+                AluRmiROpcode::Sub,
+                RegMemImm::imm(*amount),
+                Writable::from_reg(regs::rbp()),
+            )
+            .emit(&[], sink, info, state);
 
             // The total size that we're going to copy, including the return address and frame
             // pointer that are pushed on the stack alreadcy.
@@ -1817,6 +1802,15 @@ pub(crate) fn emit(
                 Writable::from_reg(regs::rsp()),
             )
             .emit(&[], sink, info, state);
+
+            // Increment FP by `amount`
+            Inst::alu_rmi_r(
+                OperandSize::Size64,
+                AluRmiROpcode::Add,
+                RegMemImm::imm(*amount),
+                Writable::from_reg(regs::rbp()),
+            )
+            .emit(&[], sink, info, state);
         }
 
         Inst::Args { .. } => {}
@@ -4349,11 +4343,6 @@ fn emit_return_call_common_sequence(
     sink: &mut MachBuffer<Inst>,
     info: &EmitInfo,
     state: &mut EmitState,
-    new_stack_arg_size: u32,
-    old_stack_arg_size: u32,
-    ret_addr: Option<Gpr>,
-    fp: Gpr,
-    tmp: WritableGpr,
     uses: &CallArgList,
 ) {
     assert!(
@@ -4366,125 +4355,18 @@ fn emit_return_call_common_sequence(
         let _ = allocs.next(u.vreg);
     }
 
-    let ret_addr = ret_addr.map(|r| Gpr::new(allocs.next(*r)).unwrap());
-
-    let fp = allocs.next(*fp);
-
-    let tmp = allocs.next(tmp.to_reg().to_reg());
-    let tmp = Gpr::new(tmp).unwrap();
-    let tmp_w = WritableGpr::from_reg(tmp);
-
-    // Copy the new frame (which is `frame_size` bytes above the SP)
-    // onto our current frame, using only volatile, non-argument
-    // registers.
-    //
-    //
-    // The current stack layout is the following:
-    //
-    //            | ...                 |
-    //            +---------------------+
-    //            | ...                 |
-    //            | stack arguments     |
-    //            | ...                 |
-    //    current | return address      |
-    //    frame   | old FP              | <-- FP
-    //            | callee saves        |
-    //            | ...                 |
-    //            | old stack slots     |
-    //            | ...                 |
-    //            +---------------------+
-    //            | ...                 |
-    //    new     | new stack arguments |
-    //    frame   | ...                 | <-- SP
-    //            +---------------------+
-    //
-    // We need to restore the old FP, copy the new stack arguments over the old
-    // stack arguments, write the return address into the correct slot just
-    // after the new stack arguments, adjust SP to point to the new return
-    // address, and then jump to the callee (which will push the old FP again).
-
-    // Restore the old FP into `rbp`.
-    Inst::Mov64MR {
-        src: SyntheticAmode::Real(Amode::ImmReg {
-            simm32: 0,
-            base: fp,
-            flags: MemFlags::trusted(),
-        }),
-        dst: Writable::from_reg(Gpr::new(regs::rbp()).unwrap()),
+    for inst in
+        X64ABIMachineSpec::gen_clobber_restore(CallConv::Tail, &info.flags, state.frame_layout())
+    {
+        inst.emit(&[], sink, info, state);
     }
-    .emit(&[], sink, info, state);
 
-    // The new lowest address (top of stack) -- relative to FP -- for
-    // our tail callee. We compute this now so that we can move our
-    // stack arguments into place.
-    let callee_sp_relative_to_fp = i64::from(old_stack_arg_size) - i64::from(new_stack_arg_size);
-
-    // Copy over each word, using `tmp` as a temporary register.
-    //
-    // Note that we have to do this from stack slots with the highest
-    // address to lowest address because in the case of when the tail
-    // callee has more stack arguments than we do, we might otherwise
-    // overwrite some of our stack arguments before they've been copied
-    // into place.
-    assert_eq!(
-        new_stack_arg_size % 8,
-        0,
-        "stack argument space sizes should always be 8-byte aligned"
-    );
-    for i in (0..new_stack_arg_size / 8).rev() {
-        Inst::Mov64MR {
-            src: SyntheticAmode::Real(Amode::ImmReg {
-                simm32: (i * 8).try_into().unwrap(),
-                base: regs::rsp(),
-                flags: MemFlags::trusted(),
-            }),
-            dst: tmp_w,
-        }
-        .emit(&[], sink, info, state);
-        Inst::MovRM {
-            size: OperandSize::Size64,
-            src: tmp,
-            dst: SyntheticAmode::Real(Amode::ImmReg {
-                // Add 2 because we need to skip over the old FP and the
-                // return address.
-                simm32: (callee_sp_relative_to_fp + i64::from((i + 2) * 8))
-                    .try_into()
-                    .unwrap(),
-                base: fp,
-                flags: MemFlags::trusted(),
-            }),
-        }
-        .emit(&[], sink, info, state);
-    }
-
-    // Initialize SP for the tail callee, deallocating the temporary
-    // stack arguments space at the same time.
-    Inst::LoadEffectiveAddress {
-        size: OperandSize::Size64,
-        addr: SyntheticAmode::Real(Amode::ImmReg {
-            // NB: We add a word to `callee_sp_relative_to_fp` here because the
-            // callee will push FP, not us.
-            simm32: callee_sp_relative_to_fp.wrapping_add(8).try_into().unwrap(),
-            base: fp,
-            flags: MemFlags::trusted(),
-        }),
-        dst: Writable::from_reg(Gpr::new(regs::rsp()).unwrap()),
-    }
-    .emit(&[], sink, info, state);
-
-    state.adjust_virtual_sp_offset(-i64::from(new_stack_arg_size));
-
-    // Write the return address into the correct stack slot.
-    if let Some(ret_addr) = ret_addr {
-        Inst::MovRM {
-            size: OperandSize::Size64,
-            src: ret_addr,
-            dst: SyntheticAmode::Real(Amode::ImmReg {
-                simm32: 0,
-                base: regs::rsp(),
-                flags: MemFlags::trusted(),
-            }),
-        }
-        .emit(&[], sink, info, state);
+    for inst in X64ABIMachineSpec::gen_epilogue_frame_restore(
+        CallConv::Tail,
+        &info.flags,
+        &info.isa_flags,
+        state.frame_layout(),
+    ) {
+        inst.emit(&[], sink, info, state);
     }
 }

cranelift/codegen/src/isa/x64/inst/emit_state.rs

@@ -14,6 +14,10 @@ pub struct EmitState {
     /// Only used during fuzz-testing. Otherwise, it is a zero-sized struct and
     /// optimized away at compiletime. See [cranelift_control].
     ctrl_plane: ControlPlane,
+
+    /// A copy of the frame layout, used during the emission of `Inst::ReturnCallKnown` and
+    /// `Inst::ReturnCallUnknown` instructions.
+    frame_layout: FrameLayout,
 }
 
 impl MachInstEmitState<Inst> for EmitState {
@@ -23,6 +27,7 @@ impl MachInstEmitState<Inst> for EmitState {
             nominal_sp_to_fp: abi.frame_size() as i64,
             stack_map: None,
             ctrl_plane,
+            frame_layout: abi.frame_layout().clone(),
         }
     }
 
@@ -62,4 +67,8 @@ impl EmitState {
     pub(crate) fn nominal_sp_to_fp(&self) -> i64 {
         self.nominal_sp_to_fp
     }
+
+    pub(crate) fn frame_layout(&self) -> &FrameLayout {
+        &self.frame_layout
+    }
 }