[wasm][liftoff] Add direct tail-calls

R=clemensb@chromium.org

Bug: v8:10693
Change-Id: I2ffc99bfb9e96afd740fc1a095ccca61b7c5ce19
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2289970
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#68990}

src/wasm/baseline/arm/liftoff-assembler-arm.h

@@ -414,6 +414,30 @@ int LiftoffAssembler::PrepareStackFrame() {
   return offset;
 }
 
+void LiftoffAssembler::PrepareTailCall(int num_callee_stack_params,
+                                       int stack_param_delta) {
+  UseScratchRegisterScope temps(this);
+  Register scratch = temps.Acquire();
+
+  // Push the return address and frame pointer to complete the stack frame.
+  sub(sp, sp, Operand(8));
+  ldr(scratch, MemOperand(fp, 4));
+  str(scratch, MemOperand(sp, 4));
+  ldr(scratch, MemOperand(fp, 0));
+  str(scratch, MemOperand(sp, 0));
+
+  // Shift the whole frame upwards.
+  int slot_count = num_callee_stack_params + 2;
+  for (int i = slot_count - 1; i >= 0; --i) {
+    ldr(scratch, MemOperand(sp, i * 4));
+    str(scratch, MemOperand(fp, (i - stack_param_delta) * 4));
+  }
+
+  // Set the new stack and frame pointer.
+  sub(sp, fp, Operand(stack_param_delta * 4));
+  Pop(lr, fp);
+}
+
 void LiftoffAssembler::PatchPrepareStackFrame(int offset, int frame_size) {
 #ifdef USE_SIMULATOR
   // When using the simulator, deal with Liftoff which allocates the stack
@@ -3581,6 +3605,10 @@ void LiftoffAssembler::CallNativeWasmCode(Address addr) {
   Call(addr, RelocInfo::WASM_CALL);
 }
 
+void LiftoffAssembler::TailCallNativeWasmCode(Address addr) {
+  Jump(addr, RelocInfo::WASM_CALL);
+}
+
 void LiftoffAssembler::CallIndirect(const wasm::FunctionSig* sig,
                                     compiler::CallDescriptor* call_descriptor,
                                     Register target) {

src/wasm/baseline/arm64/liftoff-assembler-arm64.h

@@ -181,6 +181,30 @@ int LiftoffAssembler::PrepareStackFrame() {
   return offset;
 }
 
+void LiftoffAssembler::PrepareTailCall(int num_callee_stack_params,
+                                       int stack_param_delta) {
+  UseScratchRegisterScope temps(this);
+  Register scratch = temps.AcquireX();
+
+  // Push the return address and frame pointer to complete the stack frame.
+  sub(sp, sp, 16);
+  ldr(scratch, MemOperand(fp, 8));
+  Poke(scratch, 8);
+  ldr(scratch, MemOperand(fp, 0));
+  Poke(scratch, 0);
+
+  // Shift the whole frame upwards.
+  int slot_count = num_callee_stack_params + 2;
+  for (int i = slot_count - 1; i >= 0; --i) {
+    ldr(scratch, MemOperand(sp, i * 8));
+    str(scratch, MemOperand(fp, (i - stack_param_delta) * 8));
+  }
+
+  // Set the new stack and frame pointer.
+  Sub(sp, fp, stack_param_delta * 8);
+  Pop<kAuthLR>(fp, lr);
+}
+
 void LiftoffAssembler::PatchPrepareStackFrame(int offset, int frame_size) {
   static_assert(kStackSlotSize == kXRegSize,
                 "kStackSlotSize must equal kXRegSize");
@@ -2592,6 +2616,10 @@ void LiftoffAssembler::CallNativeWasmCode(Address addr) {
   Call(addr, RelocInfo::WASM_CALL);
 }
 
+void LiftoffAssembler::TailCallNativeWasmCode(Address addr) {
+  Jump(addr, RelocInfo::WASM_CALL);
+}
+
 void LiftoffAssembler::CallIndirect(const wasm::FunctionSig* sig,
                                     compiler::CallDescriptor* call_descriptor,
                                     Register target) {

src/wasm/baseline/ia32/liftoff-assembler-ia32.h

@@ -157,6 +157,27 @@ int LiftoffAssembler::PrepareStackFrame() {
   return offset;
 }
 
+void LiftoffAssembler::PrepareTailCall(int num_callee_stack_params,
+                                       int stack_param_delta) {
+  // Push the return address and frame pointer to complete the stack frame.
+  push(Operand(ebp, 4));
+  push(Operand(ebp, 0));
+
+  // Shift the whole frame upwards.
+  Register scratch = eax;
+  push(scratch);
+  const int slot_count = num_callee_stack_params + 2;
+  for (int i = slot_count; i > 0; --i) {
+    mov(scratch, Operand(esp, i * 4));
+    mov(Operand(ebp, (i - stack_param_delta - 1) * 4), scratch);
+  }
+  pop(scratch);
+
+  // Set the new stack and frame pointers.
+  lea(esp, Operand(ebp, -stack_param_delta * 4));
+  pop(ebp);
+}
+
 void LiftoffAssembler::PatchPrepareStackFrame(int offset, int frame_size) {
   DCHECK_EQ(frame_size % kSystemPointerSize, 0);
   // We can't run out of space, just pass anything big enough to not cause the
@@ -4199,6 +4220,10 @@ void LiftoffAssembler::CallNativeWasmCode(Address addr) {
   wasm_call(addr, RelocInfo::WASM_CALL);
 }
 
+void LiftoffAssembler::TailCallNativeWasmCode(Address addr) {
+  jmp(addr, RelocInfo::WASM_CALL);
+}
+
 void LiftoffAssembler::CallIndirect(const wasm::FunctionSig* sig,
                                     compiler::CallDescriptor* call_descriptor,
                                     Register target) {

src/wasm/baseline/liftoff-assembler.h

@@ -463,6 +463,8 @@ class LiftoffAssembler : public TurboAssembler {
   // which can later be patched (via {PatchPrepareStackFrame)} when the size of
   // the frame is known.
   inline int PrepareStackFrame();
+  inline void PrepareTailCall(int num_callee_stack_params,
+                              int stack_param_delta);
   inline void PatchPrepareStackFrame(int offset, int frame_size);
   inline void FinishCode();
   inline void AbortCompilation();
@@ -1090,6 +1092,7 @@ class LiftoffAssembler : public TurboAssembler {
                     int stack_bytes, ExternalReference ext_ref);
 
   inline void CallNativeWasmCode(Address addr);
+  inline void TailCallNativeWasmCode(Address addr);
   // Indirect call: If {target == no_reg}, then pop the target from the stack.
   inline void CallIndirect(const FunctionSig* sig,
                            compiler::CallDescriptor* call_descriptor,

src/wasm/baseline/liftoff-compiler.cc

@@ -919,8 +919,6 @@ class LiftoffCompiler {
 
   void EndControl(FullDecoder* decoder, Control* c) {}
 
-  enum CCallReturn : bool { kHasReturn = true, kNoReturn = false };
-
   void GenerateCCall(const LiftoffRegister* result_regs, const FunctionSig* sig,
                      ValueType out_argument_type,
                      const LiftoffRegister* arg_regs,
@@ -2294,63 +2292,12 @@ class LiftoffCompiler {
         __ cache_state()->stack_state.begin(), assume_spilling);
   }
 
+  enum CallKind : bool { kReturnCall = true, kNoReturnCall = false };
+
   void CallDirect(FullDecoder* decoder,
                   const CallFunctionImmediate<validate>& imm,
-                  const Value args[], Value returns[]) {
-    for (ValueType ret : imm.sig->returns()) {
-      if (!CheckSupportedType(decoder, kSupportedTypes, ret, "return")) {
-        return;
-      }
-    }
-
-    auto call_descriptor =
-        compiler::GetWasmCallDescriptor(compilation_zone_, imm.sig);
-    call_descriptor =
-        GetLoweredCallDescriptor(compilation_zone_, call_descriptor);
-
-    // Place the source position before any stack manipulation, since this will
-    // be used for OSR in debugging.
-    source_position_table_builder_.AddPosition(
-        __ pc_offset(), SourcePosition(decoder->position()), true);
-
-    if (imm.index < env_->module->num_imported_functions) {
-      // A direct call to an imported function.
-      LiftoffRegList pinned;
-      Register tmp = pinned.set(__ GetUnusedRegister(kGpReg, pinned)).gp();
-      Register target = pinned.set(__ GetUnusedRegister(kGpReg, pinned)).gp();
-
-      Register imported_targets = tmp;
-      LOAD_INSTANCE_FIELD(imported_targets, ImportedFunctionTargets,
-                          kSystemPointerSize);
-      __ Load(LiftoffRegister(target), imported_targets, no_reg,
-              imm.index * sizeof(Address), kPointerLoadType, pinned);
-
-      Register imported_function_refs = tmp;
-      LOAD_TAGGED_PTR_INSTANCE_FIELD(imported_function_refs,
-                                     ImportedFunctionRefs);
-      Register imported_function_ref = tmp;
-      __ LoadTaggedPointer(
-          imported_function_ref, imported_function_refs, no_reg,
-          ObjectAccess::ElementOffsetInTaggedFixedArray(imm.index), pinned);
-
-      Register* explicit_instance = &imported_function_ref;
-      __ PrepareCall(imm.sig, call_descriptor, &target, explicit_instance);
-      __ CallIndirect(imm.sig, call_descriptor, target);
-    } else {
-      // A direct call within this module just gets the current instance.
-      __ PrepareCall(imm.sig, call_descriptor);
-
-      // Just encode the function index. This will be patched at instantiation.
-      Address addr = static_cast<Address>(imm.index);
-      __ CallNativeWasmCode(addr);
-    }
-
-    RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
-    safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
-
-    MaybeGenerateExtraSourcePos(decoder);
-
-    __ FinishCall(imm.sig, call_descriptor);
+                  const Value args[], Value[]) {
+    CallDirect(decoder, imm, args, nullptr, kNoReturnCall);
   }
 
   void CallIndirect(FullDecoder* decoder, const Value& index_val,
@@ -2490,8 +2437,9 @@ class LiftoffCompiler {
   void ReturnCall(FullDecoder* decoder,
                   const CallFunctionImmediate<validate>& imm,
                   const Value args[]) {
-    unsupported(decoder, kTailCall, "return_call");
+    CallDirect(decoder, imm, args, nullptr, kReturnCall);
   }
+
   void ReturnCallIndirect(FullDecoder* decoder, const Value& index_val,
                           const CallIndirectImmediate<validate>& imm,
                           const Value args[]) {
@@ -3761,6 +3709,78 @@ class LiftoffCompiler {
     __ nop();
   }
 
+  void CallDirect(FullDecoder* decoder,
+                  const CallFunctionImmediate<validate>& imm,
+                  const Value args[], Value returns[], CallKind call_kind) {
+    for (ValueType ret : imm.sig->returns()) {
+      if (!CheckSupportedType(decoder, kSupportedTypes, ret, "return")) {
+        return;
+      }
+    }
+
+    auto call_descriptor =
+        compiler::GetWasmCallDescriptor(compilation_zone_, imm.sig);
+    call_descriptor =
+        GetLoweredCallDescriptor(compilation_zone_, call_descriptor);
+
+    // Place the source position before any stack manipulation, since this will
+    // be used for OSR in debugging.
+    source_position_table_builder_.AddPosition(
+        __ pc_offset(), SourcePosition(decoder->position()), true);
+
+    if (imm.index < env_->module->num_imported_functions) {
+      // A direct call to an imported function.
+      LiftoffRegList pinned;
+      Register tmp = pinned.set(__ GetUnusedRegister(kGpReg, pinned)).gp();
+      Register target = pinned.set(__ GetUnusedRegister(kGpReg, pinned)).gp();
+
+      Register imported_targets = tmp;
+      LOAD_INSTANCE_FIELD(imported_targets, ImportedFunctionTargets,
+                          kSystemPointerSize);
+      __ Load(LiftoffRegister(target), imported_targets, no_reg,
+              imm.index * sizeof(Address), kPointerLoadType, pinned);
+
+      Register imported_function_refs = tmp;
+      LOAD_TAGGED_PTR_INSTANCE_FIELD(imported_function_refs,
+                                     ImportedFunctionRefs);
+      Register imported_function_ref = tmp;
+      __ LoadTaggedPointer(
+          imported_function_ref, imported_function_refs, no_reg,
+          ObjectAccess::ElementOffsetInTaggedFixedArray(imm.index), pinned);
+
+      Register* explicit_instance = &imported_function_ref;
+      // TODO(thibaudm): Implement indirect tail calls.
+      if (call_kind == kReturnCall) {
+        unsupported(decoder, kTailCall, "tail calling imported function");
+        return;
+      }
+      __ PrepareCall(imm.sig, call_descriptor, &target, explicit_instance);
+      __ CallIndirect(imm.sig, call_descriptor, target);
+    } else {
+      // A direct call within this module just gets the current instance.
+      __ PrepareCall(imm.sig, call_descriptor);
+      // Just encode the function index. This will be patched at instantiation.
+      Address addr = static_cast<Address>(imm.index);
+      if (call_kind == kReturnCall) {
+        DCHECK(descriptor_->CanTailCall(call_descriptor));
+        __ PrepareTailCall(
+            static_cast<int>(call_descriptor->StackParameterCount()),
+            static_cast<int>(
+                call_descriptor->GetStackParameterDelta(descriptor_)));
+        __ TailCallNativeWasmCode(addr);
+      } else {
+        __ CallNativeWasmCode(addr);
+      }
+    }
+
+    RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
+    safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
+
+    MaybeGenerateExtraSourcePos(decoder);
+
+    __ FinishCall(imm.sig, call_descriptor);
+  }
+
   static constexpr WasmOpcode kNoOutstandingOp = kExprUnreachable;
 
   LiftoffAssembler asm_;

src/wasm/baseline/x64/liftoff-assembler-x64.h

@@ -132,6 +132,24 @@ int LiftoffAssembler::PrepareStackFrame() {
   return offset;
 }
 
+void LiftoffAssembler::PrepareTailCall(int num_callee_stack_params,
+                                       int stack_param_delta) {
+  // Push the return address and frame pointer to complete the stack frame.
+  pushq(Operand(rbp, 8));
+  pushq(Operand(rbp, 0));
+
+  // Shift the whole frame upwards.
+  const int slot_count = num_callee_stack_params + 2;
+  for (int i = slot_count - 1; i >= 0; --i) {
+    movq(kScratchRegister, Operand(rsp, i * 8));
+    movq(Operand(rbp, (i - stack_param_delta) * 8), kScratchRegister);
+  }
+
+  // Set the new stack and frame pointer.
+  leaq(rsp, Operand(rbp, -stack_param_delta * 8));
+  popq(rbp);
+}
+
 void LiftoffAssembler::PatchPrepareStackFrame(int offset, int frame_size) {
   // Need to align sp to system pointer size.
   frame_size = RoundUp(frame_size, kSystemPointerSize);
@@ -3864,6 +3882,10 @@ void LiftoffAssembler::CallNativeWasmCode(Address addr) {
   near_call(addr, RelocInfo::WASM_CALL);
 }
 
+void LiftoffAssembler::TailCallNativeWasmCode(Address addr) {
+  near_jmp(addr, RelocInfo::WASM_CALL);
+}
+
 void LiftoffAssembler::CallIndirect(const wasm::FunctionSig* sig,
                                     compiler::CallDescriptor* call_descriptor,
                                     Register target) {