Bug 1764366 part 3 - Shrink MachineState. r=nbp,iain

The MachineState class became pretty large, especially on ARM64 it was about 1 KB due
to the large number of (float) registers.

This patch fixes that by refactoring MachineState to store a pointer to the register
dumps (for bailouts) or the register sets + spill location (for safepoints).
This lets us recover the register address when we want to read/write it.

Differential Revision: https://phabricator.services.mozilla.com/D143588

js/src/jit/Bailouts.h

@@ -15,9 +15,9 @@
 #include "jstypes.h"
 
 #include "jit/IonTypes.h"  // js::jit::Bailout{Id,Kind}, js::jit::SnapshotOffset
-#include "jit/Registers.h"  // js::jit::MachineState
-#include "js/TypeDecls.h"   // jsbytecode
-#include "vm/JSContext.h"   // JSContext
+#include "jit/MachineState.h"  // js::jit::MachineState
+#include "js/TypeDecls.h"      // jsbytecode
+#include "vm/JSContext.h"      // JSContext
 
 namespace js {
 

js/src/jit/JSJitFrameIter.cpp

@@ -206,30 +206,16 @@ MachineState JSJitFrameIter::machineState() const {
   }
 
   SafepointReader reader(ionScript(), safepoint());
+
+  FloatRegisterSet fregs = reader.allFloatSpills().set().reduceSetForPush();
+  GeneralRegisterSet regs = reader.allGprSpills().set();
+
   uintptr_t* spill = spillBase();
-  MachineState machine;
-
-  for (GeneralRegisterBackwardIterator iter(reader.allGprSpills()); iter.more();
-       ++iter) {
-    machine.setRegisterLocation(*iter, --spill);
-  }
-
-  uint8_t* spillAlign = alignDoubleSpill(reinterpret_cast<uint8_t*>(spill));
-
+  uint8_t* spillAlign =
+      alignDoubleSpill(reinterpret_cast<uint8_t*>(spill - regs.size()));
   char* floatSpill = reinterpret_cast<char*>(spillAlign);
-  FloatRegisterSet fregs = reader.allFloatSpills().set();
-  fregs = fregs.reduceSetForPush();
-  for (FloatRegisterBackwardIterator iter(fregs); iter.more(); ++iter) {
-    floatSpill -= (*iter).size();
-    for (uint32_t a = 0; a < (*iter).numAlignedAliased(); a++) {
-      // Only say that registers that actually start here start here.
-      // e.g. d0 should not start at s1, only at s0.
-      FloatRegister ftmp = (*iter).alignedAliased(a);
-      machine.setRegisterLocation(ftmp, (double*)floatSpill);
-    }
-  }
 
-  return machine;
+  return MachineState::FromSafepoint(fregs, regs, floatSpill, spill);
 }
 
 JitFrameLayout* JSJitFrameIter::jsFrame() const {

js/src/jit/JSJitFrameIter.h

@@ -12,6 +12,7 @@
 #include "jstypes.h"
 
 #include "jit/JitCode.h"
+#include "jit/MachineState.h"
 #include "jit/Snapshots.h"
 #include "js/ProfilingFrameIterator.h"
 #include "vm/JSFunction.h"

js/src/jit/JitFrames.cpp

@@ -2185,81 +2185,85 @@ bool InlineFrameIterator::isFunctionFrame() const { return !!calleeTemplate_; }
 
 bool InlineFrameIterator::isModuleFrame() const { return script()->module(); }
 
-MachineState MachineState::FromBailout(RegisterDump::GPRArray& regs,
-                                       RegisterDump::FPUArray& fpregs) {
-  MachineState machine;
+uintptr_t* MachineState::SafepointState::addressOfRegister(Register reg) const {
+  size_t offset = regs.offsetOfPushedRegister(reg);
 
-  for (unsigned i = 0; i < Registers::Total; i++) {
-    machine.setRegisterLocation(Register::FromCode(i), &regs[i].r);
-  }
-#ifdef JS_CODEGEN_ARM
-  float* fbase = (float*)&fpregs[0];
-  for (unsigned i = 0; i < FloatRegisters::TotalDouble; i++) {
-    machine.setRegisterLocation(FloatRegister(i, FloatRegister::Double),
-                                &fpregs[i].d);
-  }
-  for (unsigned i = 0; i < FloatRegisters::TotalSingle; i++) {
-    machine.setRegisterLocation(FloatRegister(i, FloatRegister::Single),
-                                (double*)&fbase[i]);
-#  ifdef ENABLE_WASM_SIMD
-#    error "More care needed here"
-#  endif
-  }
-#elif defined(JS_CODEGEN_MIPS32)
-  for (unsigned i = 0; i < FloatRegisters::TotalPhys; i++) {
-    machine.setRegisterLocation(
-        FloatRegister::FromIndex(i, FloatRegister::Double), &fpregs[i]);
-    machine.setRegisterLocation(
-        FloatRegister::FromIndex(i, FloatRegister::Single), &fpregs[i]);
-#  ifdef ENABLE_WASM_SIMD
-#    error "More care needed here"
-#  endif
-  }
-#elif defined(JS_CODEGEN_MIPS64)
-  for (unsigned i = 0; i < FloatRegisters::TotalPhys; i++) {
-    machine.setRegisterLocation(FloatRegister(i, FloatRegisters::Double),
-                                &fpregs[i]);
-    machine.setRegisterLocation(FloatRegister(i, FloatRegisters::Single),
-                                &fpregs[i]);
-#  ifdef ENABLE_WASM_SIMD
-#    error "More care needed here"
-#  endif
-  }
-#elif defined(JS_CODEGEN_X86) || defined(JS_CODEGEN_X64)
-  for (unsigned i = 0; i < FloatRegisters::TotalPhys; i++) {
-    machine.setRegisterLocation(FloatRegister(i, FloatRegisters::Single),
-                                &fpregs[i]);
-    machine.setRegisterLocation(FloatRegister(i, FloatRegisters::Double),
-                                &fpregs[i]);
-    machine.setRegisterLocation(FloatRegister(i, FloatRegisters::Simd128),
-                                &fpregs[i]);
-  }
-#elif defined(JS_CODEGEN_ARM64)
-  for (unsigned i = 0; i < FloatRegisters::TotalPhys; i++) {
-    machine.setRegisterLocation(
-        FloatRegister(FloatRegisters::Encoding(i), FloatRegisters::Single),
-        &fpregs[i]);
-    machine.setRegisterLocation(
-        FloatRegister(FloatRegisters::Encoding(i), FloatRegisters::Double),
-        &fpregs[i]);
-    // No SIMD support in bailouts, SIMD is internal to wasm
-  }
-#elif defined(JS_CODEGEN_LOONG64)
-  for (unsigned i = 0; i < FloatRegisters::TotalPhys; i++) {
-    machine.setRegisterLocation(
-        FloatRegister(FloatRegisters::Encoding(i), FloatRegisters::Single),
-        &fpregs[i]);
-    machine.setRegisterLocation(
-        FloatRegister(FloatRegisters::Encoding(i), FloatRegisters::Double),
-        &fpregs[i]);
-  }
+  MOZ_ASSERT((offset % sizeof(uintptr_t)) == 0);
+  uint32_t index = offset / sizeof(uintptr_t);
 
-#elif defined(JS_CODEGEN_NONE)
-  MOZ_CRASH();
-#else
-#  error "Unknown architecture!"
+#ifdef DEBUG
+  // Assert correctness with a slower algorithm in debug builds.
+  uint32_t expectedIndex = 0;
+  bool found = false;
+  for (GeneralRegisterBackwardIterator iter(regs); iter.more(); ++iter) {
+    expectedIndex++;
+    if (*iter == reg) {
+      found = true;
+      break;
+    }
+  }
+  MOZ_ASSERT(found);
+  MOZ_ASSERT(expectedIndex == index);
 #endif
-  return machine;
+
+  return spillBase - index;
+}
+
+char* MachineState::SafepointState::addressOfRegister(FloatRegister reg) const {
+  // Note: this could be optimized similar to the GPR case above by implementing
+  // offsetOfPushedRegister for FloatRegisterSet. Float register sets are
+  // complicated though and this case is very uncommon: it's only reachable for
+  // exception bailouts with live float registers.
+  MOZ_ASSERT(!reg.isSimd128());
+  char* ptr = floatSpillBase;
+  for (FloatRegisterBackwardIterator iter(floatRegs); iter.more(); ++iter) {
+    ptr -= (*iter).size();
+    for (uint32_t a = 0; a < (*iter).numAlignedAliased(); a++) {
+      // Only say that registers that actually start here start here.
+      // e.g. d0 should not start at s1, only at s0.
+      FloatRegister ftmp = (*iter).alignedAliased(a);
+      if (ftmp == reg) {
+        return ptr;
+      }
+    }
+  }
+  MOZ_CRASH("Invalid register");
+}
+
+uintptr_t MachineState::read(Register reg) const {
+  if (state_.is<BailoutState>()) {
+    return state_.as<BailoutState>().regs[reg.code()].r;
+  }
+  if (state_.is<SafepointState>()) {
+    uintptr_t* addr = state_.as<SafepointState>().addressOfRegister(reg);
+    return *addr;
+  }
+  MOZ_CRASH("Invalid state");
+}
+
+double MachineState::read(FloatRegister reg) const {
+  if (state_.is<BailoutState>()) {
+    uint32_t offset = reg.getRegisterDumpOffsetInBytes();
+
+    MOZ_ASSERT((offset % sizeof(FloatRegisters::RegisterContent)) == 0);
+    uint32_t index = offset / sizeof(FloatRegisters::RegisterContent);
+
+    return state_.as<BailoutState>().floatRegs[index].d;
+  }
+  if (state_.is<SafepointState>()) {
+    char* addr = state_.as<SafepointState>().addressOfRegister(reg);
+    return *reinterpret_cast<double*>(addr);
+  }
+  MOZ_CRASH("Invalid state");
+}
+
+void MachineState::write(Register reg, uintptr_t value) const {
+  if (state_.is<SafepointState>()) {
+    uintptr_t* addr = state_.as<SafepointState>().addressOfRegister(reg);
+    *addr = value;
+    return;
+  }
+  MOZ_CRASH("Invalid state");
 }
 
 bool InlineFrameIterator::isConstructing() const {

js/src/jit/JitFrames.h

@@ -13,6 +13,7 @@
 #include <stdint.h>
 
 #include "jit/CalleeToken.h"
+#include "jit/MachineState.h"
 #include "jit/Registers.h"
 #include "js/Id.h"
 #include "js/TypeDecls.h"

js/src/jit/MachineState.h

@@ -0,0 +1,109 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*-
+ * vim: set ts=8 sts=2 et sw=2 tw=80:
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef jit_MachineState_h
+#define jit_MachineState_h
+
+#include "mozilla/Assertions.h"
+#include "mozilla/Attributes.h"
+#include "mozilla/Variant.h"
+
+#include <stdint.h>
+
+#include "jit/Registers.h"
+#include "jit/RegisterSets.h"
+
+namespace js::jit {
+
+// Information needed to recover machine register state. This supports two
+// different modes:
+//
+// * Bailouts: all registers are pushed on the stack as part of the bailout
+//   process, so MachineState simply points to these FPU/GPR arrays.
+//   See RegisterDump and BailoutStack.
+//
+// * Safepoints: live registers are pushed on the stack before a VM call, so
+//   MachineState stores the register sets and a pointer to the stack memory
+//   where these registers were pushed. This is also used by exception bailouts.
+class MOZ_STACK_CLASS MachineState {
+  struct NullState {};
+
+  struct BailoutState {
+    RegisterDump::FPUArray& floatRegs;
+    RegisterDump::GPRArray& regs;
+
+    BailoutState(RegisterDump::FPUArray& floatRegs,
+                 RegisterDump::GPRArray& regs)
+        : floatRegs(floatRegs), regs(regs) {}
+  };
+
+  struct SafepointState {
+    FloatRegisterSet floatRegs;
+    GeneralRegisterSet regs;
+    // Pointers to the start of the pushed |floatRegs| and |regs| on the stack.
+    // This is the value of the stack pointer right before the first register
+    // was pushed.
+    char* floatSpillBase;
+    uintptr_t* spillBase;
+
+    SafepointState(const FloatRegisterSet& floatRegs,
+                   const GeneralRegisterSet& regs, char* floatSpillBase,
+                   uintptr_t* spillBase)
+        : floatRegs(floatRegs),
+          regs(regs),
+          floatSpillBase(floatSpillBase),
+          spillBase(spillBase) {}
+    uintptr_t* addressOfRegister(Register reg) const;
+    char* addressOfRegister(FloatRegister reg) const;
+  };
+  using State = mozilla::Variant<NullState, BailoutState, SafepointState>;
+  State state_{NullState()};
+
+ public:
+  MachineState() = default;
+  MachineState(const MachineState& other) = default;
+  MachineState& operator=(const MachineState& other) = default;
+
+  static MachineState FromBailout(RegisterDump::GPRArray& regs,
+                                  RegisterDump::FPUArray& fpregs) {
+    MachineState res;
+    res.state_.emplace<BailoutState>(fpregs, regs);
+    return res;
+  }
+
+  static MachineState FromSafepoint(const FloatRegisterSet& floatRegs,
+                                    const GeneralRegisterSet& regs,
+                                    char* floatSpillBase,
+                                    uintptr_t* spillBase) {
+    MachineState res;
+    res.state_.emplace<SafepointState>(floatRegs, regs, floatSpillBase,
+                                       spillBase);
+    return res;
+  }
+
+  bool has(Register reg) const {
+    if (state_.is<BailoutState>()) {
+      return true;
+    }
+    return state_.as<SafepointState>().regs.hasRegisterIndex(reg);
+  }
+  bool has(FloatRegister reg) const {
+    if (state_.is<BailoutState>()) {
+      return true;
+    }
+    return state_.as<SafepointState>().floatRegs.hasRegisterIndex(reg);
+  }
+
+  uintptr_t read(Register reg) const;
+  double read(FloatRegister reg) const;
+
+  // Used by moving GCs to update pointers.
+  void write(Register reg, uintptr_t value) const;
+};
+
+}  // namespace js::jit
+
+#endif /* jit_MachineState_h */

js/src/jit/RegisterSets.h

@@ -9,6 +9,7 @@
 
 #include "mozilla/Assertions.h"
 #include "mozilla/Attributes.h"
+#include "mozilla/Variant.h"
 
 #include <new>
 #include <stddef.h>
@@ -372,6 +373,11 @@ class TypedRegisterSet {
     return T::ReduceSetForPush(*this);
   }
   uint32_t getPushSizeInBytes() const { return T::GetPushSizeInBytes(*this); }
+
+  size_t offsetOfPushedRegister(RegType reg) const {
+    MOZ_ASSERT(hasRegisterIndex(reg));
+    return T::OffsetOfPushedRegister(bits(), reg);
+  }
 };
 
 using GeneralRegisterSet = TypedRegisterSet<Register>;

js/src/jit/Registers.h

@@ -92,6 +92,13 @@ struct Register {
   static uint32_t SetSize(SetType x) { return Codes::SetSize(x); }
   static uint32_t FirstBit(SetType x) { return Codes::FirstBit(x); }
   static uint32_t LastBit(SetType x) { return Codes::LastBit(x); }
+
+  // Returns the offset of |reg| on the stack, assuming all registers in |set|
+  // were pushed in order (e.g. by |PushRegsInMask|). This is computed by
+  // clearing the lower bits (registers that were pushed later).
+  static size_t OffsetOfPushedRegister(SetType set, Register reg) {
+    return sizeof(Codes::RegisterContent) * Codes::SetSize(set >> reg.code());
+  }
 };
 
 // Architectures where the stack pointer is not a plain register with a standard
@@ -220,61 +227,6 @@ class RegisterDump {
   }
 };
 
-// Information needed to recover machine register state. This records the
-// location of spilled register and not the content of the spilled
-// registers. Thus we can safely assume that this structure is unchanged, even
-// if the GC pointers mapped by this structure are relocated.
-class MachineState {
-  mozilla::Array<Registers::RegisterContent*, Registers::Total> regs_;
-  mozilla::Array<FloatRegisters::RegisterContent*, FloatRegisters::Total>
-      fpregs_;
-
- public:
-  MachineState() {
-#ifndef JS_CODEGEN_NONE
-    for (uintptr_t i = 0; i < Registers::Total; i++) {
-      regs_[i] = reinterpret_cast<Registers::RegisterContent*>(i + 0x100);
-    }
-    for (uintptr_t i = 0; i < FloatRegisters::Total; i++) {
-      fpregs_[i] =
-          reinterpret_cast<FloatRegisters::RegisterContent*>(i + 0x200);
-    }
-#endif
-  }
-
-  static MachineState FromBailout(RegisterDump::GPRArray& regs,
-                                  RegisterDump::FPUArray& fpregs);
-
-  void setRegisterLocation(Register reg, uintptr_t* up) {
-    regs_[reg.code()] = (Registers::RegisterContent*)up;
-  }
-  void setRegisterLocation(FloatRegister reg, float* fp) {
-    MOZ_ASSERT(reg.isSingle());
-    fpregs_[reg.code()] = (FloatRegisters::RegisterContent*)fp;
-  }
-  void setRegisterLocation(FloatRegister reg, double* dp) {
-    fpregs_[reg.code()] = (FloatRegisters::RegisterContent*)dp;
-  }
-  void setRegisterLocation(FloatRegister reg,
-                           FloatRegisters::RegisterContent* rp) {
-    fpregs_[reg.code()] = rp;
-  }
-
-  bool has(Register reg) const { return regs_[reg.code()] != nullptr; }
-  bool has(FloatRegister reg) const { return fpregs_[reg.code()] != nullptr; }
-  uintptr_t read(Register reg) const { return regs_[reg.code()]->r; }
-  double read(FloatRegister reg) const { return fpregs_[reg.code()]->d; }
-  void write(Register reg, uintptr_t value) const {
-    regs_[reg.code()]->r = value;
-  }
-  const Registers::RegisterContent* address(Register reg) const {
-    return regs_[reg.code()];
-  }
-  const FloatRegisters::RegisterContent* address(FloatRegister reg) const {
-    return fpregs_[reg.code()];
-  }
-};
-
 // Class for mapping each register to an offset.
 class RegisterOffsets {
   mozilla::Array<uint32_t, Registers::Total> offsets_;