[wasm][liftoff] Fix and cleanup tracing of return value

- Fix tracing of reference return values. StoreTaggedPointer should not
  use the write barrier since we are writing to the stack.
- Avoid re-allocating a slot for the return value when it is already
  spilled.

R=manoskouk@chromium.org

Change-Id: I6418c48332964a1c3d407abafaf466b0e789be69
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3882971
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Reviewed-by: Manos Koukoutos <manoskouk@chromium.org>
Cr-Commit-Position: refs/heads/main@{#83102}

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

@@ -1550,6 +1550,10 @@ void LiftoffAssembler::FillStackSlotsWithZero(int start, int size) {
   pop(r0);
 }
 
+void LiftoffAssembler::LoadSpillAddress(Register dst, int offset) {
+  sub(dst, fp, Operand(offset));
+}
+
 #define I32_BINOP(name, instruction)                             \
   void LiftoffAssembler::emit_##name(Register dst, Register lhs, \
                                      Register rhs) {             \

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

@@ -1033,6 +1033,10 @@ void LiftoffAssembler::FillStackSlotsWithZero(int start, int size) {
   }
 }
 
+void LiftoffAssembler::LoadSpillAddress(Register dst, int offset) {
+  Sub(dst, fp, offset);
+}
+
 #define I32_BINOP(name, instruction)                             \
   void LiftoffAssembler::emit_##name(Register dst, Register lhs, \
                                      Register rhs) {             \

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

@@ -1286,6 +1286,10 @@ void LiftoffAssembler::FillStackSlotsWithZero(int start, int size) {
   }
 }
 
+void LiftoffAssembler::LoadSpillAddress(Register dst, int offset) {
+  lea(dst, liftoff::GetStackSlot(offset));
+}
+
 void LiftoffAssembler::emit_i32_add(Register dst, Register lhs, Register rhs) {
   if (lhs != dst) {
     lea(dst, Operand(lhs, rhs, times_1, 0));

src/wasm/baseline/liftoff-assembler.h

@@ -658,6 +658,7 @@ class LiftoffAssembler : public TurboAssembler {
   void Spill(VarState* slot);
   void SpillLocals();
   void SpillAllRegisters();
+  inline void LoadSpillAddress(Register dst, int offset);
 
   // Clear any uses of {reg} in both the cache and in {possible_uses}.
   // Any use in the stack is spilled. If any register in {possible_uses} matches

src/wasm/baseline/liftoff-compiler.cc

@@ -2345,40 +2345,28 @@ class LiftoffCompiler {
     CODE_COMMENT("trace function exit");
     // Before making the runtime call, spill all cache registers.
     __ SpillAllRegisters();
-    LiftoffRegList pinned;
-    // Get a register to hold the stack slot for the return value.
-    LiftoffRegister info = pinned.set(__ GetUnusedRegister(kGpReg, pinned));
-    __ AllocateStackSlot(info.gp(), sizeof(int64_t));
 
     // Store the return value if there is exactly one. Multiple return values
     // are not handled yet.
     size_t num_returns = decoder->sig_->return_count();
-    if (num_returns == 1) {
-      ValueKind return_kind = decoder->sig_->GetReturn(0).kind();
-      LiftoffRegister return_reg =
-          __ LoadToRegister(__ cache_state()->stack_state.back(), pinned);
-      if (is_reference(return_kind)) {
-        __ StoreTaggedPointer(info.gp(), no_reg, 0, return_reg, pinned);
-      } else {
-        __ Store(info.gp(), no_reg, 0, return_reg,
-                 StoreType::ForValueKind(return_kind), pinned);
-      }
-    }
     // Put the parameter in its place.
     WasmTraceExitDescriptor descriptor;
     DCHECK_EQ(0, descriptor.GetStackParameterCount());
     DCHECK_EQ(1, descriptor.GetRegisterParameterCount());
     Register param_reg = descriptor.GetRegisterParameter(0);
-    if (info.gp() != param_reg) {
-      __ Move(param_reg, info.gp(), kPointerKind);
+    if (num_returns == 1) {
+      auto& return_slot = __ cache_state()->stack_state.back();
+      if (return_slot.is_const()) {
+        __ Spill(&return_slot);
+      }
+      DCHECK(return_slot.is_stack());
+      __ LoadSpillAddress(param_reg, return_slot.offset());
     }
 
     source_position_table_builder_.AddPosition(
         __ pc_offset(), SourcePosition(decoder->position()), false);
     __ CallRuntimeStub(WasmCode::kWasmTraceExit);
     DefineSafepoint();
-
-    __ DeallocateStackSlot(sizeof(int64_t));
   }
 
   void TierupCheckOnTailCall(FullDecoder* decoder) {

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

@@ -1018,6 +1018,10 @@ void LiftoffAssembler::FillStackSlotsWithZero(int start, int size) {
   }
 }
 
+void LiftoffAssembler::LoadSpillAddress(Register dst, int offset) {
+  leaq(dst, liftoff::GetStackSlot(offset));
+}
+
 void LiftoffAssembler::emit_trace_instruction(uint32_t markid) {
   Assembler::emit_trace_instruction(Immediate(markid));
 }