[wasm][liftoff] Introduce PrepareBuiltinCall

This CL provides a generic way to prepare a builtin call: The
{PrepareBuiltinCall} takes the builtin signature for 64-bit systems,
the CallDescriptor, and a Vector of VarStates for the parameters, and
moves all parameters to their correct place, which is either in a
register or on the stack.

To test the new code this CL adjusts the implementation of AtomicWait
to use PrepareBuiltinCall. Thereby AtomicWait is now also supported
on 32-bit platforms, including ia32.

R=clemensb@chromium.org

Bug: v8:10108, v8:10281
Change-Id: Ia8589166310ea2e8442531b4ed20db62d7b4aff0
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2108554
Commit-Queue: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#66810}

src/wasm/baseline/liftoff-assembler.cc

@@ -507,10 +507,8 @@ LiftoffAssembler::~LiftoffAssembler() {
   }
 }
 
-LiftoffRegister LiftoffAssembler::PopToRegister(LiftoffRegList pinned) {
-  DCHECK(!cache_state_.stack_state.empty());
-  VarState slot = cache_state_.stack_state.back();
-  cache_state_.stack_state.pop_back();
+LiftoffRegister LiftoffAssembler::LoadToRegister(VarState slot,
+                                                 LiftoffRegList pinned) {
   switch (slot.loc()) {
     case VarState::kStack: {
       LiftoffRegister reg =
@@ -532,6 +530,24 @@ LiftoffRegister LiftoffAssembler::PopToRegister(LiftoffRegList pinned) {
   UNREACHABLE();
 }
 
+LiftoffRegister LiftoffAssembler::PopToRegister(LiftoffRegList pinned) {
+  DCHECK(!cache_state_.stack_state.empty());
+  VarState slot = cache_state_.stack_state.back();
+  cache_state_.stack_state.pop_back();
+  return LoadToRegister(slot, pinned);
+}
+
+LiftoffRegister LiftoffAssembler::PeekToRegister(int index,
+                                                 LiftoffRegList pinned) {
+  DCHECK_LT(index, cache_state_.stack_state.size());
+  VarState& slot = cache_state_.stack_state.end()[-1 - index];
+  LiftoffRegister reg = LoadToRegister(slot, pinned);
+  if (!slot.is_reg()) {
+    slot.MakeRegister(reg);
+  }
+  return reg;
+}
+
 void LiftoffAssembler::MergeFullStackWith(const CacheState& target,
                                           const CacheState& source) {
   DCHECK_EQ(source.stack_height(), target.stack_height());
@@ -598,55 +614,24 @@ void LiftoffAssembler::SpillAllRegisters() {
   cache_state_.reset_used_registers();
 }
 
-void LiftoffAssembler::PrepareCall(const FunctionSig* sig,
-                                   compiler::CallDescriptor* call_descriptor,
-                                   Register* target,
-                                   Register* target_instance) {
-  uint32_t num_params = static_cast<uint32_t>(sig->parameter_count());
-  // Input 0 is the call target.
-  constexpr size_t kInputShift = 1;
-
-  // Spill all cache slots which are not being used as parameters.
-  // Don't update any register use counters, they will be reset later anyway.
-  for (uint32_t idx = 0, end = cache_state_.stack_height() - num_params;
-       idx < end; ++idx) {
-    VarState& slot = cache_state_.stack_state[idx];
-    if (!slot.is_reg()) continue;
-    Spill(slot.offset(), slot.reg(), slot.type());
-    slot.MakeStack();
-  }
-
-  LiftoffStackSlots stack_slots(this);
-  StackTransferRecipe stack_transfers(this);
-  LiftoffRegList param_regs;
-
-  // Move the target instance (if supplied) into the correct instance register.
-  compiler::LinkageLocation instance_loc =
-      call_descriptor->GetInputLocation(kInputShift);
-  DCHECK(instance_loc.IsRegister() && !instance_loc.IsAnyRegister());
-  Register instance_reg = Register::from_code(instance_loc.AsRegister());
-  param_regs.set(instance_reg);
-  if (target_instance && *target_instance != instance_reg) {
-    stack_transfers.MoveRegister(LiftoffRegister(instance_reg),
-                                 LiftoffRegister(*target_instance),
-                                 kWasmIntPtr);
-  }
-
-  // Now move all parameter values into the right slot for the call.
-  // Don't pop values yet, such that the stack height is still correct when
-  // executing the {stack_transfers}.
+namespace {
+void PrepareStackTransfers(const FunctionSig* sig,
+                           compiler::CallDescriptor* call_descriptor,
+                           const VarState* slots,
+                           LiftoffStackSlots* stack_slots,
+                           StackTransferRecipe* stack_transfers,
+                           LiftoffRegList* param_regs) {
   // Process parameters backwards, such that pushes of caller frame slots are
   // in the correct order.
-  uint32_t param_base = cache_state_.stack_height() - num_params;
   uint32_t call_desc_input_idx =
       static_cast<uint32_t>(call_descriptor->InputCount());
+  uint32_t num_params = static_cast<uint32_t>(sig->parameter_count());
   for (uint32_t i = num_params; i > 0; --i) {
     const uint32_t param = i - 1;
     ValueType type = sig->GetParam(param);
     const bool is_gp_pair = kNeedI64RegPair && type == kWasmI64;
     const int num_lowered_params = is_gp_pair ? 2 : 1;
-    const uint32_t stack_idx = param_base + param;
-    const VarState& slot = cache_state_.stack_state[stack_idx];
+    const VarState& slot = slots[param];
     const uint32_t stack_offset = slot.offset();
     // Process both halfs of a register pair separately, because they are passed
     // as separate parameters. One or both of them could end up on the stack.
@@ -679,21 +664,81 @@ void LiftoffAssembler::PrepareCall(const FunctionSig* sig,
           reg = LiftoffRegister::from_code(rc, reg_code);
         }
 
-        param_regs.set(reg);
+        param_regs->set(reg);
         if (is_gp_pair) {
-          stack_transfers.LoadI64HalfIntoRegister(reg, slot, stack_offset,
-                                                  half);
+          stack_transfers->LoadI64HalfIntoRegister(reg, slot, stack_offset,
+                                                   half);
         } else {
-          stack_transfers.LoadIntoRegister(reg, slot, stack_offset);
+          stack_transfers->LoadIntoRegister(reg, slot, stack_offset);
         }
       } else {
         DCHECK(loc.IsCallerFrameSlot());
-        stack_slots.Add(slot, stack_offset, half);
+        stack_slots->Add(slot, stack_offset, half);
       }
     }
   }
-  // {call_desc_input_idx} should point after the instance parameter now.
-  DCHECK_EQ(call_desc_input_idx, kInputShift + 1);
+}
+
+}  // namespace
+
+void LiftoffAssembler::PrepareBuiltinCall(
+    const FunctionSig* sig, compiler::CallDescriptor* call_descriptor,
+    std::initializer_list<VarState> params) {
+  LiftoffStackSlots stack_slots(this);
+  StackTransferRecipe stack_transfers(this);
+  LiftoffRegList param_regs;
+  PrepareStackTransfers(sig, call_descriptor, params.begin(), &stack_slots,
+                        &stack_transfers, &param_regs);
+  // Create all the slots.
+  stack_slots.Construct();
+  // Execute the stack transfers before filling the instance register.
+  stack_transfers.Execute();
+
+  // Reset register use counters.
+  cache_state_.reset_used_registers();
+  SpillAllRegisters();
+}
+
+void LiftoffAssembler::PrepareCall(const FunctionSig* sig,
+                                   compiler::CallDescriptor* call_descriptor,
+                                   Register* target,
+                                   Register* target_instance) {
+  uint32_t num_params = static_cast<uint32_t>(sig->parameter_count());
+  // Input 0 is the call target.
+  constexpr size_t kInputShift = 1;
+
+  // Spill all cache slots which are not being used as parameters.
+  // Don't update any register use counters, they will be reset later anyway.
+  for (uint32_t idx = 0, end = cache_state_.stack_height() - num_params;
+       idx < end; ++idx) {
+    VarState& slot = cache_state_.stack_state[idx];
+    if (!slot.is_reg()) continue;
+    Spill(slot.offset(), slot.reg(), slot.type());
+    slot.MakeStack();
+  }
+
+  LiftoffStackSlots stack_slots(this);
+  StackTransferRecipe stack_transfers(this);
+  LiftoffRegList param_regs;
+
+  // Move the target instance (if supplied) into the correct instance register.
+  compiler::LinkageLocation instance_loc =
+      call_descriptor->GetInputLocation(kInputShift);
+  DCHECK(instance_loc.IsRegister() && !instance_loc.IsAnyRegister());
+  Register instance_reg = Register::from_code(instance_loc.AsRegister());
+  param_regs.set(instance_reg);
+  if (target_instance && *target_instance != instance_reg) {
+    stack_transfers.MoveRegister(LiftoffRegister(instance_reg),
+                                 LiftoffRegister(*target_instance),
+                                 kWasmIntPtr);
+  }
+
+  if (num_params) {
+    uint32_t param_base = cache_state_.stack_height() - num_params;
+    PrepareStackTransfers(sig, call_descriptor,
+                          &cache_state_.stack_state[param_base], &stack_slots,
+                          &stack_transfers, &param_regs);
+  }
 
   // If the target register overlaps with a parameter register, then move the
   // target to another free register, or spill to the stack.

src/wasm/baseline/liftoff-assembler.h

@@ -103,6 +103,11 @@ class LiftoffAssembler : public TurboAssembler {
 
     void MakeStack() { loc_ = kStack; }
 
+    void MakeRegister(LiftoffRegister r) {
+      reg_ = r;
+      loc_ = kRegister;
+    }
+
     // Copy src to this, except for offset, since src and this could have been
     // from different stack states.
     void Copy(VarState src) {
@@ -282,6 +287,16 @@ class LiftoffAssembler : public TurboAssembler {
 
   LiftoffRegister PopToRegister(LiftoffRegList pinned = {});
 
+  // Returns the register which holds the value of stack slot {index}. If the
+  // value is not stored in a register yet, a register is allocated for it. The
+  // register is then assigned to the stack slot. The value stack height is not
+  // modified. The top of the stack is index 0, i.e. {PopToRegister()} and
+  // {PeekToRegister(0)} should result in the same register.
+  // {PeekToRegister} already decrements the used count of the register of the
+  // stack slot. Therefore the register must not be popped by {PopToRegister}
+  // but discarded with {stack_state.pop_back(count)}.
+  LiftoffRegister PeekToRegister(int index, LiftoffRegList pinned);
+
   int NextSpillOffset(ValueType type) {
     int offset = TopSpillOffset() + SlotSizeForType(type);
     if (NeedsAlignment(type)) {
@@ -373,6 +388,11 @@ class LiftoffAssembler : public TurboAssembler {
     if (offset >= max_used_spill_offset_) max_used_spill_offset_ = offset;
   }
 
+  // Load parameters into the right registers / stack slots for the call.
+  void PrepareBuiltinCall(const FunctionSig* sig,
+                          compiler::CallDescriptor* call_descriptor,
+                          std::initializer_list<VarState> params);
+
   // Load parameters into the right registers / stack slots for the call.
   // Move {*target} into another register if needed and update {*target} to that
   // register, or {no_reg} if target was spilled to the stack.
@@ -816,6 +836,8 @@ class LiftoffAssembler : public TurboAssembler {
   }
 
  private:
+  LiftoffRegister LoadToRegister(VarState slot, LiftoffRegList pinned);
+
   uint32_t num_locals_ = 0;
   static constexpr uint32_t kInlineLocalTypes = 8;
   union {

src/wasm/baseline/liftoff-compiler.cc

@@ -2592,63 +2592,85 @@ class LiftoffCompiler {
 #endif
   }
 
+  template <typename BuiltinDescriptor>
+  compiler::CallDescriptor* GetBuiltinCallDescriptor(Zone* zone) {
+    BuiltinDescriptor interface_descriptor;
+    return compiler::Linkage::GetStubCallDescriptor(
+        zone,                                           // zone
+        interface_descriptor,                           // descriptor
+        interface_descriptor.GetStackParameterCount(),  // stack parameter count
+        compiler::CallDescriptor::kNoFlags,             // flags
+        compiler::Operator::kNoProperties,              // properties
+        StubCallMode::kCallWasmRuntimeStub);            // stub call mode
+  }
+
   void AtomicWait(FullDecoder* decoder, ValueType type,
                   const MemoryAccessImmediate<validate>& imm) {
-#if V8_TARGET_ARCH_32_BIT
-    // The current implementation does not work on 32-bit platforms, as it
-    // passes i64 values over registers to builtins.
-    unsupported(decoder, kAtomics, "AtomicWait");
-    return;
-#else
     LiftoffRegList pinned;
-    LiftoffRegister timeout = pinned.set(__ PopToRegister(pinned));
-    LiftoffRegister expected_value = pinned.set(__ PopToRegister(pinned));
-    Register index = pinned.set(__ PopToRegister(pinned)).gp();
-    if (BoundsCheckMem(decoder, type.element_size_bytes(), imm.offset, index,
-                       pinned, kDoForceCheck)) {
+    Register index_reg = pinned.set(__ PeekToRegister(2, pinned)).gp();
+    if (BoundsCheckMem(decoder, type.element_size_bytes(), imm.offset,
+                       index_reg, pinned, kDoForceCheck)) {
       return;
     }
-    AlignmentCheckMem(decoder, type.element_size_bytes(), imm.offset, index,
+    AlignmentCheckMem(decoder, type.element_size_bytes(), imm.offset, index_reg,
                       pinned);
 
     uint32_t offset = imm.offset;
-    index = AddMemoryMasking(index, &offset, &pinned);
-    if (offset != 0) __ emit_i32_add(index, index, offset);
+    index_reg = AddMemoryMasking(index_reg, &offset, &pinned);
+    if (offset != 0) __ emit_i32_add(index_reg, index_reg, offset);
 
-    // TODO(ahaas): Use PrepareCall to prepare parameters.
-    __ SpillAllRegisters();
+    LiftoffAssembler::VarState timeout =
+        __ cache_state()->stack_state.end()[-1];
+    LiftoffAssembler::VarState expected_value =
+        __ cache_state()->stack_state.end()[-2];
+    LiftoffAssembler::VarState index = __ cache_state()->stack_state.end()[-3];
 
-    Register params[3]{no_reg, no_reg, no_reg};
+    // We have to set the correct register for the index. It may have changed
+    // above in {AddMemoryMasking}.
+    index.MakeRegister(LiftoffRegister(index_reg));
+
+    WasmCode::RuntimeStubId target;
+    compiler::CallDescriptor* call_descriptor;
     if (type == kWasmI32) {
-      WasmI32AtomicWait64Descriptor descriptor;
-      DCHECK_EQ(0, descriptor.GetStackParameterCount());
-      DCHECK_EQ(3, descriptor.GetRegisterParameterCount());
-      params[0] = descriptor.GetRegisterParameter(0);
-      params[1] = descriptor.GetRegisterParameter(1);
-      params[2] = descriptor.GetRegisterParameter(2);
+      if (kNeedI64RegPair) {
+        target = WasmCode::kWasmI32AtomicWait32;
+        call_descriptor =
+            GetBuiltinCallDescriptor<WasmI32AtomicWait32Descriptor>(
+                compilation_zone_);
+      } else {
+        target = WasmCode::kWasmI32AtomicWait64;
+        call_descriptor =
+            GetBuiltinCallDescriptor<WasmI32AtomicWait64Descriptor>(
+                compilation_zone_);
+      }
     } else {
-      DCHECK_EQ(kWasmI64, type);
-
-      WasmI64AtomicWait64Descriptor descriptor;
-      DCHECK_EQ(0, descriptor.GetStackParameterCount());
-      DCHECK_EQ(3, descriptor.GetRegisterParameterCount());
-      params[0] = descriptor.GetRegisterParameter(0);
-      params[1] = descriptor.GetRegisterParameter(1);
-      params[2] = descriptor.GetRegisterParameter(2);
+      if (kNeedI64RegPair) {
+        target = WasmCode::kWasmI64AtomicWait32;
+        call_descriptor =
+            GetBuiltinCallDescriptor<WasmI64AtomicWait32Descriptor>(
+                compilation_zone_);
+      } else {
+        target = WasmCode::kWasmI64AtomicWait64;
+        call_descriptor =
+            GetBuiltinCallDescriptor<WasmI64AtomicWait64Descriptor>(
+                compilation_zone_);
+      }
     }
-    LiftoffAssembler::ParallelRegisterMoveTuple reg_moves[]{
-        {LiftoffRegister(params[0]), LiftoffRegister(index), kWasmI32},
-        {LiftoffRegister(params[1]), expected_value, type},
-        {LiftoffRegister(params[2]), timeout, kWasmI64}};
-    __ ParallelRegisterMove(ArrayVector(reg_moves));
 
-    __ CallRuntimeStub(type == kWasmI32 ? WasmCode::kWasmI32AtomicWait64
-                                        : WasmCode::kWasmI64AtomicWait64);
+    ValueType sig_reps[] = {kWasmI32, type, kWasmI64};
+    FunctionSig sig(0, 3, sig_reps);
+
+    __ PrepareBuiltinCall(&sig, call_descriptor,
+                          {index, expected_value, timeout});
+    __ CallRuntimeStub(target);
+
+    // Pop parameters from the value stack.
+    __ cache_state()->stack_state.pop_back(3);
+
     RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
     safepoint_table_builder_.DefineSafepoint(&asm_, Safepoint::kNoLazyDeopt);
 
     __ PushRegister(kWasmI32, LiftoffRegister(kReturnRegister0));
-#endif
   }
 
   void AtomicNotify(FullDecoder* decoder,