Add lvIsImplicitByRef information to lvaSetStruct (#19223)

Before implicit byrefs were tracked by setting lvIsParam and lvIsTemp.
This change explicitly adds a flag for implicitByRef instead of overloading.
In addition, it fixes the decision to copy an implicitByRef for arm64 varargs.

Temporarily bump weight on byref params to match old behavior and avoid codegen
diffs.

Re-enabled various tests and parts of tests.

Closes #20046
Closes #19860

7 files changed, 80 insertions, 72 deletions

diff --git a/src/jit/compiler.cpp b/src/jit/compiler.cpp
index 15ed19c62d..378d83bd1f 100644
--- a/src/jit/compiler.cpp
+++ b/src/jit/compiler.cpp
@@ -704,7 +704,7 @@ var_types Compiler::getPrimitiveTypeForStruct(unsigned structSize, CORINFO_CLASS
 // getArgTypeForStruct:
 //     Get the type that is used to pass values of the given struct type.
 //     If you have already retrieved the struct size then it should be
-//     passed as the optional third argument, as this allows us to avoid
+//     passed as the optional fourth argument, as this allows us to avoid
 //     an extra call to getClassSize(clsHnd)
 //
 // Arguments:
diff --git a/src/jit/compiler.h b/src/jit/compiler.h
index dc1f844a31..dec2343356 100644
--- a/src/jit/compiler.h
+++ b/src/jit/compiler.h
@@ -610,8 +610,12 @@ public:
     unsigned char lvHasILStoreOp : 1;         // there is at least one STLOC or STARG on this local
     unsigned char lvHasMultipleILStoreOp : 1; // there is more than one STLOC on this local
 
-    unsigned char lvIsTemp : 1; // Short-lifetime compiler temp (if lvIsParam is false), or implicit byref parameter
-                                // (if lvIsParam is true)
+    unsigned char lvIsTemp : 1; // Short-lifetime compiler temp
+
+#if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+    unsigned char lvIsImplicitByRef : 1; // Set if the argument is an implicit byref.
+#endif                                   // defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+
 #if OPT_BOOL_OPS
     unsigned char lvIsBoolean : 1; // set if variable is boolean
 #endif
@@ -957,7 +961,7 @@ public:
 
 private:
     unsigned short m_lvRefCnt; // unweighted (real) reference count.  For implicit by reference
-                               // parameters, this gets hijacked from fgMarkImplicitByRefArgs
+                               // parameters, this gets hijacked from fgResetImplicitByRefRefCount
                                // through fgMarkDemotedImplicitByRefArgs, to provide a static
                                // appearance count (computed during address-exposed analysis)
                                // that fgMakeOutgoingStructArgCopy consults during global morph
@@ -3346,16 +3350,16 @@ public:
     BOOL lvaIsOriginalThisReadOnly();           // return TRUE if there is no place in the code
                                                 // that writes to arg0
 
-    // Struct parameters that are passed by reference are marked as both lvIsParam and lvIsTemp
-    // (this is an overload of lvIsTemp because there are no temp parameters).
     // For x64 this is 3, 5, 6, 7, >8 byte structs that are passed by reference.
     // For ARM64, this is structs larger than 16 bytes that are passed by reference.
     bool lvaIsImplicitByRefLocal(unsigned varNum)
     {
 #if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
-        LclVarDsc* varDsc = &(lvaTable[varNum]);
-        if (varDsc->lvIsParam && varDsc->lvIsTemp)
+        LclVarDsc* varDsc = lvaGetDesc(varNum);
+        if (varDsc->lvIsImplicitByRef)
         {
+            assert(varDsc->lvIsParam);
+
             assert(varTypeIsStruct(varDsc) || (varDsc->lvType == TYP_BYREF));
             return true;
         }
@@ -5667,8 +5671,8 @@ private:
     void fgMorphStructField(GenTree* tree, GenTree* parent);
     void fgMorphLocalField(GenTree* tree, GenTree* parent);
 
-    // Identify which parameters are implicit byrefs, and flag their LclVarDscs.
-    void fgMarkImplicitByRefArgs();
+    // Reset the refCount for implicit byrefs.
+    void fgResetImplicitByRefRefCount();
 
     // Change implicit byrefs' types from struct to pointer, and for any that were
     // promoted, create new promoted struct temps.
diff --git a/src/jit/compiler.hpp b/src/jit/compiler.hpp
index 0b30114f67..9b10542722 100644
--- a/src/jit/compiler.hpp
+++ b/src/jit/compiler.hpp
@@ -1756,8 +1756,15 @@ inline void LclVarDsc::incRefCnts(BasicBlock::weight_t weight, Compiler* comp, R
         if (weight != 0)
         {
             // We double the weight of internal temps
-            //
-            if (lvIsTemp && (weight * 2 > weight))
+
+            bool doubleWeight = lvIsTemp;
+
+#if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+            // and, for the time being, implict byref params
+            doubleWeight |= lvIsImplicitByRef;
+#endif // defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+
+            if (doubleWeight && (weight * 2 > weight))
             {
                 weight *= 2;
             }
diff --git a/src/jit/importer.cpp b/src/jit/importer.cpp
index 88638092fe..78c0b4fde0 100644
--- a/src/jit/importer.cpp
+++ b/src/jit/importer.cpp
@@ -9022,13 +9022,16 @@ GenTree* Compiler::impFixupStructReturnType(GenTree* op, CORINFO_CLASS_HANDLE re
             // This LCL_VAR stays as a TYP_STRUCT
             unsigned lclNum = op->gtLclVarCommon.gtLclNum;
 
-            // Make sure this struct type is not struct promoted
-            lvaTable[lclNum].lvIsMultiRegRet = true;
+            if (!lvaIsImplicitByRefLocal(lclNum))
+            {
+                // Make sure this struct type is not struct promoted
+                lvaTable[lclNum].lvIsMultiRegRet = true;
 
-            // TODO-1stClassStructs: Handle constant propagation and CSE-ing of multireg returns.
-            op->gtFlags |= GTF_DONT_CSE;
+                // TODO-1stClassStructs: Handle constant propagation and CSE-ing of multireg returns.
+                op->gtFlags |= GTF_DONT_CSE;
 
-            return op;
+                return op;
+            }
         }
 
         if (op->gtOper == GT_CALL)
diff --git a/src/jit/lclvars.cpp b/src/jit/lclvars.cpp
index d7d0f11134..515e9348c0 100644
--- a/src/jit/lclvars.cpp
+++ b/src/jit/lclvars.cpp
@@ -1280,6 +1280,10 @@ void Compiler::lvaInitVarDsc(LclVarDsc*              varDsc,
         varDsc->lvStructGcCount = 1;
     }
 
+#if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+    varDsc->lvIsImplicitByRef = 0;
+#endif // defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+
 // Set the lvType (before this point it is TYP_UNDEF).
 
 #ifdef FEATURE_HFA
@@ -2163,6 +2167,10 @@ void Compiler::StructPromotionHelper::PromoteStructVar(unsigned lclNum)
         fieldVarDsc->lvParentLcl     = lclNum;
         fieldVarDsc->lvIsParam       = varDsc->lvIsParam;
 #if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+
+        // Reset the implicitByRef flag.
+        fieldVarDsc->lvIsImplicitByRef = 0;
+
         // Do we have a parameter that can be enregistered?
         //
         if (varDsc->lvIsRegArg)
@@ -2490,11 +2498,12 @@ void Compiler::lvaSetStruct(unsigned varNum, CORINFO_CLASS_HANDLE typeHnd, bool
             varDsc->lvExactSize = info.compCompHnd->getHeapClassSize(typeHnd);
         }
 
-        size_t lvSize = varDsc->lvSize();
-        assert((lvSize % TARGET_POINTER_SIZE) ==
-               0); // The struct needs to be a multiple of TARGET_POINTER_SIZE bytes for getClassGClayout() to be valid.
-        varDsc->lvGcLayout = getAllocator(CMK_LvaTable).allocate<BYTE>(lvSize / TARGET_POINTER_SIZE);
-        unsigned  numGCVars;
+        // Normalize struct types, and fill in GC info for all types
+        unsigned lvSize = varDsc->lvSize();
+        // The struct needs to be a multiple of TARGET_POINTER_SIZE bytes for getClassGClayout() to be valid.
+        assert((lvSize % TARGET_POINTER_SIZE) == 0);
+        varDsc->lvGcLayout     = getAllocator(CMK_LvaTable).allocate<BYTE>(lvSize / TARGET_POINTER_SIZE);
+        unsigned  numGCVars    = 0;
         var_types simdBaseType = TYP_UNKNOWN;
         if (isValueClass)
         {
@@ -2510,10 +2519,27 @@ void Compiler::lvaSetStruct(unsigned varNum, CORINFO_CLASS_HANDLE typeHnd, bool
         {
             numGCVars = 7;
         }
+
         varDsc->lvStructGcCount = numGCVars;
 
         if (isValueClass)
         {
+
+#if defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+            // Mark implicit byref struct parameters
+            if (varDsc->lvIsParam && !varDsc->lvIsStructField)
+            {
+                structPassingKind howToReturnStruct;
+                getArgTypeForStruct(typeHnd, &howToReturnStruct, this->info.compIsVarArgs, varDsc->lvExactSize);
+
+                if (howToReturnStruct == SPK_ByReference)
+                {
+                    JITDUMP("Marking V%02i as a byref parameter\n", varNum);
+                    varDsc->lvIsImplicitByRef = 1;
+                }
+            }
+#endif // defined(_TARGET_AMD64_) || defined(_TARGET_ARM64_)
+
 #if FEATURE_SIMD
             if (simdBaseType != TYP_UNKNOWN)
             {
diff --git a/src/jit/morph.cpp b/src/jit/morph.cpp
index fbfa173311..83e069cebf 100644
--- a/src/jit/morph.cpp
+++ b/src/jit/morph.cpp
@@ -17061,9 +17061,8 @@ void Compiler::fgMorph()
     fgUpdateFinallyTargetFlags();
 
     /* For x64 and ARM64 we need to mark irregular parameters */
-
     lvaRefCountState = RCS_EARLY;
-    fgMarkImplicitByRefArgs();
+    fgResetImplicitByRefRefCount();
 
     /* Promote struct locals if necessary */
     fgPromoteStructs();
@@ -17491,63 +17490,29 @@ void Compiler::fgMorphLocalField(GenTree* tree, GenTree* parent)
 }
 
 //------------------------------------------------------------------------
-// fgMarkImplicitByRefArgs: Identify any by-value struct parameters which are "implicit by-reference";
-//                          i.e. which the ABI requires to be passed by making a copy in the caller and
-//                          passing its address to the callee.  Mark their `LclVarDsc`s such that
-//                          `lvaIsImplicitByRefLocal` will return true for them.
+// fgResetImplicitByRefRefCount: Clear the ref count field of all implicit byrefs
 
-void Compiler::fgMarkImplicitByRefArgs()
+void Compiler::fgResetImplicitByRefRefCount()
 {
 #if (defined(_TARGET_AMD64_) && !defined(UNIX_AMD64_ABI)) || defined(_TARGET_ARM64_)
 #ifdef DEBUG
     if (verbose)
     {
-        printf("\n*************** In fgMarkImplicitByRefs()\n");
+        printf("\n*************** In fgResetImplicitByRefRefCount()\n");
     }
 #endif // DEBUG
 
-    for (unsigned lclNum = 0; lclNum < info.compArgsCount; lclNum++)
+    for (unsigned lclNum = 0; lclNum < info.compArgsCount; ++lclNum)
     {
-        LclVarDsc* varDsc = &lvaTable[lclNum];
+        LclVarDsc* varDsc = lvaGetDesc(lclNum);
 
-        if (varDsc->lvIsParam && varTypeIsStruct(varDsc))
+        if (varDsc->lvIsImplicitByRef)
         {
-            size_t size = varDsc->lvExactSize;
-            assert(size == info.compCompHnd->getClassSize(varDsc->lvVerTypeInfo.GetClassHandle()));
-
-            bool isPassedByReference;
-#if defined(_TARGET_AMD64_)
-            isPassedByReference = (size > REGSIZE_BYTES || (size & (size - 1)) != 0);
-#elif defined(_TARGET_ARM64_)
-            if (size > TARGET_POINTER_SIZE)
-            {
-                CORINFO_CLASS_HANDLE clsHnd = varDsc->lvVerTypeInfo.GetClassHandleForValueClass();
-                structPassingKind    howToPassStruct;
-                var_types            type =
-                    getArgTypeForStruct(clsHnd, &howToPassStruct, this->info.compIsVarArgs, varDsc->lvExactSize);
-                isPassedByReference = (howToPassStruct == SPK_ByReference);
-            }
-            else
-            {
-                isPassedByReference = false;
-            }
-#endif
-
-            if (isPassedByReference)
-            {
-                // Previously nobody was ever setting lvIsParam and lvIsTemp on the same local
-                // So I am now using it to indicate that this is one of the weird implicit
-                // by ref locals.
-                // The address taken cleanup will look for references to locals marked like
-                // this, and transform them appropriately.
-                varDsc->lvIsTemp = 1;
-
-                // Clear the ref count field; fgMarkAddressTakenLocals will increment it per
-                // 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, RCS_EARLY);
-            }
+            // Clear the ref count field; fgMarkAddressTakenLocals will increment it per
+            // 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, RCS_EARLY);
         }
     }
 
diff --git a/src/jit/scopeinfo.cpp b/src/jit/scopeinfo.cpp
index be6fb8ce6c..d396b333a4 100644
--- a/src/jit/scopeinfo.cpp
+++ b/src/jit/scopeinfo.cpp
@@ -301,14 +301,17 @@ void CodeGenInterface::siVarLoc::siFillStackVarLoc(
             // size is not 1, 2, 4 or 8 bytes in size. During fgMorph, the compiler modifies
             // the IR to comply with the ABI and therefore changes the type of the lclVar
             // that holds the struct from TYP_STRUCT to TYP_BYREF but it gives us a hint that
-            // this is still a struct by setting the lvIsTemp flag.
+            // this is still a struct by setting the lvIsImplicitByref flag.
             // The same is true for ARM64 and structs > 16 bytes.
-            // (See Compiler::fgMarkImplicitByRefArgs in Morph.cpp for further detail)
+            //
+            // See lvaSetStruct for further detail.
+            //
             // Now, the VM expects a special enum for these type of local vars: VLT_STK_BYREF
             // to accomodate for this situation.
-            if (varDsc->lvType == TYP_BYREF && varDsc->lvIsTemp)
+            if (varDsc->lvIsImplicitByRef)
             {
                 assert(varDsc->lvIsParam);
+                assert(varDsc->lvType == TYP_BYREF);
                 this->vlType = VLT_STK_BYREF;
             }
             else