Enregister multireg lclVars

Allow struct lclVars that are returned in multiple registers to be
enregistered, as long as the fields are a match for the registers.

Fix dotnet#34105

src/coreclr/src/jit/codegenarm.cpp

@@ -1048,7 +1048,7 @@ void CodeGen::genCodeForStoreLclVar(GenTreeLclVar* tree)
 
     // var = call, where call returns a multi-reg return value
     // case is handled separately.
-    if (data->gtSkipReloadOrCopy()->IsMultiRegCall())
+    if (data->gtSkipReloadOrCopy()->IsMultiRegNode())
     {
         genMultiRegStoreToLocal(tree);
     }

src/coreclr/src/jit/codegenarm64.cpp

@@ -1925,21 +1925,37 @@ void CodeGen::genCodeForStoreLclVar(GenTreeLclVar* tree)
 {
     GenTree* data = tree->gtOp1;
 
-    // var = call, where call returns a multi-reg return value
-    // case is handled separately.
-    if (data->gtSkipReloadOrCopy()->IsMultiRegCall())
+    // Multi-reg nodes are handled separately.
+    if (data->gtSkipReloadOrCopy()->IsMultiRegNode())
     {
         genMultiRegStoreToLocal(tree);
+        return;
+    }
+
+    LclVarDsc* varDsc = compiler->lvaGetDesc(tree);
+    if (tree->IsMultiReg())
+    {
+        // This is the case of storing to a multi-reg HFA local from a fixed-size SIMD type.
+        assert(varTypeIsSIMD(data) && varDsc->lvIsHfa() && (varDsc->GetHfaType() == TYP_FLOAT));
+        regNumber    operandReg = genConsumeReg(data);
+        unsigned int regCount   = varDsc->lvFieldCnt;
+        for (unsigned i = 0; i < regCount; ++i)
+        {
+            regNumber varReg = tree->GetRegByIndex(i);
+            assert(varReg != REG_NA);
+            unsigned   fieldLclNum = varDsc->lvFieldLclStart + i;
+            LclVarDsc* fieldVarDsc = compiler->lvaGetDesc(fieldLclNum);
+            assert(fieldVarDsc->TypeGet() == TYP_FLOAT);
+            GetEmitter()->emitIns_R_R_I(INS_dup, emitTypeSize(TYP_FLOAT), varReg, operandReg, i);
+        }
     }
     else
     {
         regNumber targetReg = tree->GetRegNum();
         emitter*  emit      = GetEmitter();
 
-        unsigned varNum = tree->GetLclNum();
-        assert(varNum < compiler->lvaCount);
-        LclVarDsc* varDsc     = compiler->lvaGetDesc(varNum);
-        var_types  targetType = varDsc->GetRegisterType(tree);
+        unsigned  varNum     = tree->GetLclNum();
+        var_types targetType = varDsc->GetRegisterType(tree);
 
 #ifdef FEATURE_SIMD
         // storing of TYP_SIMD12 (i.e. Vector3) field

src/coreclr/src/jit/codegenarmarch.cpp

@@ -1377,7 +1377,8 @@ void CodeGen::genMultiRegStoreToLocal(GenTree* treeNode)
 
     // Assumption: current implementation requires that a multi-reg
     // var in 'var = call' is flagged as lvIsMultiRegRet to prevent it from
-    // being promoted.
+    // being promoted, unless compiler->lvaEnregMultiRegVars is true.
+
     unsigned   lclNum = treeNode->AsLclVarCommon()->GetLclNum();
     LclVarDsc* varDsc = compiler->lvaGetDesc(lclNum);
     if (op1->OperIs(GT_CALL))
@@ -1388,7 +1389,9 @@ void CodeGen::genMultiRegStoreToLocal(GenTree* treeNode)
 
     genConsumeRegs(op1);
 
-    int offset = 0;
+    int  offset        = 0;
+    bool isMultiRegVar = treeNode->IsMultiRegLclVar();
+    bool hasRegs       = false;
 
     // Check for the case of an enregistered SIMD type that's returned in multiple registers.
     if (varDsc->lvIsRegCandidate() && treeNode->GetRegNum() != REG_NA)
@@ -1442,6 +1445,11 @@ void CodeGen::genMultiRegStoreToLocal(GenTree* treeNode)
     }
     else
     {
+        if (isMultiRegVar)
+        {
+            assert(compiler->lvaEnregMultiRegVars);
+            assert(regCount == varDsc->lvFieldCnt);
+        }
         for (unsigned i = 0; i < regCount; ++i)
         {
             var_types type = actualOp1->GetRegTypeByIndex(i);
@@ -1455,13 +1463,55 @@ void CodeGen::genMultiRegStoreToLocal(GenTree* treeNode)
             }
 
             assert(reg != REG_NA);
-            GetEmitter()->emitIns_S_R(ins_Store(type), emitTypeSize(type), reg, lclNum, offset);
-            offset += genTypeSize(type);
+            regNumber varReg = REG_NA;
+            if (isMultiRegVar)
+            {
+                regNumber  varReg      = treeNode->GetRegByIndex(i);
+                unsigned   fieldLclNum = varDsc->lvFieldLclStart + i;
+                LclVarDsc* fieldVarDsc = compiler->lvaGetDesc(fieldLclNum);
+                var_types  type        = fieldVarDsc->TypeGet();
+                if (varReg != REG_NA)
+                {
+                    hasRegs = true;
+                    if (varReg != reg)
+                    {
+                        inst_RV_RV(ins_Copy(type), varReg, reg, type);
+                    }
+                    fieldVarDsc->SetRegNum(varReg);
+                }
+                else
+                {
+                    if (!treeNode->AsLclVar()->IsLastUse(i))
+                    {
+                        GetEmitter()->emitIns_S_R(ins_Store(type), emitTypeSize(type), reg, fieldLclNum, 0);
+                    }
+                    fieldVarDsc->SetRegNum(REG_STK);
+                }
+            }
+            else
+            {
+                GetEmitter()->emitIns_S_R(ins_Store(type), emitTypeSize(type), reg, lclNum, offset);
+                offset += genTypeSize(type);
+            }
         }
 
         // Update variable liveness.
-        genUpdateLife(treeNode);
-        varDsc->SetRegNum(REG_STK);
+        if (isMultiRegVar)
+        {
+            if (hasRegs)
+            {
+                genProduceReg(treeNode);
+            }
+            else
+            {
+                genUpdateLife(treeNode);
+            }
+        }
+        else
+        {
+            genUpdateLife(treeNode);
+            varDsc->SetRegNum(REG_STK);
+        }
     }
 }
 

src/coreclr/src/jit/codegencommon.cpp

@@ -11660,6 +11660,7 @@ void CodeGen::genStructReturn(GenTree* treeNode)
     {
         varDsc = compiler->lvaGetDesc(actualOp1->AsLclVar()->GetLclNum());
         retTypeDesc.InitializeStructReturnType(compiler, varDsc->lvVerTypeInfo.GetClassHandle());
+        assert(varDsc->lvIsMultiRegRet);
     }
     else
     {
@@ -11670,15 +11671,15 @@ void CodeGen::genStructReturn(GenTree* treeNode)
     assert(regCount <= MAX_RET_REG_COUNT);
 
 #if FEATURE_MULTIREG_RET
-    if (actualOp1->OperIs(GT_LCL_VAR) && (varTypeIsEnregisterable(op1)))
+    if (actualOp1->OperIs(GT_LCL_VAR) && varTypeIsEnregisterable(op1) && !actualOp1->AsLclVar()->IsMultiReg())
     {
         // Right now the only enregisterable structs supported are SIMD vector types.
         assert(varTypeIsSIMD(op1));
 #ifdef FEATURE_SIMD
         genSIMDSplitReturn(op1, &retTypeDesc);
 #endif // FEATURE_SIMD
     }
-    else if (actualOp1->OperIs(GT_LCL_VAR))
+    else if (actualOp1->OperIs(GT_LCL_VAR) && !actualOp1->AsLclVar()->IsMultiReg())
     {
         GenTreeLclVar* lclNode = actualOp1->AsLclVar();
         LclVarDsc*     varDsc  = compiler->lvaGetDesc(lclNode->GetLclNum());
@@ -11694,20 +11695,36 @@ void CodeGen::genStructReturn(GenTree* treeNode)
     }
     else
     {
-        assert(actualOp1->IsMultiRegCall());
         for (unsigned i = 0; i < regCount; ++i)
         {
             var_types type    = retTypeDesc.GetReturnRegType(i);
             regNumber toReg   = retTypeDesc.GetABIReturnReg(i);
             regNumber fromReg = op1->GetRegByIndex(i);
-            if (fromReg == REG_NA)
+            if ((fromReg == REG_NA) && op1->OperIs(GT_COPY))
             {
-                assert(op1->IsCopyOrReload());
+                // A copy that doesn't copy this field will have REG_NA.
+                // TODO-Cleanup: It would probably be better to always have a valid reg
+                // on a GT_COPY, unless the operand is actually spilled. Then we wouldn't have
+                // to check for this case (though we'd have to check in the genRegCopy that the
+                // reg is valid).
                 fromReg = actualOp1->GetRegByIndex(i);
             }
-            if (fromReg != toReg)
+            if (fromReg == REG_NA)
+            {
+                // This is a spilled field of a multi-reg lclVar.
+                // We currently only mark a lclVar operand as RegOptional, since we don't have a way
+                // to mark a multi-reg tree node as used from spill (GTF_NOREG_AT_USE) on a per-reg basis.
+                assert(varDsc != nullptr);
+                assert(varDsc->lvPromoted);
+                unsigned fieldVarNum = varDsc->lvFieldLclStart + i;
+                assert(compiler->lvaGetDesc(fieldVarNum)->lvOnFrame);
+                GetEmitter()->emitIns_R_S(ins_Load(type), emitTypeSize(type), toReg, fieldVarNum, 0);
+            }
+            else if (fromReg != toReg)
             {
-                inst_RV_RV(ins_Copy(type), toReg, fromReg, type);
+                // Note that ins_Copy(fromReg, type) will return the appropriate register to copy
+                // between register files if needed.
+                inst_RV_RV(ins_Copy(fromReg, type), toReg, fromReg, type);
             }
         }
     }
@@ -11774,6 +11791,7 @@ void CodeGen::genRegCopy(GenTree* treeNode)
             // on the source is still valid at the consumer).
             if (targetReg != REG_NA)
             {
+                regMaskTP targetRegMask = genRegMask(targetReg);
                 // We shouldn't specify a no-op move.
                 regMaskTP targetRegMask = genRegMask(targetReg);
                 assert(sourceReg != targetReg);
@@ -11852,7 +11870,7 @@ void CodeGen::genRegCopy(GenTree* treeNode)
 
             if ((lcl->gtFlags & GTF_VAR_DEATH) == 0 && (treeNode->gtFlags & GTF_VAR_DEATH) == 0)
             {
-                LclVarDsc* varDsc = compiler->lvaGetDesc(lcl);
+                LclVarDsc* varDsc = &compiler->lvaTable[lcl->GetLclNum()];
 
                 // If we didn't just spill it (in genConsumeReg, above), then update the register info
                 if (varDsc->GetRegNum() != REG_STK)

src/coreclr/src/jit/codegeninterface.h

@@ -129,6 +129,7 @@ class CodeGenInterface
     // Liveness-related fields & methods
 public:
     void genUpdateRegLife(const LclVarDsc* varDsc, bool isBorn, bool isDying DEBUGARG(GenTree* tree));
+    void genUpdateVarReg(LclVarDsc* varDsc, GenTree* tree, int regIndex);
     void genUpdateVarReg(LclVarDsc* varDsc, GenTree* tree);
 
 protected: