[wasm] Suspend wasm continuation

Save the PC in the jump buffer and implement the suspend builtin.

R=ahaas@chromium.org
CC=​fgm@chromium.org

Bug: v8:12191
Change-Id: I1a6d965d7864dce0a572f6c8d7102046dad190fd
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3345006Reviewed-by: Jakob Kummerow <jkummerow@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/main@{#78715}

src/builtins/x64/builtins-x64.cc

@@ -3616,10 +3616,23 @@ void Builtins::Generate_GenericJSToWasmWrapper(MacroAssembler* masm) {
 }
 
 namespace {
-// Helper function for WasmReturnPromiseOnSuspend.
-void LoadJumpBuffer(MacroAssembler* masm, Register jmpbuf) {
+// Helper function for wasm stack switching.
+void FillJumpBuffer(MacroAssembler* masm, Register jmpbuf, Label* pc) {
+  __ movq(MemOperand(jmpbuf, wasm::kJmpBufSpOffset), rsp);
+  __ movq(MemOperand(jmpbuf, wasm::kJmpBufFpOffset), rbp);
+  __ movq(kScratchRegister,
+          __ StackLimitAsOperand(StackLimitKind::kRealStackLimit));
+  __ movq(MemOperand(jmpbuf, wasm::kJmpBufStackLimitOffset), kScratchRegister);
+  __ leaq(kScratchRegister, MemOperand(pc, 0));
+  __ movq(MemOperand(jmpbuf, wasm::kJmpBufPcOffset), kScratchRegister);
+}
+
+void LoadJumpBuffer(MacroAssembler* masm, Register jmpbuf, bool load_pc) {
   __ movq(rsp, MemOperand(jmpbuf, wasm::kJmpBufSpOffset));
   __ movq(rbp, MemOperand(jmpbuf, wasm::kJmpBufFpOffset));
+  if (load_pc) {
+    __ jmp(MemOperand(jmpbuf, wasm::kJmpBufPcOffset));
+  }
   // The stack limit is set separately under the ExecutionAccess lock.
   // TODO(thibaudm): Reload live registers.
 }
@@ -3674,19 +3687,9 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
   __ LoadExternalPointerField(
       jmpbuf, FieldOperand(foreign_jmpbuf, Foreign::kForeignAddressOffset),
       kForeignForeignAddressTag, r8);
-  __ movq(MemOperand(jmpbuf, wasm::kJmpBufSpOffset), rsp);
-  __ movq(MemOperand(jmpbuf, wasm::kJmpBufFpOffset), rbp);
-  Register stack_limit_address = rcx;
-  __ movq(stack_limit_address,
-          FieldOperand(wasm_instance,
-                       WasmInstanceObject::kRealStackLimitAddressOffset));
-  Register stack_limit = rdx;
-  __ movq(stack_limit, MemOperand(stack_limit_address, 0));
-  __ movq(MemOperand(jmpbuf, wasm::kJmpBufStackLimitOffset), stack_limit);
-  // TODO(thibaudm): Save live registers.
+  Label suspend;
+  FillJumpBuffer(masm, jmpbuf, &suspend);
   foreign_jmpbuf = no_reg;
-  stack_limit = no_reg;
-  stack_limit_address = no_reg;
   // live: [rsi, rdi, rax]
 
   // -------------------------------------------
@@ -3725,7 +3728,7 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
       kForeignForeignAddressTag, r8);
   __ Move(GCScanSlotPlace, 0);
   // Switch stack!
-  LoadJumpBuffer(masm, target_jmpbuf);
+  LoadJumpBuffer(masm, target_jmpbuf, false);
   foreign_jmpbuf = no_reg;
   target_jmpbuf = no_reg;
   // live: [rsi, rdi]
@@ -3787,7 +3790,7 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
       jmpbuf, FieldOperand(foreign_jmpbuf, Foreign::kForeignAddressOffset),
       kForeignForeignAddressTag, r8);
   // Switch stack!
-  LoadJumpBuffer(masm, jmpbuf);
+  LoadJumpBuffer(masm, jmpbuf, false);
   __ Move(GCScanSlotPlace, 1);
   __ Push(wasm_instance);  // Spill.
   __ Move(kContextRegister, Smi::zero());
@@ -3812,17 +3815,22 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
 #ifdef DEBUG
   // Check that the parent suspender is inactive.
   Label parent_inactive;
-  __ cmpq(FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
-          Immediate(WasmSuspenderObject::Inactive));
+  Register state = rbx;
+  __ LoadTaggedSignedField(
+      state, FieldOperand(suspender, WasmSuspenderObject::kStateOffset));
+  __ SmiCompare(state, Smi::FromInt(WasmSuspenderObject::Inactive));
   __ j(equal, &parent_inactive, Label::kNear);
   __ Trap();
   __ bind(&parent_inactive);
 #endif
-  __ movq(FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
-          Immediate(WasmSuspenderObject::State::Active));
+  __ StoreTaggedSignedField(
+      FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
+      Smi::FromInt(WasmSuspenderObject::State::Active));
   __ bind(&undefined);
   __ movq(masm->RootAsOperand(RootIndex::kActiveSuspender), suspender);
 
+  __ bind(&suspend);
+
   // -------------------------------------------
   // Epilogue.
   // -------------------------------------------
@@ -3838,6 +3846,97 @@ void Builtins::Generate_WasmReturnPromiseOnSuspend(MacroAssembler* masm) {
 void Builtins::Generate_WasmSuspend(MacroAssembler* masm) {
   // Set up the stackframe.
   __ EnterFrame(StackFrame::STACK_SWITCH);
+
+  Register promise = rax;
+  Register suspender = rbx;
+
+  __ subq(rsp, Immediate(-(BuiltinWasmWrapperConstants::kGCScanSlotCountOffset -
+                           TypedFrameConstants::kFixedFrameSizeFromFp)));
+
+  // TODO(thibaudm): Throw if any of the following holds:
+  // - caller is null
+  // - ActiveSuspender is undefined
+  // - 'suspender' is not the active suspender
+
+  // -------------------------------------------
+  // Save current state in active jump buffer.
+  // -------------------------------------------
+  Label resume;
+  Register continuation = rcx;
+  __ LoadRoot(continuation, RootIndex::kActiveContinuation);
+  Register jmpbuf = rdx;
+  __ LoadAnyTaggedField(
+      jmpbuf,
+      FieldOperand(continuation, WasmContinuationObject::kJmpbufOffset));
+  __ LoadExternalPointerField(
+      jmpbuf, FieldOperand(jmpbuf, Foreign::kForeignAddressOffset),
+      kForeignForeignAddressTag, r8);
+  FillJumpBuffer(masm, jmpbuf, &resume);
+  __ StoreTaggedSignedField(
+      FieldOperand(suspender, WasmSuspenderObject::kStateOffset),
+      Smi::FromInt(WasmSuspenderObject::Suspended));
+  jmpbuf = no_reg;
+  // live: [rax, rbx, rcx]
+
+#ifdef DEBUG
+  // -------------------------------------------
+  // Check that the suspender's continuation is the active continuation.
+  // -------------------------------------------
+  // TODO(thibaudm): Once we add core stack-switching instructions, this check
+  // will not hold anymore: it's possible that the active continuation changed
+  // (due to an internal switch), so we have to update the suspender.
+  Register suspender_continuation = rdx;
+  __ LoadAnyTaggedField(
+      suspender_continuation,
+      FieldOperand(suspender, WasmSuspenderObject::kContinuationOffset));
+  __ cmpq(suspender_continuation, continuation);
+  Label ok;
+  __ j(equal, &ok);
+  __ Trap();
+  __ bind(&ok);
+#endif
+
+  // -------------------------------------------
+  // Update roots.
+  // -------------------------------------------
+  Register caller = rcx;
+  __ LoadAnyTaggedField(
+      caller,
+      FieldOperand(suspender, WasmSuspenderObject::kContinuationOffset));
+  __ LoadAnyTaggedField(
+      caller, FieldOperand(caller, WasmContinuationObject::kParentOffset));
+  __ movq(masm->RootAsOperand(RootIndex::kActiveContinuation), caller);
+  Register parent = rdx;
+  __ LoadAnyTaggedField(
+      parent, FieldOperand(suspender, WasmSuspenderObject::kParentOffset));
+  __ movq(masm->RootAsOperand(RootIndex::kActiveSuspender), parent);
+  parent = no_reg;
+  // live: [rax, rcx]
+
+  // -------------------------------------------
+  // Load jump buffer.
+  // -------------------------------------------
+  MemOperand GCScanSlotPlace =
+      MemOperand(rbp, BuiltinWasmWrapperConstants::kGCScanSlotCountOffset);
+  __ Move(GCScanSlotPlace, 2);
+  __ Push(promise);
+  __ Push(caller);
+  __ Move(kContextRegister, Smi::zero());
+  __ CallRuntime(Runtime::kWasmSyncStackLimit);
+  __ Pop(caller);
+  __ Pop(promise);
+  jmpbuf = caller;
+  __ LoadAnyTaggedField(
+      jmpbuf, FieldOperand(caller, WasmContinuationObject::kJmpbufOffset));
+  caller = no_reg;
+  __ LoadExternalPointerField(
+      jmpbuf, FieldOperand(jmpbuf, Foreign::kForeignAddressOffset),
+      kForeignForeignAddressTag, r8);
+  __ movq(kReturnRegister0, promise);
+  __ Move(GCScanSlotPlace, 0);
+  LoadJumpBuffer(masm, jmpbuf, true);
+  __ Trap();
+  __ bind(&resume);
   __ LeaveFrame(StackFrame::STACK_SWITCH);
   __ ret(0);
 }

src/codegen/interface-descriptors.h

@@ -1791,8 +1791,9 @@ class WasmFloat64ToNumberDescriptor final
 class WasmSuspendDescriptor final
     : public StaticCallInterfaceDescriptor<WasmSuspendDescriptor> {
  public:
-  DEFINE_RESULT_AND_PARAMETERS_NO_CONTEXT(0, kArg)
-  DEFINE_RESULT_AND_PARAMETER_TYPES(MachineType::AnyTagged())  // value
+  DEFINE_RESULT_AND_PARAMETERS_NO_CONTEXT(0, kArg0, kArg1)
+  DEFINE_RESULT_AND_PARAMETER_TYPES(MachineType::AnyTagged(),  // value
+                                    MachineType::AnyTagged())  // value
   DECLARE_DESCRIPTOR(WasmSuspendDescriptor)
 };
 

src/compiler/wasm-compiler.cc

@@ -7056,8 +7056,12 @@ class WasmWrapperGraphBuilder : public WasmGraphBuilder {
               Builtin::kWasmSuspend, zone_, StubCallMode::kCallWasmRuntimeStub);
           Node* call_target = mcgraph()->RelocatableIntPtrConstant(
               wasm::WasmCode::kWasmSuspend, RelocInfo::WASM_STUB_CALL);
-
-          gasm_->Call(call_descriptor, call_target, call);
+          Node* suspender =
+              gasm_->Load(MachineType::TaggedPointer(), Param(0),
+                          wasm::ObjectAccess::ToTagged(
+                              WasmApiFunctionRef::kSuspenderOffset));
+          // TODO(thibaudm): Create and pass the chained promise.
+          gasm_->Call(call_descriptor, call_target, call, suspender);
         }
         break;
       }

src/diagnostics/objects-printer.cc

@@ -1884,6 +1884,8 @@ void WasmContinuationObject::WasmContinuationObjectPrint(std::ostream& os) {
 void WasmSuspenderObject::WasmSuspenderObjectPrint(std::ostream& os) {
   PrintHeader(os, "WasmSuspenderObject");
   os << "\n - continuation: " << continuation();
+  os << "\n - parent: " << parent();
+  os << "\n - state: " << state();
   os << "\n";
 }
 

src/runtime/runtime-wasm.cc

@@ -741,6 +741,7 @@ RUNTIME_FUNCTION(Runtime_WasmAllocateContinuation) {
         .set_state(WasmSuspenderObject::Inactive);
   }
   suspender->set_state(WasmSuspenderObject::State::Active);
+  suspender->set_continuation(*target);
   active_suspender_slot.store(*suspender);
 
   SyncStackLimit(isolate);

src/wasm/stacks.h

@@ -21,12 +21,14 @@ namespace wasm {
 struct JumpBuffer {
   Address sp;
   Address fp;
+  Address pc;
   void* stack_limit;
   // TODO(thibaudm/fgm): Add general-purpose registers.
 };
 
 constexpr int kJmpBufSpOffset = offsetof(JumpBuffer, sp);
 constexpr int kJmpBufFpOffset = offsetof(JumpBuffer, fp);
+constexpr int kJmpBufPcOffset = offsetof(JumpBuffer, pc);
 constexpr int kJmpBufStackLimitOffset = offsetof(JumpBuffer, stack_limit);
 
 class StackMemory {

test/mjsunit/wasm/stack-switching.js

@@ -41,16 +41,14 @@ load("test/mjsunit/wasm/wasm-module-builder.js");
   let suspender = new WebAssembly.Suspender();
   function js_import() {
     // TODO(thibaudm): Return the value as a promise.
-    return 42;
+    return new Promise((resolve) => { resolve(42); });
   };
   let wasm_js_import = new WebAssembly.Function({parameters: [], results: ['i32']}, js_import);
   let suspending_wasm_js_import = suspender.suspendOnReturnedPromise(wasm_js_import);
   let instance = builder.instantiate({m: {import: suspending_wasm_js_import}});
   let wrapped_export = suspender.returnPromiseOnSuspend(instance.exports.test);
   let combined_promise = wrapped_export();
-  // TODO(thibaudm): Once we generate the actual wrapper, we expect 0 here
-  // The global will only be set after the promise resolves.
-  assertEquals(42, instance.exports.g.value);
+  assertEquals(0, instance.exports.g.value);
 })();
 
 (function TestStackSwitchGC() {