Merge pull request #11883 from JosephTremoulet/PromotedStructRefs

Update full-struct references to promoted IBR args

2 files changed, 106 insertions, 61 deletions

diff --git a/src/jit/compiler.h b/src/jit/compiler.h
index 6f38e00160..c8ead479b6 100644
--- a/src/jit/compiler.h
+++ b/src/jit/compiler.h
@@ -255,7 +255,8 @@ 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
+    unsigned char lvIsTemp : 1; // Short-lifetime compiler temp (if lvIsParam is false), or implicit byref parameter
+                                // (if lvIsParam is true)
 #if OPT_BOOL_OPS
     unsigned char lvIsBoolean : 1; // set if variable is boolean
 #endif
@@ -286,7 +287,9 @@ public:
                                // checks)
     unsigned char lvIsUnsafeBuffer : 1; // Does this contain an unsafe buffer requiring buffer overflow security checks?
     unsigned char lvPromoted : 1; // True when this local is a promoted struct, a normed struct, or a "split" long on a
-                                  // 32-bit target.
+                                  // 32-bit target.  For implicit byref parameters, this gets hijacked between
+                                  // fgRetypeImplicitByRefArgs and fgMarkDemotedImplicitByRefArgs to indicate whether
+                                  // references to the arg are being rewritten as references to a promoted shadow local.
     unsigned char lvIsStructField : 1;          // Is this local var a field of a promoted struct local?
     unsigned char lvContainsFloatingFields : 1; // Does this struct contains floating point fields?
     unsigned char lvOverlappingFields : 1;      // True when we have a struct with possibly overlapping fields
@@ -332,7 +335,9 @@ public:
 
     union {
         unsigned lvFieldLclStart; // The index of the local var representing the first field in the promoted struct
-                                  // local.
+                                  // local.  For implicit byref parameters, this gets hijacked between
+                                  // fgRetypeImplicitByRefArgs and fgMarkDemotedImplicitByRefArgs to point to the
+                                  // struct local created to model the parameter's struct promotion, if any.
         unsigned lvParentLcl; // The index of the local var representing the parent (i.e. the promoted struct local).
                               // Valid on promoted struct local fields.
     };
@@ -656,11 +661,16 @@ public:
 
     regMaskSmall lvPrefReg; // set of regs it prefers to live in
 
-    unsigned short lvVarIndex;  // variable tracking index
-    unsigned short lvRefCnt;    // unweighted (real) reference count
-    unsigned       lvRefCntWtd; // weighted reference count
-    int            lvStkOffs;   // stack offset of home
-    unsigned       lvExactSize; // (exact) size of the type in bytes
+    unsigned short lvVarIndex; // variable tracking index
+    unsigned short lvRefCnt;   // unweighted (real) reference count.  For implicit by reference
+                               // parameters, this gets hijacked from fgMarkImplicitByRefArgs
+                               // through fgMarkDemotedImplicitByRefArgs, to provide a static
+                               // appearance count (computed during address-exposed analysis)
+                               // that fgMakeOutgoingStructArgCopy consults during global morph
+                               // to determine if eliding its copy is legal.
+    unsigned lvRefCntWtd;      // weighted reference count
+    int      lvStkOffs;        // stack offset of home
+    unsigned lvExactSize;      // (exact) size of the type in bytes
 
     // Is this a promoted struct?
     // This method returns true only for structs (including SIMD structs), not for
@@ -4879,8 +4889,7 @@ private:
     bool fgMorphImplicitByRefArgs(GenTreePtr tree);
     GenTreePtr fgMorphImplicitByRefArgs(GenTreePtr tree, bool isAddr);
 
-    // Clear up annotations for any struct promotion temps created for implicit byrefs that
-    // wound up unused (due e.g. to being address-exposed and not worth promoting).
+    // Clear up annotations for any struct promotion temps created for implicit byrefs.
     void fgMarkDemotedImplicitByRefArgs();
 
     static fgWalkPreFn  fgMarkAddrTakenLocalsPreCB;
diff --git a/src/jit/morph.cpp b/src/jit/morph.cpp
index 32d3f8fc79..319d1f98ac 100644
--- a/src/jit/morph.cpp
+++ b/src/jit/morph.cpp
@@ -17916,19 +17916,32 @@ void Compiler::fgRetypeImplicitByRefArgs()
                     fieldVarDsc->lvPrefReg = 0;
                 }
 
-                if (undoPromotion)
-                {
-                    // Hijack lvFieldLclStart to record the new temp number.
-                    // It will get fixed up in fgMarkDemotedImplicitByRefArgs.
-                    varDsc->lvFieldLclStart = newLclNum;
-                }
-                else
-                {
-                    // Unmark the parameter as promoted (it's a pointer now).
-                    varDsc->lvPromoted      = false;
-                    varDsc->lvFieldCnt      = 0;
-                    varDsc->lvFieldLclStart = 0;
-                }
+                // Hijack lvFieldLclStart to record the new temp number.
+                // It will get fixed up in fgMarkDemotedImplicitByRefArgs.
+                varDsc->lvFieldLclStart = newLclNum;
+                // Go ahead and clear lvFieldCnt -- either we're promoting
+                // a replacement temp or we're not promoting this arg, and
+                // in either case the parameter is now a pointer that doesn't
+                // have these fields.
+                varDsc->lvFieldCnt = 0;
+
+                // Hijack lvPromoted to communicate to fgMorphImplicitByRefArgs
+                // whether references to the struct should be rewritten as
+                // indirections off the pointer (not promoted) or references
+                // to the new struct local (promoted).
+                varDsc->lvPromoted = !undoPromotion;
+            }
+            else
+            {
+                // The "undo promotion" path above clears lvPromoted for args that struct
+                // promotion wanted to promote but that aren't considered profitable to
+                // rewrite.  It hijacks lvFieldLclStart to communicate to
+                // fgMarkDemotedImplicitByRefArgs that it needs to clean up annotations left
+                // on such args for fgMorphImplicitByRefArgs to consult in the interim.
+                // Here we have an arg that was simply never promoted, so make sure it doesn't
+                // have nonzero lvFieldLclStart, since that would confuse fgMorphImplicitByRefArgs
+                // and fgMarkDemotedImplicitByRefArgs.
+                assert(varDsc->lvFieldLclStart == 0);
             }
 
             // Since the parameter in this position is really a pointer, its type is TYP_BYREF.
@@ -17965,10 +17978,9 @@ void Compiler::fgRetypeImplicitByRefArgs()
 
 //------------------------------------------------------------------------
 // fgMarkDemotedImplicitByRefArgs: Clear annotations for any implicit byrefs that struct promotion
-//                                 asked to promote but for which fgRetypeImplicitByRefArgs decided
-//                                 to discard the promotion.  Appearances of these have now been
-//                                 rewritten (by fgMorphImplicitByRefArgs) using indirections from
-//                                 the pointer parameter.
+//                                 asked to promote.  Appearances of these have now been rewritten
+//                                 (by fgMorphImplicitByRefArgs) using indirections from the pointer
+//                                 parameter or references to the promotion temp, as appropriate.
 
 void Compiler::fgMarkDemotedImplicitByRefArgs()
 {
@@ -17978,45 +17990,60 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
     {
         LclVarDsc* varDsc = &lvaTable[lclNum];
 
-        if (lvaIsImplicitByRefLocal(lclNum) && varDsc->lvPromoted)
-        {
-            // We stashed the pointer to the real promotion temp in lvFieldLclStart
-            unsigned structLclNum = varDsc->lvFieldLclStart;
-
-            // Unmark the parameter as promoted.
-            varDsc->lvPromoted      = false;
-            varDsc->lvFieldCnt      = 0;
-            varDsc->lvFieldLclStart = 0;
-            // Clear its ref count; this was set during address-taken analysis so that
-            // 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->lvRefCnt = 0;
-
-            // 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->lvRefCnt      = 0;
-            structVarDsc->lvAddrExposed = false;
+        if (lvaIsImplicitByRefLocal(lclNum))
+        {
+            if (varDsc->lvPromoted)
+            {
+                // The parameter is simply a pointer now, so clear lvPromoted.  It was left set
+                // by fgRetypeImplicitByRefArgs to communicate to fgMorphImplicitByRefArgs that
+                // appearances of this arg needed to be rewritten to a new promoted struct local.
+                varDsc->lvPromoted = false;
+
+                // Clear the lvFieldLclStart value that was set by fgRetypeImplicitByRefArgs
+                // to tell fgMorphImplicitByRefArgs which local is the new promoted struct one.
+                varDsc->lvFieldLclStart = 0;
+            }
+            else if (varDsc->lvFieldLclStart != 0)
+            {
+                // We created new temps to represent a promoted struct corresponding to this
+                // parameter, but decided not to go through with the promotion and have
+                // rewritten all uses as indirections off the pointer parameter.
+                // We stashed the pointer to the new struct temp in lvFieldLclStart; make
+                // note of that and clear the annotation.
+                unsigned structLclNum   = varDsc->lvFieldLclStart;
+                varDsc->lvFieldLclStart = 0;
+
+                // Clear the arg's ref count; this was set during address-taken analysis so that
+                // 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->lvRefCnt = 0;
+
+                // 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->lvRefCnt      = 0;
+                structVarDsc->lvAddrExposed = false;
 #ifdef DEBUG
-            structVarDsc->lvUnusedStruct = true;
+                structVarDsc->lvUnusedStruct = true;
 #endif // DEBUG
 
-            unsigned fieldLclStart = structVarDsc->lvFieldLclStart;
-            unsigned fieldCount    = structVarDsc->lvFieldCnt;
-            unsigned fieldLclStop  = fieldLclStart + fieldCount;
+                unsigned fieldLclStart = structVarDsc->lvFieldLclStart;
+                unsigned fieldCount    = structVarDsc->lvFieldCnt;
+                unsigned fieldLclStop  = fieldLclStart + fieldCount;
 
-            for (unsigned fieldLclNum = fieldLclStart; fieldLclNum < fieldLclStop; ++fieldLclNum)
-            {
-                // Fix the pointer to the parent local.
-                LclVarDsc* fieldVarDsc = &lvaTable[fieldLclNum];
-                assert(fieldVarDsc->lvParentLcl == lclNum);
-                fieldVarDsc->lvParentLcl = structLclNum;
+                for (unsigned fieldLclNum = fieldLclStart; fieldLclNum < fieldLclStop; ++fieldLclNum)
+                {
+                    // Fix the pointer to the parent local.
+                    LclVarDsc* fieldVarDsc = &lvaTable[fieldLclNum];
+                    assert(fieldVarDsc->lvParentLcl == lclNum);
+                    fieldVarDsc->lvParentLcl = structLclNum;
 
-                // The field local is now unused; set flags appropriately so that
-                // we won't allocate stack space for it.
-                fieldVarDsc->lvRefCnt      = 0;
-                fieldVarDsc->lvAddrExposed = false;
+                    // The field local is now unused; set flags appropriately so that
+                    // we won't allocate stack space for it.
+                    fieldVarDsc->lvRefCnt      = 0;
+                    fieldVarDsc->lvAddrExposed = false;
+                }
             }
         }
     }
@@ -18094,8 +18121,17 @@ GenTreePtr Compiler::fgMorphImplicitByRefArgs(GenTreePtr tree, bool isAddr)
         if (!varTypeIsStruct(lclVarTree))
         {
             assert(lclVarTree->TypeGet() == TYP_BYREF);
+
             return nullptr;
         }
+        else if (lclVarDsc->lvPromoted)
+        {
+            // fgRetypeImplicitByRefArgs created a new promoted struct local to represent this
+            // arg.  Rewrite this to refer to the new local.
+            assert(lclVarDsc->lvFieldLclStart != 0);
+            lclVarTree->AsLclVarCommon()->SetLclNum(lclVarDsc->lvFieldLclStart);
+            return tree;
+        }
 
         fieldHnd = nullptr;
     }