[WIP] Enregister small structs.

src/coreclr/jit/codegencommon.cpp

@@ -4381,7 +4381,7 @@ void CodeGen::genFnPrologCalleeRegArgs(regNumber xtraReg, bool* pXtraRegClobbere
 
             /* Register argument - hopefully it stays in the same register */
             regNumber destRegNum  = REG_NA;
-            var_types destMemType = varDsc->TypeGet();
+            var_types destMemType = varDsc->GetRegisterType();
 
             if (regArgTab[argNum].slot == 1)
             {
@@ -4657,8 +4657,6 @@ void CodeGen::genEnregisterIncomingStackArgs()
             continue;
         }
 
-        var_types type = genActualType(varDsc->TypeGet());
-
         /* Is the variable dead on entry */
 
         if (!VarSetOps::IsMember(compiler, compiler->fgFirstBB->bbLiveIn, varDsc->lvVarIndex))
@@ -4673,7 +4671,9 @@ void CodeGen::genEnregisterIncomingStackArgs()
         regNumber regNum = varDsc->GetArgInitReg();
         assert(regNum != REG_STK);
 
-        GetEmitter()->emitIns_R_S(ins_Load(type), emitTypeSize(type), regNum, varNum, 0);
+        var_types regType = genActualType(varDsc->GetRegisterType());
+
+        GetEmitter()->emitIns_R_S(ins_Load(regType), emitTypeSize(regType), regNum, varNum, 0);
         regSet.verifyRegUsed(regNum);
 #ifdef USING_SCOPE_INFO
         psiMoveToReg(varNum);
@@ -4794,14 +4794,41 @@ void CodeGen::genCheckUseBlockInit()
                 }
             }
 
-            /* With compInitMem, all untracked vars will have to be init'ed */
-            /* VSW 102460 - Do not force initialization of compiler generated temps,
-                unless they are untracked GC type or structs that contain GC pointers */
-            CLANG_FORMAT_COMMENT_ANCHOR;
+            bool mustInitThisVar = false;
+
+            const bool isTemp      = varDsc->lvIsTemp;
+            const bool hasGCPtr    = varDsc->HasGCPtr();
+            const bool isTracked   = varDsc->lvTracked;
+            const bool isStruct    = varTypeIsStruct(varDsc);
+            const bool compInitMem = compiler->info.compInitMem;
 
-            if ((!varDsc->lvTracked || (varDsc->lvType == TYP_STRUCT)) && varDsc->lvOnFrame)
+            if (hasGCPtr && !isTracked)
             {
+                JITDUMP("must init V%02u because it has a GC ref\n", varNum);
+                mustInitThisVar = true;
+            }
+            else if (hasGCPtr && isStruct)
+            {
+                // TODO-1stClassStructs: support precise liveness reporting for such structs.
+                JITDUMP("must init a tracked V%02u because it a struct with a GC ref\n", varNum);
+                mustInitThisVar = true;
+            }
+            else
+            {
+                // We are done with tracked or GC vars, now look at untracked vars without GC refs.
+                if (!isTracked)
+                {
+                    assert(!hasGCPtr);
+                    if (compInitMem && !isTemp)
+                    {
+                        JITDUMP("must init V%02u because compInitMem is set and it is not a temp\n", varNum);
+                        mustInitThisVar = true;
+                    }
+                }
+            }
 
+            if (mustInitThisVar)
+            {
                 varDsc->lvMustInit = true;
 
                 if (!counted)

src/coreclr/jit/codegenlinear.cpp

@@ -864,14 +864,14 @@ void CodeGen::genSpillVar(GenTree* tree)
         // therefore be store-normalized (rather than load-normalized). In fact, not performing store normalization
         // can lead to problems on architectures where a lclVar may be allocated to a register that is not
         // addressable at the granularity of the lclVar's defined type (e.g. x86).
-        var_types lclTyp = genActualType(varDsc->TypeGet());
-        emitAttr  size   = emitTypeSize(lclTyp);
+        var_types spillType = genActualType(varDsc->GetRegisterType(tree->AsLclVar()));
+        emitAttr  size      = emitTypeSize(spillType);
 
         // If this is a write-thru variable, we don't actually spill at a use, but we will kill the var in the reg
         // (below).
         if (!varDsc->lvLiveInOutOfHndlr)
         {
-            instruction storeIns = ins_Store(lclTyp, compiler->isSIMDTypeLocalAligned(varNum));
+            instruction storeIns = ins_Store(spillType, compiler->isSIMDTypeLocalAligned(varNum));
             assert(varDsc->GetRegNum() == tree->GetRegNum());
             inst_TT_RV(storeIns, size, tree, tree->GetRegNum());
         }
@@ -1181,7 +1181,8 @@ void CodeGen::genUnspillRegIfNeeded(GenTree* tree)
 
             GenTreeLclVar* lcl       = unspillTree->AsLclVar();
             LclVarDsc*     varDsc    = compiler->lvaGetDesc(lcl->GetLclNum());
-            var_types      spillType = unspillTree->TypeGet();
+            var_types      spillType = varDsc->GetRegisterType(lcl);
+            assert(spillType != TYP_UNDEF);
 
 // TODO-Cleanup: The following code could probably be further merged and cleaned up.
 #ifdef TARGET_XARCH
@@ -1196,11 +1197,12 @@ void CodeGen::genUnspillRegIfNeeded(GenTree* tree)
             // later used as a long, we will have incorrectly truncated the long.
             // In the normalizeOnLoad case ins_Load will return an appropriate sign- or zero-
             // extending load.
-
-            if (spillType != genActualType(varDsc->lvType) && !varTypeIsGC(spillType) && !varDsc->lvNormalizeOnLoad())
+            var_types lclActualType = genActualType(varDsc->GetRegisterType());
+            assert(lclActualType != TYP_UNDEF);
+            if (spillType != lclActualType && !varTypeIsGC(spillType) && !varDsc->lvNormalizeOnLoad())
             {
                 assert(!varTypeIsGC(varDsc));
-                spillType = genActualType(varDsc->lvType);
+                spillType = lclActualType;
             }
 #elif defined(TARGET_ARM64)
             var_types targetType = unspillTree->gtType;
@@ -2060,8 +2062,9 @@ void CodeGen::genSpillLocal(unsigned varNum, var_types type, GenTreeLclVar* lclN
     {
         // Store local variable to its home location.
         // Ensure that lclVar stores are typed correctly.
-        GetEmitter()->emitIns_S_R(ins_Store(type, compiler->isSIMDTypeLocalAligned(varNum)), emitTypeSize(type), regNum,
-                                  varNum, 0);
+        var_types spillType = genActualType(varDsc->GetRegisterType(lclNode));
+        GetEmitter()->emitIns_S_R(ins_Store(spillType, compiler->isSIMDTypeLocalAligned(varNum)),
+                                  emitTypeSize(spillType), regNum, varNum, 0);
     }
 }
 

src/coreclr/jit/codegenxarch.cpp

@@ -4512,7 +4512,7 @@ void CodeGen::genCodeForStoreLclVar(GenTreeLclVar* lclNode)
             else if (op1->GetRegNum() != targetReg)
             {
                 assert(op1->GetRegNum() != REG_NA);
-                emit->emitInsBinary(ins_Move_Extend(targetType, true), emitTypeSize(lclNode), lclNode, op1);
+                emit->emitInsBinary(ins_Move_Extend(targetType, true), emitTypeSize(targetType), lclNode, op1);
             }
         }
         if (targetReg != REG_NA)

src/coreclr/jit/compiler.h

@@ -971,7 +971,7 @@ class LclVarDsc
     var_types lvaArgType();
 
     // Returns true if this variable contains GC pointers (including being a GC pointer itself).
-    bool HasGCPtr()
+    bool HasGCPtr() const
     {
         return varTypeIsGC(lvType) || ((lvType == TYP_STRUCT) && m_layout->HasGCPtr());
     }
@@ -1001,33 +1001,13 @@ class LclVarDsc
         return lvPerSsaData.GetSsaDef(ssaNum);
     }
 
-    //------------------------------------------------------------------------
-    // GetRegisterType: Determine register type for that local var.
-    //
-    // Arguments:
-    //    tree - node that uses the local, its type is checked first.
-    //
-    // Return Value:
-    //    TYP_UNDEF if the layout is enregistrable, register type otherwise.
-    //
-    var_types GetRegisterType(const GenTreeLclVarCommon* tree) const
-    {
-        var_types targetType = tree->gtType;
+    var_types GetRegisterType(const GenTreeLclVarCommon* tree) const;
 
-#ifdef DEBUG
-        // Ensure that lclVar nodes are typed correctly.
-        if (tree->OperIs(GT_STORE_LCL_VAR) && lvNormalizeOnStore())
-        {
-            // TODO: update that assert to work with TypeGet() == TYP_STRUCT case.
-            // assert(targetType == genActualType(TypeGet()));
-        }
-#endif
+    var_types GetRegisterType() const;
 
-        if (targetType != TYP_STRUCT)
-        {
-            return targetType;
-        }
-        return GetLayout()->GetRegisterType();
+    bool IsEnregisterable() const
+    {
+        return GetRegisterType() != TYP_UNDEF;
     }
 
     bool CanBeReplacedWithItsField(Compiler* comp) const;
@@ -9645,6 +9625,11 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 #endif // FEATURE_MULTIREG_RET
     }
 
+    bool compEnregStructLocals()
+    {
+        return JitConfig.JitEnregStructLocals();
+    }
+
     // Returns true if the method returns a value in more than one return register,
     // it should replace/be  merged with compMethodReturnsMultiRegRetType when #36868 is fixed.
     // The difference from original `compMethodReturnsMultiRegRetType` is in ARM64 SIMD* handling,

src/coreclr/jit/jitconfigvalues.h

@@ -515,6 +515,13 @@ CONFIG_INTEGER(JitSaveFpLrWithCalleeSavedRegisters, W("JitSaveFpLrWithCalleeSave
 #endif // defined(TARGET_ARM64)
 #endif // DEBUG
 
+// Allow to enregister locals with struct type.
+#if defined(TARGET_AMD64) && !defined(UNIX_AMD64_ABI)
+CONFIG_INTEGER(JitEnregStructLocals, W("JitEnregStructLocals"), 1)
+#else  // !TARGET_AMD64 && !UNIX_AMD64_ABI
+CONFIG_INTEGER(JitEnregStructLocals, W("JitEnregStructLocals"), 0)
+#endif // !TARGET_AMD64 && !UNIX_AMD64_ABI
+
 #undef CONFIG_INTEGER
 #undef CONFIG_STRING
 #undef CONFIG_METHODSET

src/coreclr/jit/lclmorph.cpp

@@ -539,7 +539,7 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
                             !m_compiler->lvaIsImplicitByRefLocal(lclVar->GetLclNum()))
                         {
                             LclVarDsc* varDsc = m_compiler->lvaGetDesc(lclNum);
-                            if (varDsc->lvFieldCnt > 1)
+                            if ((varDsc->lvFieldCnt > 1) && !varDsc->lvRegStruct)
                             {
                                 m_compiler->lvaSetVarDoNotEnregister(lclNum DEBUGARG(Compiler::DNER_BlockOp));
                             }

src/coreclr/jit/lclvars.cpp

@@ -2107,14 +2107,22 @@ bool Compiler::StructPromotionHelper::ShouldPromoteStructVar(unsigned lclNum)
         // TODO-1stClassStructs: a temporary solution to keep diffs small, it will be fixed later.
         shouldPromote = false;
     }
+    else if (compiler->compEnregStructLocals() && (structPromotionInfo.fieldCnt > 1) && !varDsc->lvFieldAccessed &&
+             varDsc->IsEnregisterable())
+    {
+        JITDUMP("Not promoting promotable struct local V%02u: #fields = %d because can put the whole struct into a "
+                "register\n",
+                lclNum, structPromotionInfo.fieldCnt);
+        shouldPromote = false;
+    }
 #if defined(DEBUG)
     else if (compiler->compPromoteFewerStructs(lclNum))
     {
         // Do not promote some structs, that can be promoted, to stress promoted/unpromoted moves.
         JITDUMP("Not promoting promotable struct local V%02u, because of STRESS_PROMOTE_FEWER_STRUCTS\n", lclNum);
         shouldPromote = false;
     }
-#endif
+#endif // DEBUG
 
     //
     // If the lvRefCnt is zero and we have a struct promoted parameter we can end up with an extra store of
@@ -3829,6 +3837,57 @@ var_types LclVarDsc::lvaArgType()
     return type;
 }
 
+//------------------------------------------------------------------------
+// GetRegisterType: Determine register type for this local var.
+//
+// Arguments:
+//    tree - node that uses the local, its type is checked first.
+//
+// Return Value:
+//    TYP_UNDEF if the layout is enregistrable, register type otherwise.
+//
+var_types LclVarDsc::GetRegisterType(const GenTreeLclVarCommon* tree) const
+{
+    var_types targetType = tree->gtType;
+
+#ifdef DEBUG
+    // Ensure that lclVar nodes are typed correctly.
+    if (tree->OperIs(GT_STORE_LCL_VAR) && lvNormalizeOnStore())
+    {
+        // TODO: update that assert to work with TypeGet() == TYP_STRUCT case.
+        // assert(targetType == genActualType(TypeGet()));
+    }
+#endif
+
+    if (targetType != TYP_STRUCT)
+    {
+        return targetType;
+    }
+    return GetLayout()->GetRegisterType();
+}
+
+//------------------------------------------------------------------------
+// GetRegisterType: Determine register type for this local var.
+//
+// Return Value:
+//    TYP_UNDEF if the layout is enregistrable, register type otherwise.
+//
+var_types LclVarDsc::GetRegisterType() const
+{
+    if (TypeGet() != TYP_STRUCT)
+    {
+#if !defined(TARGET_64BIT)
+        if (TypeGet() == TYP_LONG)
+        {
+            return TYP_UNDEF;
+        }
+#endif
+        return TypeGet();
+    }
+    assert(m_layout != nullptr);
+    return m_layout->GetRegisterType();
+}
+
 //----------------------------------------------------------------------------------------------
 // CanBeReplacedWithItsField: check if a whole struct reference could be replaced by a field.
 //
@@ -7337,6 +7396,11 @@ void Compiler::lvaDumpEntry(unsigned lclNum, FrameLayoutState curState, size_t r
 
         printf(" (%3u,%*s)", varDsc->lvRefCnt(), (int)refCntWtdWidth, refCntWtd2str(varDsc->lvRefCntWtd()));
 
+        if (varTypeIsStruct(type) && varDsc->lvRegStruct)
+        {
+            printf("reg ");
+        }
+
         printf(" %7s ", varTypeName(type));
         if (genTypeSize(type) == 0)
         {

src/coreclr/jit/lower.cpp

@@ -3166,9 +3166,12 @@ void Lowering::LowerStoreLocCommon(GenTreeLclVarCommon* lclStore)
             }
 #endif // !WINDOWS_AMD64_ABI
         }
-        else if (!src->OperIs(GT_LCL_VAR) || (varDsc->GetLayout()->GetRegisterType() == TYP_UNDEF))
+        else if (!src->OperIs(GT_LCL_VAR) || !varDsc->IsEnregisterable())
         {
-            GenTreeLclVar* addr = comp->gtNewLclVarAddrNode(lclStore->GetLclNum(), TYP_BYREF);
+            const unsigned lclNum = lclStore->GetLclNum();
+            GenTreeLclVar* addr   = comp->gtNewLclVarAddrNode(lclNum, TYP_BYREF);
+            // TODO-1stClassStructs: always keep it as a STORE_LCL_VAR.
+            comp->lvaSetVarDoNotEnregister(lclNum DEBUGARG(Compiler::DNER_BlockOp));
 
             addr->gtFlags |= GTF_VAR_DEF;
             assert(!addr->IsPartialLclFld(comp));
@@ -5915,17 +5918,35 @@ void Lowering::CheckNode(Compiler* compiler, GenTree* node)
         {
             unsigned   lclNum = node->AsLclVarCommon()->GetLclNum();
             LclVarDsc* lclVar = &compiler->lvaTable[lclNum];
-#ifdef DEBUG
             if (node->TypeIs(TYP_SIMD12))
             {
                 assert(compiler->lvaIsFieldOfDependentlyPromotedStruct(lclVar) || (lclVar->lvSize() == 12));
             }
-#endif // DEBUG
         }
         break;
 #endif // TARGET_64BIT
 #endif // SIMD
 
+        case GT_LCL_VAR_ADDR:
+        case GT_LCL_FLD_ADDR:
+        {
+            const GenTreeLclVarCommon* lclVarAddr = node->AsLclVarCommon();
+            const LclVarDsc*           varDsc     = compiler->lvaGetDesc(lclVarAddr);
+            if (((lclVarAddr->gtFlags & GTF_VAR_DEF) != 0) && varDsc->HasGCPtr())
+            {
+                // Emitter does not correctly handle live updates for LCL_VAR_ADDR
+                // when they are not contained, for example, `STOREIND byref(GT_LCL_VAR_ADDR not-contained)`
+                // would generate:
+                // add     r1, sp, 48   // r1 contains address of a lclVar V01.
+                // str     r0, [r1]     // a gc ref becomes live in V01, but emitter would not report it.
+                // Make sure that we use uncontained address nodes only for variables
+                // that will be marked as mustInit and will be alive throught the whole block even when tracked.
+                assert(lclVarAddr->isContained() || !varDsc->lvTracked || varTypeIsStruct(varDsc));
+                // TODO: support this assert for uses, see https://github.com/dotnet/runtime/issues/51900.
+            }
+            break;
+        }
+
         default:
             break;
     }

src/coreclr/jit/lowerxarch.cpp

@@ -230,6 +230,17 @@ void Lowering::LowerBlockStore(GenTreeBlk* blkNode)
             // address, not knowing that GT_IND is part of a block op that has containment restrictions.
             src->AsIndir()->Addr()->ClearContained();
         }
+        else if (src->OperIs(GT_LCL_VAR))
+        {
+            const GenTreeLclVar* lclVar = src->AsLclVar();
+            const LclVarDsc*     varDsc = comp->lvaGetDesc(src->AsLclVar());
+            if (!varDsc->lvDoNotEnregister)
+            {
+                // TODO-1stClassStructs: delete this pessimization, support LCL_VAR in a reg as STORE_BLK source.
+                comp->lvaSetVarDoNotEnregister(lclVar->GetLclNum() DEBUG_ARG(Compiler::DNER_BlockOp));
+            }
+            assert(src->isContained());
+        }
 
         if (blkNode->OperIs(GT_STORE_OBJ))
         {