[wasm][arm64] Fix crash on atomic cmpxchg with large offset

Liftoff can currently run out of registers when compiling an atomic
compare-exchange instruction. In order to see this crash, the following
conditions must be met:

- The offset in the instruction doesn't fit in a 12-bit immediate
- Either FLAG_untrusted_code_mitigations is false, or trap handlers are
  enabled, so that AddMemoryMasking decides to do nothing

The fix proposed in this CL is just to defer allocation of a temporary
register until after CalculateActualAddress has finished, because it
might have also needed a temporary register.

Change-Id: I28225614dcdbe2bcc9e52208f1e806baac89c5f1
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2488840
Commit-Queue: Seth Brenith <seth.brenith@microsoft.com>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Cr-Commit-Position: refs/heads/master@{#70687}

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

@@ -687,11 +687,12 @@ void LiftoffAssembler::AtomicCompareExchange(
   }
 
   UseScratchRegisterScope temps(this);
-  Register store_result = temps.AcquireW();
 
   Register actual_addr = liftoff::CalculateActualAddress(
       this, dst_addr, offset_reg, offset_imm, temps.AcquireX());
 
+  Register store_result = temps.AcquireW();
+
   Label retry;
   Label done;
   Bind(&retry);

test/mjsunit/wasm/atomics.js

@@ -47,7 +47,7 @@ function GetAtomicCmpExchangeFunction(wasmExpression, alignment, offset) {
       kExprLocalGet, 1,
       kExprLocalGet, 2,
       kAtomicPrefix,
-      wasmExpression, alignment, offset])
+      wasmExpression, alignment, ...wasmSignedLeb(offset, 5)])
     .exportAs("main");
 
   // Instantiate module, get function exports
@@ -251,15 +251,15 @@ function Test8Op(operation, func) {
   Test8Op(Exchange, wasmExchange);
 })();
 
-function TestCmpExchange(func, buffer, params, size) {
-  for (let i = 0; i < buffer.length; i = inc(i)) {
+function TestCmpExchange(func, buffer, params, size, offset = 0) {
+  for (let i = 0; i + (offset / size) < buffer.length; i = inc(i)) {
     for (let j = 0; j < params.length; j++) {
       for (let k = 0; k < params.length; k++) {
-        buffer[i] = params[j];
+        buffer[i + (offset / size)] = params[j];
         let loaded = func(i * size, params[k], params[j]) >>> 0;
         let expected = (params[k] == loaded) ? params[j] : loaded;
         assertEquals(loaded, params[j]);
-        assertEquals(expected, buffer[i]);
+        assertEquals(expected, buffer[i + (offset / size)]);
       }
     }
   }
@@ -268,11 +268,14 @@ function TestCmpExchange(func, buffer, params, size) {
 
 (function TestAtomicCompareExchange() {
   print(arguments.callee.name);
+  // Offset is big enough to not fit in a 12-bit immediate on arm64, but small
+  // enough to fit in the maxSize wasm pages.
+  const offset = 0x1234;
   let wasmCmpExchange =
-      GetAtomicCmpExchangeFunction(kExprI32AtomicCompareExchange, 2, 0);
+      GetAtomicCmpExchangeFunction(kExprI32AtomicCompareExchange, 2, offset);
   let i32 = new Uint32Array(memory.buffer);
   let params = [0x00000001, 0x00000555, 0x00099999, 0xffffffff];
-  TestCmpExchange(wasmCmpExchange, i32, params, kMemtypeSize32);
+  TestCmpExchange(wasmCmpExchange, i32, params, kMemtypeSize32, offset);
 })();
 
 (function TestAtomicCompareExchange16U() {