[wasm][liftoff] Delay use counter decrement of PeekToRegister

This CL fixes a bug in the code generation for I32AtomicCompareExchange
in Liftoff on ia32. The problem is the inconsistency that
LiftoffAssembler::PeekToRegister(...) introduces to the cache state.
PeekToRegister loads the value from the value stack into a register, but
does not pop the value off the stack. When the value was already stored
in a register, the use counter of that register gets decreased, even
though the value is still on the stack.

The problem arises when this register later gets reused, which is
necessary unfortunately on ia32. When SpillRegister is called for this
register, all stack values that are stored in this register get written
to memory. SpillRegister uses the use counter of the register to detect
when the register was spilled to all stack slots that were cached by
this register. However, as described above, the value stack and the use
counter are inconsistent at that moment, so SpillRegister finishes
early and does not spill the register to all stack values, and this
causes the bug later.

With this CL the decrement of the use counter gets delayed until when
the value actually gets popped off the stack.

R=clemensb@chromium.org

Bug: chromium:1145135
Change-Id: I07cb256a7e5135dbce41b246c120650635ad2758
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2602464Reviewed-by: Clemens Backes <clemensb@chromium.org>
Commit-Queue: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#72018}

src/wasm/baseline/liftoff-assembler.cc

@@ -604,14 +604,25 @@ LiftoffRegister LiftoffAssembler::PeekToRegister(int index,
   DCHECK_LT(index, cache_state_.stack_state.size());
   VarState& slot = cache_state_.stack_state.end()[-1 - index];
   if (slot.is_reg()) {
-    cache_state_.dec_used(slot.reg());
     return slot.reg();
   }
   LiftoffRegister reg = LoadToRegister(slot, pinned);
+  cache_state_.inc_used(reg);
   slot.MakeRegister(reg);
   return reg;
 }
 
+void LiftoffAssembler::DropValues(int count) {
+  for (int i = 0; i < count; ++i) {
+    DCHECK(!cache_state_.stack_state.empty());
+    VarState slot = cache_state_.stack_state.back();
+    cache_state_.stack_state.pop_back();
+    if (slot.is_reg()) {
+      cache_state_.dec_used(slot.reg());
+    }
+  }
+}
+
 void LiftoffAssembler::PrepareLoopArgs(int num) {
   for (int i = 0; i < num; ++i) {
     VarState& slot = cache_state_.stack_state.end()[-1 - i];

src/wasm/baseline/liftoff-assembler.h

@@ -360,11 +360,12 @@ class LiftoffAssembler : public TurboAssembler {
   // 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)}.
+  // When the value is finally popped, the use counter of its register has to be
+  // decremented. This can be done by popping the value with {DropValues}.
   LiftoffRegister PeekToRegister(int index, LiftoffRegList pinned);
 
+  void DropValues(int count);
+
   // Ensure that the loop inputs are either in a register or spilled to the
   // stack, so that we can merge different values on the back-edge.
   void PrepareLoopArgs(int num);

src/wasm/baseline/liftoff-compiler.cc

@@ -1713,12 +1713,7 @@ class LiftoffCompiler {
     __ PushRegister(ValueType::Ref(arg.type.heap_type(), kNonNullable), obj);
   }
 
-  void Drop(FullDecoder* decoder) {
-    auto& slot = __ cache_state()->stack_state.back();
-    // If the dropped slot contains a register, decrement it's use count.
-    if (slot.is_reg()) __ cache_state()->dec_used(slot.reg());
-    __ cache_state()->stack_state.pop_back();
-  }
+  void Drop(FullDecoder* decoder) { __ DropValues(1); }
 
   void TraceFunctionExit(FullDecoder* decoder) {
     DEBUG_CODE_COMMENT("trace function exit");
@@ -3415,7 +3410,7 @@ class LiftoffCompiler {
     LiftoffRegister expected = pinned.set(__ PopToRegister(pinned));
 
     // Pop the index from the stack.
-    __ cache_state()->stack_state.pop_back(1);
+    __ DropValues(1);
 
     LiftoffRegister result = expected;
 
@@ -3531,7 +3526,7 @@ class LiftoffCompiler {
     __ CallRuntimeStub(target);
     DefineSafepoint();
     // Pop parameters from the value stack.
-    __ cache_state()->stack_state.pop_back(3);
+    __ DropValues(3);
 
     RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
 
@@ -3574,7 +3569,7 @@ class LiftoffCompiler {
     __ CallRuntimeStub(WasmCode::kWasmAtomicNotify);
     DefineSafepoint();
     // Pop parameters from the value stack.
-    __ cache_state()->stack_state.pop_back(2);
+    __ DropValues(2);
 
     RegisterDebugSideTableEntry(DebugSideTableBuilder::kDidSpill);
 

test/mjsunit/regress/wasm/regress-1145135.js

+// Copyright 2021 the V8 project authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+// Flags: --wasm-staging
+
+load('test/mjsunit/wasm/wasm-module-builder.js');
+
+const builder = new WasmModuleBuilder();
+builder.addMemory(1, 1, false, true);
+builder.addGlobal(kWasmI32, 1);
+builder.addFunction(undefined, kSig_v_i)
+  .addLocals(kWasmI32, 5)
+  .addBody([
+// signature: v_i
+// body:
+kExprGlobalGet, 0x00,  // global.get
+kExprI32Const, 0x10,  // i32.const
+kExprI32Sub,  // i32.sub
+kExprLocalTee, 0x02,  // local.tee
+kExprGlobalSet, 0x00,  // global.set
+kExprBlock, kWasmStmt,  // block @12
+  kExprLocalGet, 0x00,  // local.get
+  kExprI32LoadMem, 0x02, 0x00,  // i32.load
+  kExprI32Eqz,  // i32.eqz
+  kExprIf, kWasmStmt,  // if @20
+    kExprLocalGet, 0x02,  // local.get
+    kExprI32Const, 0x00,  // i32.const
+    kExprI32StoreMem, 0x02, 0x0c,  // i32.store
+    kExprLocalGet, 0x00,  // local.get
+    kExprI32Const, 0x20,  // i32.const
+    kExprI32Add,  // i32.add
+    kExprLocalSet, 0x05,  // local.set
+    kExprLocalGet, 0x00,  // local.get
+    kExprI32Const, 0x00,  // i32.const
+    kExprI32Const, 0x01,  // i32.const
+    kAtomicPrefix, kExprI32AtomicCompareExchange, 0x02, 0x20,  // i32.atomic.cmpxchng32
+]);
+assertThrows(() => builder.toModule(), WebAssembly.CompileError);