Bug 1571446 part 3 - Combine JSScript's jitScript_ and warmUpCount_ fields in a single warmUpData_ field. r=tcampbell
authorJan de Mooij <jdemooij@mozilla.com>
Wed, 09 Oct 2019 09:51:19 +0000
changeset 496932 e753d23c1237313b3cc94ee509cbba9a5b6246d8
parent 496931 3a4012655a64037ce150dadaf10e6a77059f2c9e
child 496933 8d85038a038b2edb41ea3ec15d312ddf7ab6a2aa
push id36671
push usershindli@mozilla.com
push dateWed, 09 Oct 2019 16:04:03 +0000
treeherdermozilla-central@0efb4f268d16 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerstcampbell
bugs1571446
milestone71.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1571446 part 3 - Combine JSScript's jitScript_ and warmUpCount_ fields in a single warmUpData_ field. r=tcampbell The warm-up count is stored in ScriptWarmUpData until the script is warm enough for the Baseline Interpreter and the JitScript is created. At that point we use the warm-up count stored in JitScript. ScriptWarmUpData uses pointer tagging. This should make it easy to add new types for LazyScript data in the future. Differential Revision: https://phabricator.services.mozilla.com/D42319
js/src/jit/BaselineCodeGen.cpp
js/src/jit/CodeGenerator.cpp
js/src/jit/IonOptimizationLevels.cpp
js/src/jit/JitScript.cpp
js/src/jit/JitScript.h
js/src/jit/MacroAssembler-inl.h
js/src/jit/MacroAssembler.cpp
js/src/jit/MacroAssembler.h
js/src/vm/HelperThreads.cpp
js/src/vm/JSScript-inl.h
js/src/vm/JSScript.cpp
js/src/vm/JSScript.h
--- a/js/src/jit/BaselineCodeGen.cpp
+++ b/js/src/jit/BaselineCodeGen.cpp
@@ -1121,17 +1121,17 @@ void BaselineInterpreterCodeGen::emitIni
     // CalleeToken_Script.
     masm.andPtr(Imm32(uint32_t(CalleeTokenMask)), scratch1);
     masm.storePtr(nonFunctionEnv, frame.addressOfEnvironmentChain());
   }
   masm.bind(&done);
   masm.storePtr(scratch1, frame.addressOfInterpreterScript());
 
   // Initialize interpreterICEntry.
-  masm.loadPtr(Address(scratch1, JSScript::offsetOfJitScript()), scratch2);
+  masm.loadJitScript(scratch1, scratch2);
   masm.computeEffectiveAddress(
       Address(scratch2, JitScript::offsetOfICEntries()), scratch2);
   masm.storePtr(scratch2, frame.addressOfInterpreterICEntry());
 
   // Initialize interpreter pc.
   masm.loadPtr(Address(scratch1, JSScript::offsetOfScriptData()), scratch1);
   masm.loadPtr(Address(scratch1, RuntimeScriptData::offsetOfISD()), scratch1);
   masm.addPtr(Imm32(ImmutableScriptData::offsetOfCode()), scratch1);
@@ -1279,19 +1279,22 @@ bool BaselineCompilerCodeGen::emitWarmUp
   // Emit no warm-up counter increments if Ion is not enabled or if the script
   // will never be Ion-compileable.
   if (!handler.maybeIonCompileable()) {
     return true;
   }
 
   Register scriptReg = R2.scratchReg();
   Register countReg = R0.scratchReg();
-  Address warmUpCounterAddr(scriptReg, JSScript::offsetOfWarmUpCounter());
-
-  masm.movePtr(ImmGCPtr(script), scriptReg);
+
+  // Load the JitScript* in scriptReg.
+  masm.movePtr(ImmPtr(script->jitScript()), scriptReg);
+
+  // Bump warm-up counter.
+  Address warmUpCounterAddr(scriptReg, JitScript::offsetOfWarmUpCount());
   masm.load32(warmUpCounterAddr, countReg);
   masm.add32(Imm32(1), countReg);
   masm.store32(countReg, warmUpCounterAddr);
 
   if (JSOp(*pc) == JSOP_LOOPENTRY) {
     // If this is a loop inside a catch or finally block, increment the warmup
     // counter but don't attempt OSR (Ion only compiles the try block).
     if (handler.analysis().info(pc).loopEntryInCatchOrFinally) {
@@ -1308,17 +1311,16 @@ bool BaselineCompilerCodeGen::emitWarmUp
 
   const OptimizationInfo* info =
       IonOptimizations.get(IonOptimizations.firstLevel());
   uint32_t warmUpThreshold = info->compilerWarmUpThreshold(script, pc);
   masm.branch32(Assembler::LessThan, countReg, Imm32(warmUpThreshold), &done);
 
   // Do nothing if Ion is already compiling this script off-thread or if Ion has
   // been disabled for this script.
-  masm.movePtr(ImmPtr(script->jitScript()), scriptReg);
   masm.loadPtr(Address(scriptReg, JitScript::offsetOfIonScript()), scriptReg);
   masm.branchPtr(Assembler::Equal, scriptReg, ImmPtr(IonCompilingScriptPtr),
                  &done);
   masm.branchPtr(Assembler::Equal, scriptReg, ImmPtr(IonDisabledScriptPtr),
                  &done);
 
   // Try to compile and/or finish a compilation.
   if (JSOp(*pc) == JSOP_LOOPENTRY) {
@@ -1414,29 +1416,31 @@ bool BaselineCompilerCodeGen::emitWarmUp
   return true;
 }
 
 template <>
 bool BaselineInterpreterCodeGen::emitWarmUpCounterIncrement() {
   Register scriptReg = R2.scratchReg();
   Register countReg = R0.scratchReg();
 
+  // Load the JitScript* in scriptReg.
+  loadScript(scriptReg);
+  masm.loadJitScript(scriptReg, scriptReg);
+
   // Bump warm-up counter.
-  Address warmUpCounterAddr(scriptReg, JSScript::offsetOfWarmUpCounter());
-  loadScript(scriptReg);
+  Address warmUpCounterAddr(scriptReg, JitScript::offsetOfWarmUpCount());
   masm.load32(warmUpCounterAddr, countReg);
   masm.add32(Imm32(1), countReg);
   masm.store32(countReg, warmUpCounterAddr);
 
   // If the script is warm enough for Baseline compilation, call into the VM to
   // compile it.
   Label done;
   masm.branch32(Assembler::BelowOrEqual, countReg,
                 Imm32(JitOptions.baselineJitWarmUpThreshold), &done);
-  masm.loadPtr(Address(scriptReg, JSScript::offsetOfJitScript()), scriptReg);
   masm.branchPtr(Assembler::Equal,
                  Address(scriptReg, JitScript::offsetOfBaselineScript()),
                  ImmPtr(BaselineDisabledScriptPtr), &done);
   {
     prepareVMCall();
 
     masm.PushBaselineFramePtr(BaselineFrameReg, R0.scratchReg());
 
@@ -6046,17 +6050,17 @@ bool BaselineCodeGen<Handler>::emitEnter
                                                       Register resumeIndex,
                                                       Register scratch) {
   // Resume in either the BaselineScript (if present) or Baseline Interpreter.
 
   static_assert(BaselineDisabledScript == 0x1,
                 "Comparison below requires specific sentinel encoding");
 
   Label noBaselineScript;
-  masm.loadPtr(Address(script, JSScript::offsetOfJitScript()), scratch);
+  masm.loadJitScript(script, scratch);
   masm.loadPtr(Address(scratch, JitScript::offsetOfBaselineScript()), scratch);
   masm.branchPtr(Assembler::BelowOrEqual, scratch,
                  ImmPtr(BaselineDisabledScriptPtr), &noBaselineScript);
   masm.load32(Address(scratch, BaselineScript::offsetOfResumeEntriesOffset()),
               script);
   masm.addPtr(scratch, script);
   masm.loadPtr(
       BaseIndex(script, resumeIndex, ScaleFromElemWidth(sizeof(uintptr_t))),
@@ -6104,25 +6108,23 @@ bool BaselineCodeGen<Handler>::emitGener
 
   // Branch to |interpret| to resume the generator in the C++ interpreter if the
   // script does not have a JitScript. Note that we don't relazify generator
   // scripts (asserted in JSFunction::maybeRelazify) so the function is
   // guaranteed to be non-lazy.
   Label interpret;
   Register scratch1 = regs.takeAny();
   masm.loadPtr(Address(callee, JSFunction::offsetOfScript()), scratch1);
-  masm.branchPtr(Assembler::Equal,
-                 Address(scratch1, JSScript::offsetOfJitScript()),
-                 ImmPtr(nullptr), &interpret);
+  masm.branchIfScriptHasNoJitScript(scratch1, &interpret);
 
 #ifdef JS_TRACE_LOGGING
   if (JS::TraceLoggerSupported()) {
     // TODO (bug 1565788): add Baseline Interpreter support.
     MOZ_CRASH("Unimplemented Baseline Interpreter TraceLogger support");
-    masm.loadPtr(Address(scratch1, JSScript::offsetOfJitScript()), scratch1);
+    masm.loadJitScript(scratch1, scratch1);
     Address baselineAddr(scratch1, JitScript::offsetOfBaselineScript());
     masm.loadPtr(baselineAddr, scratch1);
     if (!emitTraceLoggerResume(scratch1, regs)) {
       return false;
     }
   }
 #endif
 
@@ -6427,17 +6429,17 @@ bool BaselineInterpreterCodeGen::emit_JS
   // scratch1 := scratch1 * sizeof(ICEntry)
   static_assert(sizeof(ICEntry) == 8 || sizeof(ICEntry) == 16,
                 "shift below depends on ICEntry size");
   uint32_t shift = (sizeof(ICEntry) == 16) ? 4 : 3;
   masm.lshiftPtr(Imm32(shift), scratch1);
 
   // Compute ICEntry* and store to frame->interpreterICEntry.
   loadScript(scratch2);
-  masm.loadPtr(Address(scratch2, JSScript::offsetOfJitScript()), scratch2);
+  masm.loadJitScript(scratch2, scratch2);
   masm.computeEffectiveAddress(
       BaseIndex(scratch2, scratch1, TimesOne, JitScript::offsetOfICEntries()),
       scratch2);
   masm.storePtr(scratch2, frame.addressOfInterpreterICEntry());
   return true;
 }
 
 template <typename Handler>
--- a/js/src/jit/CodeGenerator.cpp
+++ b/js/src/jit/CodeGenerator.cpp
@@ -13326,19 +13326,32 @@ void CodeGenerator::visitRecompileCheck(
           oolCallVM<Fn, IonForcedInvalidation>(ins, ArgList(), StoreNothing());
     } else if (ins->mir()->forceRecompilation()) {
       ool = oolCallVM<Fn, IonForcedRecompile>(ins, ArgList(), StoreNothing());
     } else {
       ool = oolCallVM<Fn, IonRecompile>(ins, ArgList(), StoreNothing());
     }
   }
 
+  JitScript* jitScript = ins->mir()->script()->jitScript();
+
+  // The code depends on the JitScript* not being discarded without also
+  // invalidating Ion code. Assert this.
+#ifdef DEBUG
+  Label ok;
+  masm.movePtr(ImmGCPtr(ins->mir()->script()), tmp);
+  masm.loadJitScript(tmp, tmp);
+  masm.branchPtr(Assembler::Equal, tmp, ImmPtr(jitScript), &ok);
+  masm.assumeUnreachable("Didn't find JitScript?");
+  masm.bind(&ok);
+#endif
+
   // Check if warm-up counter is high enough.
   AbsoluteAddress warmUpCount =
-      AbsoluteAddress(ins->mir()->script()->addressOfWarmUpCounter());
+      AbsoluteAddress(jitScript->addressOfWarmUpCount());
   if (ins->mir()->increaseWarmUpCounter()) {
     masm.load32(warmUpCount, tmp);
     masm.add32(Imm32(1), tmp);
     masm.store32(tmp, warmUpCount);
     if (ins->mir()->checkCounter()) {
       masm.branch32(Assembler::BelowOrEqual, tmp,
                     Imm32(ins->mir()->recompileThreshold()), &done);
     }
--- a/js/src/jit/IonOptimizationLevels.cpp
+++ b/js/src/jit/IonOptimizationLevels.cpp
@@ -4,16 +4,18 @@
  * 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/. */
 
 #include "jit/IonOptimizationLevels.h"
 
 #include "jit/Ion.h"
 #include "vm/JSScript.h"
 
+#include "vm/JSScript-inl.h"
+
 using namespace js;
 using namespace js::jit;
 
 namespace js {
 namespace jit {
 
 const OptimizationLevelInfo IonOptimizations;
 
--- a/js/src/jit/JitScript.cpp
+++ b/js/src/jit/JitScript.cpp
@@ -50,16 +50,19 @@ JitScript::JitScript(JSScript* script, u
       typeSetOffset_(typeSetOffset),
       bytecodeTypeMapOffset_(bytecodeTypeMapOffset),
       allocBytes_(allocBytes) {
   setTypesGeneration(script->zone()->types.generation);
 
   uint8_t* base = reinterpret_cast<uint8_t*>(this);
   DefaultInitializeElements<StackTypeSet>(base + typeSetOffset, numTypeSets());
 
+  // Initialize the warm-up count from the count stored in the script.
+  warmUpCount_ = script->getWarmUpCount();
+
   // Ensure the baselineScript_ and ionScript_ fields match the BaselineDisabled
   // and IonDisabled script flags.
   if (!script->canBaselineCompile()) {
     setBaselineScriptImpl(script, BaselineDisabledScriptPtr);
   }
   if (!script->canIonCompile()) {
     setIonScriptImpl(script, IonDisabledScriptPtr);
   }
@@ -137,17 +140,17 @@ bool JSScript::createJitScript(JSContext
       [&] { jitScript->prepareForDestruction(cx->zone()); });
 
   if (!jitScript->initICEntriesAndBytecodeTypeMap(cx, this)) {
     return false;
   }
 
   MOZ_ASSERT(!hasJitScript());
   prepareForDestruction.release();
-  jitScript_ = jitScript.release();
+  warmUpData_.setJitScript(jitScript.release());
   AddCellMemory(this, allocSize.value(), MemoryUse::JitScript);
 
   // We have a JitScript so we can set the script's jitCodeRaw_ pointer to the
   // Baseline Interpreter code.
   updateJitCodeRaw(cx->runtime());
 
 #ifdef DEBUG
   AutoSweepJitScript sweep(this);
@@ -186,17 +189,17 @@ void JSScript::maybeReleaseJitScript(JSF
 void JSScript::releaseJitScript(JSFreeOp* fop) {
   MOZ_ASSERT(hasJitScript());
   MOZ_ASSERT(!hasBaselineScript());
   MOZ_ASSERT(!hasIonScript());
 
   fop->removeCellMemory(this, jitScript()->allocBytes(), MemoryUse::JitScript);
 
   JitScript::Destroy(zone(), jitScript());
-  jitScript_ = nullptr;
+  warmUpData_.clearJitScript();
   updateJitCodeRaw(fop->runtime());
 }
 
 void JSScript::releaseJitScriptOnFinalize(JSFreeOp* fop) {
   MOZ_ASSERT(hasJitScript());
 
   if (hasIonScript()) {
     IonScript* ion = jitScript()->clearIonScript(fop, this);
--- a/js/src/jit/JitScript.h
+++ b/js/src/jit/JitScript.h
@@ -2,16 +2,18 @@
  * 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_JitScript_h
 #define jit_JitScript_h
 
+#include "mozilla/Atomics.h"
+
 #include "jit/BaselineIC.h"
 #include "js/UniquePtr.h"
 #include "vm/TypeInference.h"
 
 class JSScript;
 
 namespace js {
 namespace jit {
@@ -173,16 +175,23 @@ class alignas(uintptr_t) JitScript final
   // Baseline code for the script. Either nullptr, BaselineDisabledScriptPtr or
   // a valid BaselineScript*.
   BaselineScript* baselineScript_ = nullptr;
 
   // Ion code for this script. Either nullptr, IonDisabledScriptPtr,
   // IonCompilingScriptPtr or a valid IonScript*.
   IonScript* ionScript_ = nullptr;
 
+  // Number of times the script has been called or has had backedges taken.
+  // Reset if the script's JIT code is forcibly discarded. See also the
+  // ScriptWarmUpData class.
+  mozilla::Atomic<uint32_t, mozilla::Relaxed,
+                  mozilla::recordreplay::Behavior::DontPreserve>
+      warmUpCount_ = {};
+
   // Offset of the StackTypeSet array.
   uint32_t typeSetOffset_ = 0;
 
   // Offset of the bytecode type map.
   uint32_t bytecodeTypeMapOffset_ = 0;
 
   // This field is used to avoid binary searches for the sought entry when
   // bytecode map queries are in linear order.
@@ -387,20 +396,30 @@ class alignas(uintptr_t) JitScript final
 
   static void Destroy(Zone* zone, JitScript* script);
 
   static constexpr size_t offsetOfICEntries() { return sizeof(JitScript); }
 
   static constexpr size_t offsetOfJitCodeSkipArgCheck() {
     return offsetof(JitScript, jitCodeSkipArgCheck_);
   }
-  static size_t offsetOfBaselineScript() {
+  static constexpr size_t offsetOfBaselineScript() {
     return offsetof(JitScript, baselineScript_);
   }
-  static size_t offsetOfIonScript() { return offsetof(JitScript, ionScript_); }
+  static constexpr size_t offsetOfIonScript() {
+    return offsetof(JitScript, ionScript_);
+  }
+  static constexpr size_t offsetOfWarmUpCount() {
+    return offsetof(JitScript, warmUpCount_);
+  }
+
+  uint32_t warmUpCount() const { return warmUpCount_; }
+  uint32_t* addressOfWarmUpCount() {
+    return reinterpret_cast<uint32_t*>(&warmUpCount_);
+  }
 
 #ifdef DEBUG
   void printTypes(JSContext* cx, HandleScript script);
 #endif
 
   void prepareForDestruction(Zone* zone) {
     // When the script contains pointers to nursery things, the store buffer can
     // contain entries that point into the fallback stub space. Since we can
--- a/js/src/jit/MacroAssembler-inl.h
+++ b/js/src/jit/MacroAssembler-inl.h
@@ -385,16 +385,46 @@ void MacroAssembler::branchIfInterpreted
                                          Label* label) {
   int32_t flags = FunctionFlags::INTERPRETED | FunctionFlags::INTERPRETED_LAZY;
   if (!isConstructing) {
     flags |= FunctionFlags::WASM_JIT_ENTRY;
   }
   branchTestFunctionFlags(fun, flags, Assembler::NonZero, label);
 }
 
+void MacroAssembler::branchIfScriptHasJitScript(Register script, Label* label) {
+  static_assert(ScriptWarmUpData::JitScriptTag == 0,
+                "Code below depends on tag value");
+  branchTestPtr(Assembler::Zero,
+                Address(script, JSScript::offsetOfWarmUpData()),
+                Imm32(ScriptWarmUpData::TagMask), label);
+}
+
+void MacroAssembler::branchIfScriptHasNoJitScript(Register script,
+                                                  Label* label) {
+  static_assert(ScriptWarmUpData::JitScriptTag == 0,
+                "Code below depends on tag value");
+  branchTestPtr(Assembler::NonZero,
+                Address(script, JSScript::offsetOfWarmUpData()),
+                Imm32(ScriptWarmUpData::TagMask), label);
+}
+
+void MacroAssembler::loadJitScript(Register script, Register dest) {
+#ifdef DEBUG
+  Label ok;
+  branchIfScriptHasJitScript(script, &ok);
+  assumeUnreachable("Script has no JitScript!");
+  bind(&ok);
+#endif
+
+  static_assert(ScriptWarmUpData::JitScriptTag == 0,
+                "Code below depends on tag value");
+  loadPtr(Address(script, JSScript::offsetOfWarmUpData()), dest);
+}
+
 void MacroAssembler::branchIfObjectEmulatesUndefined(Register objReg,
                                                      Register scratch,
                                                      Label* slowCheck,
                                                      Label* label) {
   // The branches to out-of-line code here implement a conservative version
   // of the JSObject::isWrapper test performed in EmulatesUndefined.
   loadObjClassUnsafe(objReg, scratch);
 
--- a/js/src/jit/MacroAssembler.cpp
+++ b/js/src/jit/MacroAssembler.cpp
@@ -1624,17 +1624,17 @@ void MacroAssembler::loadJitCodeRaw(Regi
           SelfHostedLazyScript::offsetOfJitCodeRaw(),
       "SelfHostedLazyScript and JSScript must use same layout for jitCodeRaw_");
   loadPtr(Address(func, JSFunction::offsetOfScript()), dest);
   loadPtr(Address(dest, JSScript::offsetOfJitCodeRaw()), dest);
 }
 
 void MacroAssembler::loadJitCodeNoArgCheck(Register func, Register dest) {
   loadPtr(Address(func, JSFunction::offsetOfScript()), dest);
-  loadPtr(Address(dest, JSScript::offsetOfJitScript()), dest);
+  loadJitScript(dest, dest);
   loadPtr(Address(dest, JitScript::offsetOfJitCodeSkipArgCheck()), dest);
 }
 
 void MacroAssembler::loadBaselineFramePtr(Register framePtr, Register dest) {
   if (framePtr != dest) {
     movePtr(framePtr, dest);
   }
   subPtr(Imm32(BaselineFrame::Size()), dest);
--- a/js/src/jit/MacroAssembler.h
+++ b/js/src/jit/MacroAssembler.h
@@ -1318,16 +1318,20 @@ class MacroAssembler : public MacroAssem
                                       Condition cond, Label* label);
 
   inline void branchIfFunctionHasNoJitEntry(Register fun, bool isConstructing,
                                             Label* label);
   inline void branchIfFunctionHasNoScript(Register fun, Label* label);
   inline void branchIfInterpreted(Register fun, bool isConstructing,
                                   Label* label);
 
+  inline void branchIfScriptHasJitScript(Register script, Label* label);
+  inline void branchIfScriptHasNoJitScript(Register script, Label* label);
+  inline void loadJitScript(Register script, Register dest);
+
   inline void branchFunctionKind(Condition cond,
                                  FunctionFlags::FunctionKind kind, Register fun,
                                  Register scratch, Label* label);
 
   void branchIfNotInterpretedConstructor(Register fun, Register scratch,
                                          Label* label);
 
   inline void branchIfObjectEmulatesUndefined(Register objReg, Register scratch,
--- a/js/src/vm/HelperThreads.cpp
+++ b/js/src/vm/HelperThreads.cpp
@@ -1607,18 +1607,20 @@ static bool IonBuilderHasHigherPriority(
   }
 
   // A script without an IonScript has precedence on one with.
   if (first->scriptHasIonScript() != second->scriptHasIonScript()) {
     return !first->scriptHasIonScript();
   }
 
   // A higher warm-up counter indicates a higher priority.
-  return first->script()->getWarmUpCount() / first->script()->length() >
-         second->script()->getWarmUpCount() / second->script()->length();
+  jit::JitScript* firstJitScript = first->script()->jitScript();
+  jit::JitScript* secondJitScript = second->script()->jitScript();
+  return firstJitScript->warmUpCount() / first->script()->length() >
+         secondJitScript->warmUpCount() / second->script()->length();
 }
 
 bool GlobalHelperThreadState::canStartIonCompile(
     const AutoLockHelperThreadState& lock) {
   return !ionWorklist(lock).empty() &&
          checkTaskThreadLimit<jit::IonBuilder*>(maxIonCompilationThreads());
 }
 
--- a/js/src/vm/JSScript-inl.h
+++ b/js/src/vm/JSScript-inl.h
@@ -226,9 +226,33 @@ inline void JSScript::disableIon() {
 inline js::jit::BaselineScript* JSScript::baselineScript() const {
   return jitScript()->baselineScript();
 }
 
 inline js::jit::IonScript* JSScript::ionScript() const {
   return jitScript()->ionScript();
 }
 
+inline uint32_t JSScript::getWarmUpCount() const {
+  if (warmUpData_.isWarmUpCount()) {
+    return warmUpData_.toWarmUpCount();
+  }
+  return warmUpData_.toJitScript()->warmUpCount_;
+}
+
+inline void JSScript::incWarmUpCounter(uint32_t amount) {
+  if (warmUpData_.isWarmUpCount()) {
+    warmUpData_.incWarmUpCount(amount);
+  } else {
+    warmUpData_.toJitScript()->warmUpCount_ += amount;
+  }
+}
+
+inline void JSScript::resetWarmUpCounterForGC() {
+  incWarmUpResetCounter();
+  if (warmUpData_.isWarmUpCount()) {
+    warmUpData_.resetWarmUpCount(0);
+  } else {
+    warmUpData_.toJitScript()->warmUpCount_ = 0;
+  }
+}
+
 #endif /* vm_JSScript_inl_h */
--- a/js/src/vm/JSScript.cpp
+++ b/js/src/vm/JSScript.cpp
@@ -4821,16 +4821,25 @@ void RuntimeScriptData::traceChildren(JS
 }
 
 void RuntimeScriptData::markForCrossZone(JSContext* cx) {
   for (uint32_t i = 0; i < natoms(); ++i) {
     cx->markAtom(atoms()[i]);
   }
 }
 
+void ScriptWarmUpData::trace(JSTracer* trc) {
+  if (isJitScript()) {
+    toJitScript()->trace(trc);
+    return;
+  }
+
+  MOZ_ASSERT(isWarmUpCount());
+}
+
 void JSScript::traceChildren(JSTracer* trc) {
   // NOTE: this JSScript may be partially initialized at this point.  E.g. we
   // may have created it and partially initialized it with
   // JSScript::Create(), but not yet finished initializing it with
   // fullyInitFromEmitter().
 
   // Trace base class fields.
   BaseScript::traceChildren(trc);
@@ -4842,19 +4851,17 @@ void JSScript::traceChildren(JSTracer* t
   if (data_) {
     data_->trace(trc);
   }
 
   if (scriptData()) {
     scriptData()->traceChildren(trc);
   }
 
-  if (jit::JitScript* jitScript = maybeJitScript()) {
-    jitScript->trace(trc);
-  }
+  warmUpData_.trace(trc);
 
   if (maybeLazyScript()) {
     TraceManuallyBarrieredEdge(trc, &lazyScript, "lazyScript");
   }
 
   if (trc->isMarkingTracer()) {
     GCMarker::fromTracer(trc)->markImplicitEdges(this);
   }
@@ -5382,19 +5389,24 @@ bool JSScript::mayReadFrameArgsDirectly(
 }
 
 void JSScript::resetWarmUpCounterToDelayIonCompilation() {
   // Reset the warm-up count only if it's greater than the BaselineCompiler
   // threshold. We do this to ensure this has no effect on Baseline compilation
   // because we don't want scripts to get stuck in the (Baseline) interpreter in
   // pathological cases.
 
-  if (warmUpCount > jit::JitOptions.baselineJitWarmUpThreshold) {
+  if (getWarmUpCount() > jit::JitOptions.baselineJitWarmUpThreshold) {
     incWarmUpResetCounter();
-    warmUpCount = jit::JitOptions.baselineJitWarmUpThreshold;
+    uint32_t newCount = jit::JitOptions.baselineJitWarmUpThreshold;
+    if (warmUpData_.isWarmUpCount()) {
+      warmUpData_.resetWarmUpCount(newCount);
+    } else {
+      warmUpData_.toJitScript()->warmUpCount_ = newCount;
+    }
   }
 }
 
 void JSScript::AutoDelazify::holdScript(JS::HandleFunction fun) {
   if (fun) {
     if (fun->realm()->isSelfHostingRealm()) {
       // The self-hosting realm is shared across runtimes, so we can't use
       // JSAutoRealm: it could cause races. Functions in the self-hosting
--- a/js/src/vm/JSScript.h
+++ b/js/src/vm/JSScript.h
@@ -2323,16 +2323,83 @@ struct RuntimeScriptDataHasher {
 
 class AutoLockScriptData;
 
 using RuntimeScriptDataTable =
     HashSet<RuntimeScriptData*, RuntimeScriptDataHasher, SystemAllocPolicy>;
 
 extern void SweepScriptData(JSRuntime* rt);
 
+// ScriptWarmUpData represents a pointer-sized field in JSScript that stores
+// one of the following:
+//
+// * The script's warm-up count. This is only used until the script has a
+//   JitScript. The Baseline Interpreter and JITs use the warm-up count stored
+//   in JitScript.
+//
+// * A pointer to the JitScript, when the script is warm enough for the Baseline
+//   Interpreter.
+//
+// Pointer tagging is used to distinguish those states.
+class ScriptWarmUpData {
+  static constexpr uintptr_t NumTagBits = 2;
+  static constexpr uint32_t MaxWarmUpCount = UINT32_MAX >> NumTagBits;
+
+ public:
+  // Public only for the JITs.
+  static constexpr uintptr_t TagMask = (1 << NumTagBits) - 1;
+  static constexpr uintptr_t JitScriptTag = 0;
+  static constexpr uintptr_t WarmUpCountTag = 1;
+
+ private:
+  uintptr_t data_ = 0 | WarmUpCountTag;
+
+  void setWarmUpCount(uint32_t count) {
+    count = std::min(count, MaxWarmUpCount);
+    data_ = (uintptr_t(count) << NumTagBits) | WarmUpCountTag;
+  }
+
+ public:
+  void trace(JSTracer* trc);
+
+  bool isWarmUpCount() const { return (data_ & TagMask) == WarmUpCountTag; }
+  bool isJitScript() const { return (data_ & TagMask) == JitScriptTag; }
+
+  uint32_t toWarmUpCount() const {
+    MOZ_ASSERT(isWarmUpCount());
+    return data_ >> NumTagBits;
+  }
+  void resetWarmUpCount(uint32_t count) {
+    MOZ_ASSERT(isWarmUpCount());
+    setWarmUpCount(count);
+  }
+  void incWarmUpCount(uint32_t amount) {
+    MOZ_ASSERT(isWarmUpCount());
+    data_ += uintptr_t(amount) << NumTagBits;
+  }
+
+  jit::JitScript* toJitScript() const {
+    MOZ_ASSERT(isJitScript());
+    static_assert(JitScriptTag == 0, "Code depends on JitScriptTag being zero");
+    return reinterpret_cast<jit::JitScript*>(data_);
+  }
+  void setJitScript(jit::JitScript* jitScript) {
+    MOZ_ASSERT(isWarmUpCount());
+    MOZ_ASSERT((uintptr_t(jitScript) & TagMask) == 0);
+    data_ = uintptr_t(jitScript) | JitScriptTag;
+  }
+  void clearJitScript() {
+    MOZ_ASSERT(isJitScript());
+    setWarmUpCount(0);
+  }
+};
+
+static_assert(sizeof(ScriptWarmUpData) == sizeof(uintptr_t),
+              "JIT code depends on ScriptWarmUpData being pointer-sized");
+
 } /* namespace js */
 
 namespace JS {
 
 // Define a GCManagedDeletePolicy to allow deleting type outside of normal
 // sweeping.
 template <>
 struct DeletePolicy<js::PrivateScriptData>
@@ -2344,31 +2411,21 @@ class JSScript : public js::BaseScript {
  private:
   // Shareable script data
   RefPtr<js::RuntimeScriptData> scriptData_ = {};
 
   // Unshared variable-length data
   js::PrivateScriptData* data_ = nullptr;
 
  private:
-  // JIT and type inference data for this script. May be purged on GC.
-  js::jit::JitScript* jitScript_ = nullptr;
+  js::ScriptWarmUpData warmUpData_ = {};
 
   /* Information used to re-lazify a lazily-parsed interpreted function. */
   js::LazyScript* lazyScript = nullptr;
 
-  // 32-bit fields.
-
-  // Number of times the script has been called or has had backedges taken.
-  // When running in ion, also increased for any inlined scripts. Reset if
-  // the script's JIT code is forcibly discarded.
-  mozilla::Atomic<uint32_t, mozilla::Relaxed,
-                  mozilla::recordreplay::Behavior::DontPreserve>
-      warmUpCount = {};
-
   //
   // End of fields.  Start methods.
   //
 
  private:
   template <js::XDRMode mode>
   friend js::XDRResult js::XDRScript(js::XDRState<mode>* xdr,
                                      js::HandleScope enclosingScope,
@@ -2644,18 +2701,18 @@ class JSScript : public js::BaseScript {
   }
 
   static constexpr size_t offsetOfScriptData() {
     return offsetof(JSScript, scriptData_);
   }
   static constexpr size_t offsetOfPrivateScriptData() {
     return offsetof(JSScript, data_);
   }
-  static constexpr size_t offsetOfJitScript() {
-    return offsetof(JSScript, jitScript_);
+  static constexpr size_t offsetOfWarmUpData() {
+    return offsetof(JSScript, warmUpData_);
   }
 
   void updateJitCodeRaw(JSRuntime* rt);
 
   // We don't relazify functions with a JitScript or JIT code, but some
   // callers (XDR, testing functions) want to know whether this script is
   // relazifiable ignoring (or after) discarding JIT code.
   bool isRelazifiableIgnoringJitCode() const {
@@ -2753,21 +2810,21 @@ class JSScript : public js::BaseScript {
    * If this script has a function associated to it, then it is not the
    * top-level of a file.
    */
   bool isTopLevel() { return code() && !functionNonDelazifying(); }
 
   /* Ensure the script has a JitScript. */
   inline bool ensureHasJitScript(JSContext* cx, js::jit::AutoKeepJitScripts&);
 
-  bool hasJitScript() const { return jitScript_ != nullptr; }
+  bool hasJitScript() const { return warmUpData_.isJitScript(); }
 
   js::jit::JitScript* jitScript() const {
     MOZ_ASSERT(hasJitScript());
-    return jitScript_;
+    return warmUpData_.toJitScript();
   }
   js::jit::JitScript* maybeJitScript() const {
     return hasJitScript() ? jitScript() : nullptr;
   }
 
   void maybeReleaseJitScript(JSFreeOp* fop);
   void releaseJitScript(JSFreeOp* fop);
   void releaseJitScriptOnFinalize(JSFreeOp* fop);
@@ -2845,30 +2902,19 @@ class JSScript : public js::BaseScript {
   bool createScriptData(JSContext* cx, uint32_t natoms);
   bool createImmutableScriptData(JSContext* cx, uint32_t codeLength,
                                  uint32_t noteLength, uint32_t numResumeOffsets,
                                  uint32_t numScopeNotes, uint32_t numTryNotes);
   bool shareScriptData(JSContext* cx);
   void freeScriptData();
 
  public:
-  uint32_t getWarmUpCount() const { return warmUpCount; }
-  uint32_t incWarmUpCounter(uint32_t amount = 1) {
-    return warmUpCount += amount;
-  }
-  uint32_t* addressOfWarmUpCounter() {
-    return reinterpret_cast<uint32_t*>(&warmUpCount);
-  }
-  static size_t offsetOfWarmUpCounter() {
-    return offsetof(JSScript, warmUpCount);
-  }
-  void resetWarmUpCounterForGC() {
-    incWarmUpResetCounter();
-    warmUpCount = 0;
-  }
+  inline uint32_t getWarmUpCount() const;
+  inline void incWarmUpCounter(uint32_t amount = 1);
+  inline void resetWarmUpCounterForGC();
 
   void resetWarmUpCounterToDelayIonCompilation();
 
   unsigned getWarmUpResetCount() const {
     constexpr uint32_t MASK = uint32_t(MutableFlags::WarmupResets_MASK);
     return mutableFlags_ & MASK;
   }
   void incWarmUpResetCounter() {