[wasm] Reload shared memory size after stack guard

The stack guard may update a shared memory's size. Therefore, we need to
update the size in the instance cache in Turbofan when StackCheck is
invoked for loops.

Change-Id: I1b000adad991a6b799ad37ba36c9a33c67559d3a
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3423780Reviewed-by: Clemens Backes <clemensb@chromium.org>
Reviewed-by: Tobias Tebbi <tebbi@chromium.org>
Commit-Queue: Manos Koukoutos <manoskouk@chromium.org>
Cr-Commit-Position: refs/heads/main@{#78925}

src/compiler/loop-unrolling.cc

@@ -77,19 +77,14 @@ void UnrollLoop(Node* loop_node, ZoneUnorderedSet<Node*>* loop, uint32_t depth,
             // ( LoadFromObject )
             //   |  |
             // {stack_check}
-            // |  *  |  *
-            // |  |  |  |
-            // |  ( Call )
-            // |   |    *
-            // |   |    |
+            // |   |
+            // |   *
+            // |
+            // | * *
+            // | | |
             // {use}: EffectPhi (stack check effect that we need to replace)
             DCHECK_EQ(use->opcode(), IrOpcode::kEffectPhi);
-            DCHECK_EQ(NodeProperties::GetEffectInput(use, 1)->opcode(),
-                      IrOpcode::kCall);
             DCHECK_EQ(NodeProperties::GetEffectInput(use), stack_check);
-            DCHECK_EQ(NodeProperties::GetEffectInput(
-                          NodeProperties::GetEffectInput(use, 1)),
-                      stack_check);
             DCHECK_EQ(NodeProperties::GetEffectInput(stack_check)->opcode(),
                       IrOpcode::kLoadFromObject);
             Node* replacing_effect = NodeProperties::GetEffectInput(

src/compiler/wasm-compiler.cc

@@ -717,7 +717,9 @@ Node* WasmGraphBuilder::UndefinedValue() {
   return LOAD_ROOT(UndefinedValue, undefined_value);
 }
 
-void WasmGraphBuilder::StackCheck(wasm::WasmCodePosition position) {
+void WasmGraphBuilder::StackCheck(
+    WasmInstanceCacheNodes* shared_memory_instance_cache,
+    wasm::WasmCodePosition position) {
   DCHECK_NOT_NULL(env_);  // Wrappers don't get stack checks.
   if (!FLAG_wasm_stack_checks || !env_->runtime_exception_support) {
     return;
@@ -731,8 +733,9 @@ void WasmGraphBuilder::StackCheck(wasm::WasmCodePosition position) {
       mcgraph()->machine()->StackPointerGreaterThan(StackCheckKind::kWasm),
       limit, effect()));
 
-  Diamond stack_check(graph(), mcgraph()->common(), check, BranchHint::kTrue);
-  stack_check.Chain(control());
+  Node* if_true;
+  Node* if_false;
+  BranchExpectTrue(check, &if_true, &if_false);
 
   if (stack_check_call_operator_ == nullptr) {
     // Build and cache the stack check call operator and the constant
@@ -753,17 +756,31 @@ void WasmGraphBuilder::StackCheck(wasm::WasmCodePosition position) {
     stack_check_call_operator_ = mcgraph()->common()->Call(call_descriptor);
   }
 
-  Node* call = graph()->NewNode(stack_check_call_operator_.get(),
-                                stack_check_code_node_.get(), effect(),
-                                stack_check.if_false);
-
+  Node* call =
+      graph()->NewNode(stack_check_call_operator_.get(),
+                       stack_check_code_node_.get(), effect(), if_false);
   SetSourcePosition(call, position);
 
   DCHECK_GT(call->op()->EffectOutputCount(), 0);
   DCHECK_EQ(call->op()->ControlOutputCount(), 0);
-  Node* ephi = stack_check.EffectPhi(effect(), call);
 
-  SetEffectControl(ephi, stack_check.merge);
+  SetEffectControl(call, if_false);
+
+  Node* merge = Merge(if_true, control());
+  Node* ephi_inputs[] = {check, effect(), merge};
+  Node* ephi = EffectPhi(2, ephi_inputs);
+
+  // We only need to refresh the size of a shared memory, as its start can never
+  // change.
+  if (shared_memory_instance_cache != nullptr) {
+    Node* new_memory_size =
+        LOAD_MUTABLE_INSTANCE_FIELD(MemorySize, MachineType::UintPtr());
+    shared_memory_instance_cache->mem_size = CreateOrMergeIntoPhi(
+        MachineType::PointerRepresentation(), merge,
+        shared_memory_instance_cache->mem_size, new_memory_size);
+  }
+
+  SetEffectControl(ephi, merge);
 }
 
 void WasmGraphBuilder::PatchInStackCheckIfNeeded() {
@@ -775,9 +792,11 @@ void WasmGraphBuilder::PatchInStackCheckIfNeeded() {
   SetEffectControl(dummy);
   // The function-prologue stack check is associated with position 0, which
   // is never a position of any instruction in the function.
-  StackCheck(0);
+  // We pass the null instance cache, as we are at the beginning of the function
+  // and do not need to update it.
+  StackCheck(nullptr, 0);
 
-  // In testing, no steck checks were emitted. Nothing to rewire then.
+  // In testing, no stack checks were emitted. Nothing to rewire then.
   if (effect() == dummy) return;
 
   // Now patch all control uses of {start} to use {control} and all effect uses

src/compiler/wasm-compiler.h

@@ -299,7 +299,8 @@ class WasmGraphBuilder {
   void AppendToMerge(Node* merge, Node* from);
   void AppendToPhi(Node* phi, Node* from);
 
-  void StackCheck(wasm::WasmCodePosition);
+  void StackCheck(WasmInstanceCacheNodes* shared_memory_instance_cache,
+                  wasm::WasmCodePosition);
 
   void PatchInStackCheckIfNeeded();
 

src/wasm/graph-builder-interface.cc

@@ -235,11 +235,16 @@ class WasmGraphBuildingInterface {
     BitVector* assigned = WasmDecoder<validate>::AnalyzeLoopAssignment(
         decoder, decoder->pc(), decoder->num_locals(), decoder->zone());
     if (decoder->failed()) return;
+    int instance_cache_index = decoder->num_locals();
+    // If the module has shared memory, the stack guard might reallocate the
+    // shared memory. We have to assume the instance cache will be updated.
+    if (decoder->module_->has_shared_memory) {
+      assigned->Add(instance_cache_index);
+    }
     DCHECK_NOT_NULL(assigned);
     decoder->control_at(0)->loop_assignments = assigned;
 
     // Only introduce phis for variables assigned in this loop.
-    int instance_cache_index = decoder->num_locals();
     for (int i = decoder->num_locals() - 1; i >= 0; i--) {
       if (!assigned->Contains(i)) continue;
       TFNode* inputs[] = {ssa_env_->locals[i], control()};
@@ -253,7 +258,10 @@ class WasmGraphBuildingInterface {
 
     // Now we setup a new environment for the inside of the loop.
     SetEnv(Split(decoder->zone(), ssa_env_));
-    builder_->StackCheck(decoder->position());
+    builder_->StackCheck(decoder->module_->has_shared_memory
+                             ? &ssa_env_->instance_cache
+                             : nullptr,
+                         decoder->position());
     ssa_env_->SetNotMerged();
 
     // Wrap input merge into phis.

test/mjsunit/wasm/grow-shared-memory.js

@@ -2,18 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-// Flags: --wasm-grow-shared-memory --experimental-wasm-threads
-
 d8.file.execute("test/mjsunit/wasm/wasm-module-builder.js");
 
-(function TestGrowSharedMemoryWithoutPostMessage() {
-  print(arguments.callee.name);
-  let memory = new WebAssembly.Memory({initial: 1, maximum: 5, shared: true});
-  assertEquals(memory.buffer.byteLength, kPageSize);
-  assertEquals(1, memory.grow(1));
-  assertEquals(memory.buffer.byteLength, 2 * kPageSize);
-})();
-
 function assertIsWasmSharedMemory(memory) {
  assertTrue(memory instanceof Object,
      "Memory is not an object");
@@ -34,6 +24,14 @@ function assertTrue(value, msg) {
 
 let workerHelpers = assertTrue.toString() + assertIsWasmSharedMemory.toString();
 
+(function TestGrowSharedMemoryWithoutPostMessage() {
+  print(arguments.callee.name);
+  let memory = new WebAssembly.Memory({initial: 1, maximum: 5, shared: true});
+  assertEquals(memory.buffer.byteLength, kPageSize);
+  assertEquals(1, memory.grow(1));
+  assertEquals(memory.buffer.byteLength, 2 * kPageSize);
+})();
+
 (function TestPostMessageWithGrow() {
   print(arguments.callee.name);
   function workerCode(workerHelpers) {
@@ -381,3 +379,75 @@ let workerHelpers = assertTrue.toString() + assertIsWasmSharedMemory.toString();
     "initial": 1, "maximum": 2, "shared": true });
   assertEquals(memory.grow(0), 1);
 })();
+
+// Tests that a function receives the update of a shared memory's size if a
+// loop's stack guard gets invoked. This is not strictly required by spec, but
+// we implement it as an optimization.
+(function TestStackGuardUpdatesMemorySize() {
+  print(arguments.callee.name);
+
+  let initial_size = 1;
+  let final_size = 2;
+
+  let memory = new WebAssembly.Memory({initial: 1, maximum: 5, shared: true});
+
+  let sync_index = 64;
+  let sync_value = 42;
+
+  let builder = new WasmModuleBuilder();
+  builder.addImportedMemory("mod", "mem", 1, 5, true);
+  // int x;
+  // while (true) {
+  //   memory[sync_index] = sync_value;
+  //   x = memory_size();
+  //   if (x != 1) break;
+  // }
+  // return x;
+  builder.addFunction("main", kSig_i_v)
+    .addLocals(kWasmI32, 1)
+    .addBody([
+      kExprLoop, kWasmVoid,
+        ...wasmI32Const(sync_index),
+        ...wasmI32Const(sync_value),
+        kAtomicPrefix, kExprI32AtomicStore, 0, 0,
+        kExprMemorySize, 0, kExprLocalTee, 0,
+        kExprI32Const, initial_size,
+        kExprI32Eq,
+        kExprBrIf, 0,
+      kExprEnd,
+      kExprLocalGet, 0])
+    .exportFunc();
+
+  builder.addFunction("setter", kSig_v_ii)
+      .addBody([kExprLocalGet, 0, kExprLocalGet, 1,
+                kAtomicPrefix, kExprI32AtomicStore, 0, 0])
+      .exportFunc();
+
+  builder.addFunction("getter", kSig_i_i)
+      .addBody([kExprLocalGet, 0, kAtomicPrefix, kExprI32AtomicLoad, 0, 0])
+      .exportFunc();
+
+  let module = new WebAssembly.Module(builder.toBuffer());
+
+  function workerCode() {
+    onmessage = function(obj) {
+      let instance = new WebAssembly.Instance(
+          obj.module, {mod: {mem: obj.memory}});
+      let res = instance.exports.main();
+      postMessage(res);
+    }
+  }
+
+  let worker = new Worker(workerCode,
+                          {type: 'function', arguments: []});
+  worker.postMessage({module: module, memory: memory});
+
+  let instance = new WebAssembly.Instance(module, {mod: {mem: memory}});
+
+  // Make sure the worker thread has entered the loop.
+  while (instance.exports.getter(sync_index) != sync_value) {}
+
+  memory.grow(final_size - initial_size);
+
+  assertEquals(final_size, worker.getMessage());
+})();