JIT: stateful local ref counts and weights (#19068)

Introduce a notion of state for local var ref counts and weighted ref counts.
Accesses and current state must agree.

State is invalid initially, enabled for an early period around bits of morph,
invalid again for a time, and then enabled normally once lvaMarkRefs is called.

Accesses normally specify RCS_NORMAL as the desired state, but in the accesses
of selected ref counts in morph, specify RCS_EARLY.

Revise how we decide if normal ref counting is active by changing
`lvaLocalVarRefCounted` into a method.

Update `gtIsLikelyRegVar` to not access ref counts when they're not in a valid
state.

Change weight APIs over to use `weight_t`.

9 files changed, 319 insertions, 140 deletions

diff --git a/src/jit/assertionprop.cpp b/src/jit/assertionprop.cpp
index 93ca437519..51b39470b1 100644
--- a/src/jit/assertionprop.cpp
+++ b/src/jit/assertionprop.cpp
@@ -2698,7 +2698,7 @@ GenTree* Compiler::optConstantAssertionProp(AssertionDsc* curAssertion,
         gtDispTree(newTree, nullptr, nullptr, true);
     }
 #endif
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         lvaTable[lclNum].decRefCnts(compCurBB->getBBWeight(this), this);
     }
@@ -2812,7 +2812,7 @@ GenTree* Compiler::optCopyAssertionProp(AssertionDsc* curAssertion,
     }
 
     // If global assertion prop, by now we should have ref counts, fix them.
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         lvaTable[lclNum].decRefCnts(compCurBB->getBBWeight(this), this);
         lvaTable[copyLclNum].incRefCnts(compCurBB->getBBWeight(this), this);
diff --git a/src/jit/compiler.cpp b/src/jit/compiler.cpp
index edfe431357..b66236b032 100644
--- a/src/jit/compiler.cpp
+++ b/src/jit/compiler.cpp
@@ -4695,7 +4695,7 @@ void Compiler::compCompile(void** methodCodePtr, ULONG* methodCodeSize, JitFlags
     // You can test the value of the following variable to see if
     // the local variable ref counts must be updated
     //
-    assert(lvaLocalVarRefCounted == true);
+    assert(lvaLocalVarRefCounted());
 
     if (!opts.MinOpts() && !opts.compDbgCode)
     {
diff --git a/src/jit/compiler.h b/src/jit/compiler.h
index 570d8d9c8e..7a914d42af 100644
--- a/src/jit/compiler.h
+++ b/src/jit/compiler.h
@@ -183,6 +183,13 @@ public:
 
 typedef JitExpandArray<LclSsaVarDsc> PerSsaArray;
 
+enum RefCountState
+{
+    RCS_INVALID, // not valid to get/set ref counts
+    RCS_EARLY,   // early counts for struct promotion and struct passing
+    RCS_NORMAL,  // normal ref counts (from lvaMarkRefs onward)
+};
+
 class LclVarDsc
 {
 public:
@@ -602,64 +609,19 @@ private:
                                // appearance count (computed during address-exposed analysis)
                                // that fgMakeOutgoingStructArgCopy consults during global morph
                                // to determine if eliding its copy is legal.
-    unsigned m_lvRefCntWtd;    // weighted reference count
-
-public:
-    unsigned short lvRefCnt() const
-    {
-        if (lvImplicitlyReferenced && (m_lvRefCnt == 0))
-        {
-            return 1;
-        }
-
-        return m_lvRefCnt;
-    }
-
-    void incLvRefCnt(unsigned short delta)
-    {
-        unsigned short oldRefCnt = m_lvRefCnt;
-        m_lvRefCnt += delta;
-        assert(m_lvRefCnt >= oldRefCnt);
-    }
-
-    void decLvRefCnt(unsigned short delta)
-    {
-        assert(m_lvRefCnt >= delta);
-        m_lvRefCnt -= delta;
-    }
-
-    void setLvRefCnt(unsigned short newValue)
-    {
-        m_lvRefCnt = newValue;
-    }
-
-    unsigned lvRefCntWtd() const
-    {
-        if (lvImplicitlyReferenced && (m_lvRefCntWtd == 0))
-        {
-            return BB_UNITY_WEIGHT;
-        }
 
-        return m_lvRefCntWtd;
-    }
-
-    void incLvRefCntWtd(unsigned delta)
-    {
-        unsigned oldRefCntWtd = m_lvRefCntWtd;
-        m_lvRefCntWtd += delta;
-        assert(m_lvRefCntWtd >= oldRefCntWtd);
-    }
+    BasicBlock::weight_t m_lvRefCntWtd; // weighted reference count
 
-    void decLvRefCntWtd(unsigned delta)
-    {
-        assert(m_lvRefCntWtd >= delta);
-        m_lvRefCntWtd -= delta;
-    }
+public:
+    unsigned short lvRefCnt(RefCountState state = RCS_NORMAL) const;
+    void incLvRefCnt(unsigned short delta, RefCountState state = RCS_NORMAL);
+    void decLvRefCnt(unsigned short delta, RefCountState state = RCS_NORMAL);
+    void setLvRefCnt(unsigned short newValue, RefCountState state = RCS_NORMAL);
 
-    void setLvRefCntWtd(unsigned newValue)
-    {
-        m_lvRefCntWtd = newValue;
-    }
+    BasicBlock::weight_t lvRefCntWtd(RefCountState state = RCS_NORMAL) const;
+    void incLvRefCntWtd(BasicBlock::weight_t delta, RefCountState state = RCS_NORMAL);
+    void decLvRefCntWtd(BasicBlock::weight_t delta, RefCountState state = RCS_NORMAL);
+    void setLvRefCntWtd(BasicBlock::weight_t newValue, RefCountState state = RCS_NORMAL);
 
     int      lvStkOffs;   // stack offset of home
     unsigned lvExactSize; // (exact) size of the type in bytes
@@ -750,9 +712,15 @@ public:
                !(lvIsParam || lvAddrExposed || lvIsStructField);
     }
 
-    void lvaResetSortAgainFlag(Compiler* pComp);
-    void decRefCnts(BasicBlock::weight_t weight, Compiler* pComp, bool propagate = true);
-    void incRefCnts(BasicBlock::weight_t weight, Compiler* pComp, bool propagate = true);
+    void lvaResetSortAgainFlag(Compiler* pComp, RefCountState = RCS_NORMAL);
+    void decRefCnts(BasicBlock::weight_t weight,
+                    Compiler*            pComp,
+                    RefCountState        state     = RCS_NORMAL,
+                    bool                 propagate = true);
+    void incRefCnts(BasicBlock::weight_t weight,
+                    Compiler*            pComp,
+                    RefCountState        state     = RCS_NORMAL,
+                    bool                 propagate = true);
     void setPrefReg(regNumber regNum, Compiler* pComp);
     void addPrefReg(regMaskTP regMask, Compiler* pComp);
     bool IsFloatRegType() const
@@ -2567,11 +2535,16 @@ public:
     };
 
 public:
-    bool     lvaRefCountingStarted; // Set to true when we have started counting the local vars
-    bool     lvaLocalVarRefCounted; // Set to true after we have called lvaMarkLocalVars()
-    bool     lvaSortAgain;          // true: We need to sort the lvaTable
-    bool     lvaTrackedFixed;       // true: We cannot add new 'tracked' variable
-    unsigned lvaCount;              // total number of locals
+    RefCountState lvaRefCountState; // Current local ref count state
+
+    bool lvaLocalVarRefCounted() const
+    {
+        return lvaRefCountState == RCS_NORMAL;
+    }
+
+    bool     lvaSortAgain;    // true: We need to sort the lvaTable
+    bool     lvaTrackedFixed; // true: We cannot add new 'tracked' variable
+    unsigned lvaCount;        // total number of locals
 
     unsigned   lvaRefCount; // total number of references to locals
     LclVarDsc* lvaTable;    // variable descriptor table
diff --git a/src/jit/compiler.hpp b/src/jit/compiler.hpp
index da00abcd03..0385e1035c 100644
--- a/src/jit/compiler.hpp
+++ b/src/jit/compiler.hpp
@@ -1810,7 +1810,7 @@ inline unsigned Compiler::lvaGrabTempWithImplicitUse(bool shortLifetime DEBUGARG
  *   and zero lvRefCntWtd when lvRefCnt is zero
  */
 
-inline void LclVarDsc::lvaResetSortAgainFlag(Compiler* comp)
+inline void LclVarDsc::lvaResetSortAgainFlag(Compiler* comp, RefCountState state)
 {
     if (!comp->lvaTrackedFixed)
     {
@@ -1818,9 +1818,9 @@ inline void LclVarDsc::lvaResetSortAgainFlag(Compiler* comp)
         comp->lvaSortAgain = true;
     }
     /* Set weighted ref count to zero if  ref count is zero */
-    if (lvRefCnt() == 0)
+    if (lvRefCnt(state) == 0)
     {
-        setLvRefCntWtd(0);
+        setLvRefCntWtd(0, state);
     }
 }
 
@@ -1829,7 +1829,7 @@ inline void LclVarDsc::lvaResetSortAgainFlag(Compiler* comp)
  *  Decrement the ref counts for a local variable
  */
 
-inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, bool propagate)
+inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, RefCountState state, bool propagate)
 {
     /* Decrement lvRefCnt and lvRefCntWtd */
     Compiler::lvaPromotionType promotionType = DUMMY_INIT(Compiler::PROMOTION_TYPE_NONE);
@@ -1843,17 +1843,17 @@ inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
     //
     if (lvType != TYP_STRUCT || promotionType != Compiler::PROMOTION_TYPE_INDEPENDENT)
     {
-        assert(lvRefCnt()); // Can't decrement below zero
+        assert(lvRefCnt(state)); // Can't decrement below zero
 
         // TODO: Well, the assert above could be bogus.
         // If lvRefCnt has overflowed before, then might drop to 0.
         // Therefore we do need the following check to keep lvRefCnt from underflow:
-        if (lvRefCnt() > 0)
+        if (lvRefCnt(state) > 0)
         {
             //
             // Decrement lvRefCnt
             //
-            decLvRefCnt(1);
+            decLvRefCnt(1, state);
 
             //
             // Decrement lvRefCntWtd
@@ -1865,13 +1865,13 @@ inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
                     weight *= 2;
                 }
 
-                if (lvRefCntWtd() <= weight)
+                if (lvRefCntWtd(state) <= weight)
                 { // Can't go below zero
-                    setLvRefCntWtd(0);
+                    setLvRefCntWtd(0, state);
                 }
                 else
                 {
-                    decLvRefCntWtd(weight);
+                    decLvRefCntWtd(weight, state);
                 }
             }
         }
@@ -1885,7 +1885,7 @@ inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
         {
             for (unsigned i = lvFieldLclStart; i < lvFieldLclStart + lvFieldCnt; ++i)
             {
-                comp->lvaTable[i].decRefCnts(comp->lvaMarkRefsWeight, comp, false); // Don't propagate
+                comp->lvaTable[i].decRefCnts(comp->lvaMarkRefsWeight, comp, state, false); // Don't propagate
             }
         }
     }
@@ -1898,19 +1898,19 @@ inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
         assert(!parentvarDsc->lvRegStruct);
         if (promotionType == Compiler::PROMOTION_TYPE_DEPENDENT)
         {
-            parentvarDsc->decRefCnts(comp->lvaMarkRefsWeight, comp, false); // Don't propagate
+            parentvarDsc->decRefCnts(comp->lvaMarkRefsWeight, comp, state, false); // Don't propagate
         }
     }
 
-    lvaResetSortAgainFlag(comp);
+    lvaResetSortAgainFlag(comp, state);
 
 #ifdef DEBUG
     if (comp->verbose)
     {
         unsigned varNum = (unsigned)(this - comp->lvaTable);
         assert(&comp->lvaTable[varNum] == this);
-        printf("New refCnts for V%02u: refCnt = %2u, refCntWtd = %s\n", varNum, lvRefCnt(),
-               refCntWtd2str(lvRefCntWtd()));
+        printf("New refCnts for V%02u: refCnt = %2u, refCntWtd = %s\n", varNum, lvRefCnt(state),
+               refCntWtd2str(lvRefCntWtd(state)));
     }
 #endif
 }
@@ -1920,7 +1920,7 @@ inline void LclVarDsc::decRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
  *  Increment the ref counts for a local variable
  */
 
-inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, bool propagate)
+inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, RefCountState state, bool propagate)
 {
     Compiler::lvaPromotionType promotionType = DUMMY_INIT(Compiler::PROMOTION_TYPE_NONE);
     if (varTypeIsStruct(lvType))
@@ -1936,10 +1936,10 @@ inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
         //
         // Increment lvRefCnt
         //
-        int newRefCnt = lvRefCnt() + 1;
+        int newRefCnt = lvRefCnt(state) + 1;
         if (newRefCnt == (unsigned short)newRefCnt) // lvRefCnt is an "unsigned short". Don't overflow it.
         {
-            setLvRefCnt((unsigned short)newRefCnt);
+            setLvRefCnt((unsigned short)newRefCnt, state);
         }
 
         // This fires when an uninitialize value for 'weight' is used (see lvaMarkRefsWeight)
@@ -1956,14 +1956,14 @@ inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
                 weight *= 2;
             }
 
-            unsigned newWeight = lvRefCntWtd() + weight;
-            if (newWeight >= lvRefCntWtd())
+            unsigned newWeight = lvRefCntWtd(state) + weight;
+            if (newWeight >= lvRefCntWtd(state))
             { // lvRefCntWtd is an "unsigned".  Don't overflow it
-                setLvRefCntWtd(newWeight);
+                setLvRefCntWtd(newWeight, state);
             }
             else
             { // On overflow we assign ULONG_MAX
-                setLvRefCntWtd(ULONG_MAX);
+                setLvRefCntWtd(ULONG_MAX, state);
             }
         }
     }
@@ -1976,7 +1976,7 @@ inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
         {
             for (unsigned i = lvFieldLclStart; i < lvFieldLclStart + lvFieldCnt; ++i)
             {
-                comp->lvaTable[i].incRefCnts(comp->lvaMarkRefsWeight, comp, false); // Don't propagate
+                comp->lvaTable[i].incRefCnts(comp->lvaMarkRefsWeight, comp, state, false); // Don't propagate
             }
         }
     }
@@ -1989,19 +1989,19 @@ inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, b
         assert(!parentvarDsc->lvRegStruct);
         if (promotionType == Compiler::PROMOTION_TYPE_DEPENDENT)
         {
-            parentvarDsc->incRefCnts(comp->lvaMarkRefsWeight, comp, false); // Don't propagate
+            parentvarDsc->incRefCnts(comp->lvaMarkRefsWeight, comp, state, false); // Don't propagate
         }
     }
 
-    lvaResetSortAgainFlag(comp);
+    lvaResetSortAgainFlag(comp, state);
 
 #ifdef DEBUG
     if (comp->verbose)
     {
         unsigned varNum = (unsigned)(this - comp->lvaTable);
         assert(&comp->lvaTable[varNum] == this);
-        printf("New refCnts for V%02u: refCnt = %2u, refCntWtd = %s\n", varNum, lvRefCnt(),
-               refCntWtd2str(lvRefCntWtd()));
+        printf("New refCnts for V%02u: refCnt = %2u, refCntWtd = %s\n", varNum, lvRefCnt(state),
+               refCntWtd2str(lvRefCntWtd(state)));
     }
 #endif
 }
@@ -4946,6 +4946,202 @@ inline void DEBUG_DESTROY_NODE(GenTree* tree)
 #endif
 }
 
+//------------------------------------------------------------------------------
+// lvRefCnt: access reference count for this local var
+//
+// Arguments:
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Return Value:
+//    Ref count for the local.
+
+inline unsigned short LclVarDsc::lvRefCnt(RefCountState state) const
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    if (lvImplicitlyReferenced && (m_lvRefCnt == 0))
+    {
+        return 1;
+    }
+
+    return m_lvRefCnt;
+}
+
+//------------------------------------------------------------------------------
+// incLvRefCnt: increment reference count for this local var
+//
+// Arguments:
+//    delta: the amount of the increment
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    It is currently the caller's responsibilty to ensure this increment
+//    will not cause overflow.
+
+inline void LclVarDsc::incLvRefCnt(unsigned short delta, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    unsigned short oldRefCnt = m_lvRefCnt;
+    m_lvRefCnt += delta;
+    assert(m_lvRefCnt >= oldRefCnt);
+}
+
+//------------------------------------------------------------------------------
+// decLvRefCnt: decrement reference count for this local var
+//
+// Arguments:
+//    delta: the amount of the decrement
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    It is currently the caller's responsibilty to ensure this decrement
+//    will not cause underflow.
+
+inline void LclVarDsc::decLvRefCnt(unsigned short delta, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    assert(m_lvRefCnt >= delta);
+    m_lvRefCnt -= delta;
+}
+
+//------------------------------------------------------------------------------
+// setLvRefCnt: set the reference count for this local var
+//
+// Arguments:
+//    newValue: the desired new reference count
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    Generally after calling v->setLvRefCnt(Y), v->lvRefCnt() == Y.
+//    However this may not be true when v->lvImplicitlyReferenced == 1.
+
+inline void LclVarDsc::setLvRefCnt(unsigned short newValue, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    m_lvRefCnt = newValue;
+}
+
+//------------------------------------------------------------------------------
+// lvRefCntWtd: access wighted reference count for this local var
+//
+// Arguments:
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Return Value:
+//    Weighted ref count for the local.
+
+inline BasicBlock::weight_t LclVarDsc::lvRefCntWtd(RefCountState state) const
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    if (lvImplicitlyReferenced && (m_lvRefCntWtd == 0))
+    {
+        return BB_UNITY_WEIGHT;
+    }
+
+    return m_lvRefCntWtd;
+}
+
+//------------------------------------------------------------------------------
+// incLvRefCntWtd: increment weighted reference count for this local var
+//
+// Arguments:
+//    delta: the amount of the increment
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    It is currently the caller's responsibilty to ensure this increment
+//    will not cause overflow.
+
+inline void LclVarDsc::incLvRefCntWtd(BasicBlock::weight_t delta, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    BasicBlock::weight_t oldRefCntWtd = m_lvRefCntWtd;
+    m_lvRefCntWtd += delta;
+    assert(m_lvRefCntWtd >= oldRefCntWtd);
+}
+
+//------------------------------------------------------------------------------
+// decLvRefCntWtd: decrement weighted reference count for this local var
+//
+// Arguments:
+//    delta: the amount of the decrement
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    It is currently the caller's responsibilty to ensure this decrement
+//    will not cause underflow.
+
+inline void LclVarDsc::decLvRefCntWtd(BasicBlock::weight_t delta, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    assert(m_lvRefCntWtd >= delta);
+    m_lvRefCntWtd -= delta;
+}
+
+//------------------------------------------------------------------------------
+// setLvRefCntWtd: set the weighted reference count for this local var
+//
+// Arguments:
+//    newValue: the desired new weighted reference count
+//    state: the requestor's expected ref count state; defaults to RCS_NORMAL
+//
+// Notes:
+//    Generally after calling v->setLvRefCntWtd(Y), v->lvRefCntWtd() == Y.
+//    However this may not be true when v->lvImplicitlyReferenced == 1.
+
+inline void LclVarDsc::setLvRefCntWtd(BasicBlock::weight_t newValue, RefCountState state)
+{
+
+#if defined(DEBUG)
+    assert(state != RCS_INVALID);
+    Compiler* compiler = JitTls::GetCompiler();
+    assert(compiler->lvaRefCountState == state);
+#endif
+
+    m_lvRefCntWtd = newValue;
+}
+
 /*****************************************************************************/
 #endif //_COMPILER_HPP_
 /*****************************************************************************/
diff --git a/src/jit/flowgraph.cpp b/src/jit/flowgraph.cpp
index 46eefa4167..ef2a4b3a37 100644
--- a/src/jit/flowgraph.cpp
+++ b/src/jit/flowgraph.cpp
@@ -9686,7 +9686,7 @@ void Compiler::fgUpdateRefCntForClone(BasicBlock* addedToBlock, GenTree* clonedT
 {
     assert(clonedTree->gtOper != GT_STMT);
 
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         compCurBB = addedToBlock;
         IncLclVarRefCountsVisitor::WalkTree(this, clonedTree);
@@ -9698,7 +9698,7 @@ void Compiler::fgUpdateRefCntForClone(BasicBlock* addedToBlock, GenTree* clonedT
 
 void Compiler::fgUpdateRefCntForExtract(GenTree* wholeTree, GenTree* keptTree)
 {
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         /*  Update the refCnts of removed lcl vars - The problem is that
          *  we have to consider back the side effects trees so we first
diff --git a/src/jit/gentree.cpp b/src/jit/gentree.cpp
index 1f60fe6169..946b52aec3 100644
--- a/src/jit/gentree.cpp
+++ b/src/jit/gentree.cpp
@@ -2983,6 +2983,15 @@ bool Compiler::gtIsLikelyRegVar(GenTree* tree)
         return false;
     }
 
+    // Be pessimistic if ref counts are not yet set up.
+    //
+    // Perhaps we should be optimistic though.
+    // See notes in GitHub issue 18969.
+    if (!lvaLocalVarRefCounted())
+    {
+        return false;
+    }
+
     if (varDsc->lvRefCntWtd() < (BB_UNITY_WEIGHT * 3))
     {
         return false;
@@ -12219,7 +12228,7 @@ GenTree* Compiler::gtFoldExprCompare(GenTree* tree)
         cons->gtNext = tree->gtNext;
         cons->gtPrev = tree->gtPrev;
     }
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         lvaRecursiveDecRefCounts(tree);
     }
@@ -12660,7 +12669,7 @@ GenTree* Compiler::gtFoldExprSpecial(GenTree* tree)
                 /* Multiply by zero - return the 'zero' node, but not if side effects */
                 if (!(op->gtFlags & GTF_SIDE_EFFECT))
                 {
-                    if (lvaLocalVarRefCounted)
+                    if (lvaLocalVarRefCounted())
                     {
                         lvaRecursiveDecRefCounts(op);
                     }
@@ -12692,7 +12701,7 @@ GenTree* Compiler::gtFoldExprSpecial(GenTree* tree)
 
                 if (!(op->gtFlags & GTF_SIDE_EFFECT))
                 {
-                    if (lvaLocalVarRefCounted)
+                    if (lvaLocalVarRefCounted())
                     {
                         lvaRecursiveDecRefCounts(op);
                     }
@@ -12732,7 +12741,7 @@ GenTree* Compiler::gtFoldExprSpecial(GenTree* tree)
 
                 if (!(op->gtFlags & GTF_SIDE_EFFECT))
                 {
-                    if (lvaLocalVarRefCounted)
+                    if (lvaLocalVarRefCounted())
                     {
                         lvaRecursiveDecRefCounts(op);
                     }
@@ -12755,7 +12764,7 @@ GenTree* Compiler::gtFoldExprSpecial(GenTree* tree)
                 }
                 else if (!(op->gtFlags & GTF_SIDE_EFFECT))
                 {
-                    if (lvaLocalVarRefCounted)
+                    if (lvaLocalVarRefCounted())
                     {
                         lvaRecursiveDecRefCounts(op);
                     }
@@ -12783,7 +12792,7 @@ GenTree* Compiler::gtFoldExprSpecial(GenTree* tree)
                 op         = op2->AsColon()->ElseNode();
                 opToDelete = op2->AsColon()->ThenNode();
             }
-            if (lvaLocalVarRefCounted)
+            if (lvaLocalVarRefCounted())
             {
                 lvaRecursiveDecRefCounts(opToDelete);
             }
diff --git a/src/jit/lclvars.cpp b/src/jit/lclvars.cpp
index 00110b0e6a..ea1976008a 100644
--- a/src/jit/lclvars.cpp
+++ b/src/jit/lclvars.cpp
@@ -35,8 +35,7 @@ unsigned Compiler::s_lvaDoubleAlignedProcsCount = 0;
 void Compiler::lvaInit()
 {
     /* We haven't allocated stack variables yet */
-    lvaRefCountingStarted = false;
-    lvaLocalVarRefCounted = false;
+    lvaRefCountState = RCS_INVALID;
 
     lvaGenericsContextUseCount = 0;
 
@@ -2025,8 +2024,9 @@ void Compiler::lvaPromoteStructVar(unsigned lclNum, lvaStructPromotionInfo* Stru
             }
 #endif // FEATURE_MULTIREG_ARGS && defined(FEATURE_SIMD)
 
-            lvaMarkRefsWeight = BB_UNITY_WEIGHT;            // incRefCnts can use this compiler global variable
-            fieldVarDsc->incRefCnts(BB_UNITY_WEIGHT, this); // increment the ref count for prolog initialization
+            lvaMarkRefsWeight = BB_UNITY_WEIGHT; // incRefCnts can use this compiler global variable
+            fieldVarDsc->incRefCnts(BB_UNITY_WEIGHT, this,
+                                    RCS_EARLY); // increment the ref count for prolog initialization
         }
 #endif
 
@@ -2803,7 +2803,7 @@ BasicBlock::weight_t BasicBlock::getBBWeight(Compiler* comp)
 // Decrement the ref counts for all locals contained in the tree and its children.
 void Compiler::lvaRecursiveDecRefCounts(GenTree* tree)
 {
-    assert(lvaLocalVarRefCounted);
+    assert(lvaLocalVarRefCounted());
 
     // We could just use the recursive walker for all cases but that is a
     // fairly heavyweight thing to spin up when we're usually just handling a leaf.
@@ -2856,7 +2856,7 @@ void Compiler::lvaDecRefCnts(BasicBlock* block, GenTree* tree)
     unsigned   lclNum;
     LclVarDsc* varDsc;
 
-    noway_assert(lvaRefCountingStarted || lvaLocalVarRefCounted);
+    noway_assert(lvaLocalVarRefCounted());
 
     if ((tree->gtOper == GT_CALL) && (tree->gtFlags & GTF_CALL_UNMANAGED))
     {
@@ -2898,7 +2898,7 @@ void Compiler::lvaDecRefCnts(BasicBlock* block, GenTree* tree)
 // Increment the ref counts for all locals contained in the tree and its children.
 void Compiler::lvaRecursiveIncRefCounts(GenTree* tree)
 {
-    assert(lvaLocalVarRefCounted);
+    assert(lvaLocalVarRefCounted());
 
     // We could just use the recursive walker for all cases but that is a
     // fairly heavyweight thing to spin up when we're usually just handling a leaf.
@@ -2942,7 +2942,7 @@ void Compiler::lvaIncRefCnts(GenTree* tree)
     unsigned   lclNum;
     LclVarDsc* varDsc;
 
-    noway_assert(lvaRefCountingStarted || lvaLocalVarRefCounted);
+    noway_assert(lvaLocalVarRefCounted());
 
     if ((tree->gtOper == GT_CALL) && (tree->gtFlags & GTF_CALL_UNMANAGED))
     {
@@ -4107,9 +4107,10 @@ void Compiler::lvaMarkLocalVars()
         }
     }
 
-    /* Mark all local variable references */
+    // Ref counting is now enabled normally.
+    lvaRefCountState = RCS_NORMAL;
 
-    lvaRefCountingStarted = true;
+    /* Mark all local variable references */
     for (block = fgFirstBB; block; block = block->bbNext)
     {
         lvaMarkLocalVars(block);
@@ -4158,9 +4159,6 @@ void Compiler::lvaMarkLocalVars()
         lvaTable[info.compTypeCtxtArg].lvImplicitlyReferenced = 1;
     }
 
-    lvaLocalVarRefCounted = true;
-    lvaRefCountingStarted = false;
-
     lvaSortByRefCount();
 }
 
diff --git a/src/jit/morph.cpp b/src/jit/morph.cpp
index 5f4cfb5d99..44034e7336 100644
--- a/src/jit/morph.cpp
+++ b/src/jit/morph.cpp
@@ -2627,7 +2627,7 @@ GenTree* Compiler::fgMakeMultiUse(GenTree** pOp)
     if (tree->IsLocal())
     {
         auto result = gtClone(tree);
-        if (lvaLocalVarRefCounted)
+        if (lvaLocalVarRefCounted())
         {
             lvaTable[tree->gtLclVarCommon.gtLclNum].incRefCnts(compCurBB->getBBWeight(this), this);
         }
@@ -2640,7 +2640,7 @@ GenTree* Compiler::fgMakeMultiUse(GenTree** pOp)
         // At this point, *pOp is GT_COMMA(GT_ASG(V01, *pOp), V01) and result = V01
         // Therefore, the ref count has to be incremented 3 times for *pOp and result, if result will
         // be added by the caller.
-        if (lvaLocalVarRefCounted)
+        if (lvaLocalVarRefCounted())
         {
             lvaTable[result->gtLclVarCommon.gtLclNum].incRefCnts(compCurBB->getBBWeight(this), this);
             lvaTable[result->gtLclVarCommon.gtLclNum].incRefCnts(compCurBB->getBBWeight(this), this);
@@ -5226,9 +5226,9 @@ void Compiler::fgMakeOutgoingStructArgCopy(GenTreeCall*         call,
             // on the caller's frame. If an argument lives on the caller caller's frame, it may get
             // overwritten if that frame is reused for the tail call. Therefore, we should always copy
             // struct parameters if they are passed as arguments to a tail call.
-            if (!call->IsTailCallViaHelper() && (varDsc->lvRefCnt() == 1) && !fgMightHaveLoop())
+            if (!call->IsTailCallViaHelper() && (varDsc->lvRefCnt(RCS_EARLY) == 1) && !fgMightHaveLoop())
             {
-                varDsc->setLvRefCnt(0);
+                varDsc->setLvRefCnt(0, RCS_EARLY);
                 args->gtOp.gtOp1 = lcl;
                 fp->node         = lcl;
 
@@ -5904,7 +5904,7 @@ GenTree* Compiler::fgMorphArrayIndex(GenTree* tree)
         bndsChk = arrBndsChk;
 
         // Make sure to increment ref-counts if already ref-counted.
-        if (lvaLocalVarRefCounted)
+        if (lvaLocalVarRefCounted())
         {
             lvaRecursiveIncRefCounts(index);
             lvaRecursiveIncRefCounts(arrRef);
@@ -13745,7 +13745,7 @@ DONE_MORPHING_CHILDREN:
                 else
                 {
                     /* The left operand is worthless, throw it away */
-                    if (lvaLocalVarRefCounted)
+                    if (lvaLocalVarRefCounted())
                     {
                         lvaRecursiveDecRefCounts(op1);
                     }
@@ -14317,7 +14317,7 @@ GenTree* Compiler::fgMorphModToSubMulDiv(GenTreeOp* tree)
     {
         numerator = fgMakeMultiUse(&tree->gtOp1);
     }
-    else if (lvaLocalVarRefCounted && numerator->OperIsLocal())
+    else if (lvaLocalVarRefCounted() && numerator->OperIsLocal())
     {
         // Morphing introduces new lclVar references. Increase ref counts
         lvaIncRefCnts(numerator);
@@ -14327,7 +14327,7 @@ GenTree* Compiler::fgMorphModToSubMulDiv(GenTreeOp* tree)
     {
         denominator = fgMakeMultiUse(&tree->gtOp2);
     }
-    else if (lvaLocalVarRefCounted && denominator->OperIsLocal())
+    else if (lvaLocalVarRefCounted() && denominator->OperIsLocal())
     {
         // Morphing introduces new lclVar references. Increase ref counts
         lvaIncRefCnts(denominator);
@@ -15687,7 +15687,7 @@ bool Compiler::fgMorphBlockStmt(BasicBlock* block, GenTreeStmt* stmt DEBUGARG(co
 
     stmt->gtStmtExpr = morph;
 
-    if (lvaLocalVarRefCounted)
+    if (lvaLocalVarRefCounted())
     {
         // fgMorphTree may have introduced new lclVar references. Bump the ref counts if requested.
         lvaRecursiveIncRefCounts(stmt->gtStmtExpr);
@@ -17093,6 +17093,8 @@ void Compiler::fgMorph()
     fgUpdateFinallyTargetFlags();
 
     /* For x64 and ARM64 we need to mark irregular parameters */
+
+    lvaRefCountState = RCS_EARLY;
     fgMarkImplicitByRefArgs();
 
     /* Promote struct locals if necessary */
@@ -17121,6 +17123,7 @@ void Compiler::fgMorph()
 
     /* Fix any LclVar annotations on discarded struct promotion temps for implicit by-ref args */
     fgMarkDemotedImplicitByRefArgs();
+    lvaRefCountState = RCS_INVALID;
 
     EndPhase(PHASE_MORPH_GLOBAL);
 
@@ -17346,8 +17349,8 @@ Compiler::fgWalkResult Compiler::fgMorphStructField(GenTree* tree, fgWalkData* f
                     // chance, so have to check now.
                     JITDUMP(
                         "Incrementing ref count from %d to %d for V%02d in fgMorphStructField for promoted struct\n",
-                        varDsc->lvRefCnt(), varDsc->lvRefCnt() + 1, lclNum);
-                    varDsc->incLvRefCnt(1);
+                        varDsc->lvRefCnt(RCS_EARLY), varDsc->lvRefCnt(RCS_EARLY) + 1, lclNum);
+                    varDsc->incLvRefCnt(1, RCS_EARLY);
                 }
 
                 tree->SetOper(GT_LCL_VAR);
@@ -17437,8 +17440,8 @@ Compiler::fgWalkResult Compiler::fgMorphStructField(GenTree* tree, fgWalkData* f
                     // lclVars, but here we're about to return SKIP_SUBTREES and rob it of the
                     // chance, so have to check now.
                     JITDUMP("Incrementing ref count from %d to %d for V%02d in fgMorphStructField for normed struct\n",
-                            varDsc->lvRefCnt(), varDsc->lvRefCnt() + 1, lclNum);
-                    varDsc->incLvRefCnt(1);
+                            varDsc->lvRefCnt(RCS_EARLY), varDsc->lvRefCnt(RCS_EARLY) + 1, lclNum);
+                    varDsc->incLvRefCnt(1, RCS_EARLY);
                 }
 
                 tree->ChangeOper(GT_LCL_VAR);
@@ -17594,7 +17597,7 @@ void Compiler::fgMarkImplicitByRefArgs()
                 // appearance of implicit-by-ref param so that call arg morphing can do an
                 // optimization for single-use implicit-by-ref params whose single use is as
                 // an outgoing call argument.
-                varDsc->setLvRefCnt(0);
+                varDsc->setLvRefCnt(0, RCS_EARLY);
             }
         }
     }
@@ -17681,7 +17684,7 @@ void Compiler::fgRetypeImplicitByRefArgs()
                 // parameter if it weren't promoted at all (otherwise the initialization
                 // of the new temp would just be a needless memcpy at method entry).
                 bool undoPromotion = (lvaGetPromotionType(newVarDsc) == PROMOTION_TYPE_DEPENDENT) ||
-                                     (varDsc->lvRefCnt() <= varDsc->lvFieldCnt);
+                                     (varDsc->lvRefCnt(RCS_EARLY) <= varDsc->lvFieldCnt);
 
                 if (!undoPromotion)
                 {
@@ -17719,7 +17722,7 @@ void Compiler::fgRetypeImplicitByRefArgs()
                         // to the implicit byref parameter when morphing calls that pass the implicit byref
                         // out as an outgoing argument value, but that doesn't pertain to this field local
                         // which is now a field of a non-arg local.
-                        fieldVarDsc->setLvRefCnt(0);
+                        fieldVarDsc->setLvRefCnt(0, RCS_EARLY);
                     }
 
                     fieldVarDsc->lvIsParam = false;
@@ -17836,12 +17839,12 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
                 // call morphing could identify single-use implicit byrefs; we're done with
                 // that, and want it to be in its default state of zero when we go to set
                 // real ref counts for all variables.
-                varDsc->setLvRefCnt(0);
+                varDsc->setLvRefCnt(0, RCS_EARLY);
 
                 // The temp struct is now unused; set flags appropriately so that we
                 // won't allocate space for it on the stack.
                 LclVarDsc* structVarDsc = &lvaTable[structLclNum];
-                structVarDsc->setLvRefCnt(0);
+                structVarDsc->setLvRefCnt(0, RCS_EARLY);
                 structVarDsc->lvAddrExposed = false;
 #ifdef DEBUG
                 structVarDsc->lvUnusedStruct = true;
@@ -17860,7 +17863,7 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
 
                     // The field local is now unused; set flags appropriately so that
                     // we won't allocate stack space for it.
-                    fieldVarDsc->setLvRefCnt(0);
+                    fieldVarDsc->setLvRefCnt(0, RCS_EARLY);
                     fieldVarDsc->lvAddrExposed = false;
                 }
             }
@@ -18255,10 +18258,10 @@ Compiler::fgWalkResult Compiler::fgMarkAddrTakenLocalsPreCB(GenTree** pTree, fgW
                 // checks the ref counts for implicit byref params when deciding if it's legal
                 // to elide certain copies of them.
                 LclVarDsc* varDsc = &comp->lvaTable[lclNum];
-                JITDUMP("Incrementing ref count from %d to %d for V%02d in fgMorphStructField\n", varDsc->lvRefCnt(),
-                        varDsc->lvRefCnt() + 1, lclNum);
+                JITDUMP("Incrementing ref count from %d to %d for V%02d in fgMorphStructField\n",
+                        varDsc->lvRefCnt(RCS_EARLY), varDsc->lvRefCnt(RCS_EARLY) + 1, lclNum);
 
-                varDsc->incLvRefCnt(1);
+                varDsc->incLvRefCnt(1, RCS_EARLY);
             }
             // This recognizes certain forms, and does all the work.  In that case, returns WALK_SKIP_SUBTREES,
             // else WALK_CONTINUE.  We do the same here.
@@ -18293,10 +18296,10 @@ Compiler::fgWalkResult Compiler::fgMarkAddrTakenLocalsPreCB(GenTree** pTree, fgW
                 // byref (here during address-exposed analysis); fgMakeOutgoingStructArgCopy
                 // checks the ref counts for implicit byref params when deciding if it's legal
                 // to elide certain copies of them.
-                JITDUMP("Incrementing ref count from %d to %d for V%02d in fgMorphStructField\n", varDsc->lvRefCnt(),
-                        varDsc->lvRefCnt() + 1, lclNum);
+                JITDUMP("Incrementing ref count from %d to %d for V%02d in fgMorphStructField\n",
+                        varDsc->lvRefCnt(RCS_EARLY), varDsc->lvRefCnt(RCS_EARLY) + 1, lclNum);
 
-                varDsc->incLvRefCnt(1);
+                varDsc->incLvRefCnt(1, RCS_EARLY);
             }
 
             if (axc == AXC_Addr || axc == AXC_AddrWide)
diff --git a/src/jit/optimizer.cpp b/src/jit/optimizer.cpp
index 6158a82cd9..a77a8696ac 100644
--- a/src/jit/optimizer.cpp
+++ b/src/jit/optimizer.cpp
@@ -7941,7 +7941,7 @@ Compiler::fgWalkResult Compiler::optRemoveTreeVisitor(GenTree** pTree, fgWalkDat
 
     // Look for any local variable references
 
-    if (tree->gtOper == GT_LCL_VAR && comp->lvaLocalVarRefCounted)
+    if (tree->gtOper == GT_LCL_VAR && comp->lvaLocalVarRefCounted())
     {
         unsigned   lclNum;
         LclVarDsc* varDsc;