JIT: Generalize assignment decomposition in physical promotion (#85323)

Generalize assignment decomposition to handle arbitrary combinations of
physically promoted structs. Do this by introducing a DecompositionPlan
class that keeps track of the copies to do that involve replacement
fields. The first step is then to fill out this plan. In the general
case where both the source and destination are physically promoted this
involves iterating the replacements in lockstep. For promotions that map
exactly, a direct copy between their locals is queued into the plan; in
other cases (e.g. partial overlap) it may involve writing the source
back to the struct local.

The plan is used to generate the IR and to figure out the best strategy to
use for the remaining parts of the struct. Additional it is used to check for
some optimization opportunities (e.g. we avoid superfluous write barriers
in some cases).

src/coreclr/jit/compiler.h

@@ -2518,6 +2518,8 @@ class Compiler
 
     GenTree* gtNewOneConNode(var_types type, var_types simdBaseType = TYP_UNDEF);
 
+    GenTree* gtNewConWithPattern(var_types type, uint8_t pattern);
+
     GenTreeLclVar* gtNewStoreLclVarNode(unsigned lclNum, GenTree* data);
 
     GenTreeLclFld* gtNewStoreLclFldNode(unsigned lclNum, var_types type, unsigned offset, GenTree* data);
@@ -6162,8 +6164,10 @@ class Compiler
     bool gtTreeContainsOper(GenTree* tree, genTreeOps op);
     ExceptionSetFlags gtCollectExceptions(GenTree* tree);
 
+public:
     bool fgIsBigOffset(size_t offset);
 
+private:
     bool fgNeedReturnSpillTemp();
 
     /*

src/coreclr/jit/gentree.cpp

@@ -7689,6 +7689,71 @@ GenTree* Compiler::gtNewOneConNode(var_types type, var_types simdBaseType /* = T
     }
 }
 
+//------------------------------------------------------------------------
+// CreateInitValue:
+//   Create an IR node representing a constant value with the specified 8
+//   byte character broadcast into all of its bytes.
+//
+// Parameters:
+//   type    - The primitive type. For small types the constant will be
+//             zero/sign-extended and a TYP_INT node will be returned.
+//   pattern - A byte pattern.
+//
+// Returns:
+//   An IR node representing the constant.
+//
+// Remarks:
+//   Should only be called when that pattern can actually be represented; for
+//   example, GC pointers only support an init pattern of zero.
+//
+GenTree* Compiler::gtNewConWithPattern(var_types type, uint8_t pattern)
+{
+    switch (type)
+    {
+        case TYP_BOOL:
+        case TYP_UBYTE:
+            return gtNewIconNode(pattern);
+        case TYP_BYTE:
+            return gtNewIconNode((int8_t)pattern);
+        case TYP_SHORT:
+            return gtNewIconNode((int16_t)(pattern * 0x0101));
+        case TYP_USHORT:
+            return gtNewIconNode((uint16_t)(pattern * 0x0101));
+        case TYP_INT:
+            return gtNewIconNode(pattern * 0x01010101);
+        case TYP_LONG:
+            return gtNewLconNode(pattern * 0x0101010101010101LL);
+        case TYP_FLOAT:
+            float floatPattern;
+            memset(&floatPattern, pattern, sizeof(floatPattern));
+            return gtNewDconNode(floatPattern, TYP_FLOAT);
+        case TYP_DOUBLE:
+            double doublePattern;
+            memset(&doublePattern, pattern, sizeof(doublePattern));
+            return gtNewDconNode(doublePattern);
+        case TYP_REF:
+        case TYP_BYREF:
+            assert(pattern == 0);
+            return gtNewZeroConNode(type);
+#ifdef FEATURE_SIMD
+        case TYP_SIMD8:
+        case TYP_SIMD12:
+        case TYP_SIMD16:
+#if defined(TARGET_XARCH)
+        case TYP_SIMD32:
+        case TYP_SIMD64:
+#endif // TARGET_XARCH
+#endif // FEATURE_SIMD
+        {
+            GenTreeVecCon* node = gtNewVconNode(type);
+            memset(&node->gtSimdVal, pattern, sizeof(node->gtSimdVal));
+            return node;
+        }
+        default:
+            unreached();
+    }
+}
+
 GenTreeLclVar* Compiler::gtNewStoreLclVarNode(unsigned lclNum, GenTree* data)
 {
     LclVarDsc*     varDsc = lvaGetDesc(lclNum);

src/coreclr/jit/jitstd/vector.h

@@ -479,7 +479,7 @@ typename vector<T, Allocator>::iterator
     assert(last.m_pElem >= m_pArray);
     assert(first.m_pElem <= m_pArray + m_nSize);
     assert(last.m_pElem <= m_pArray + m_nSize);
-    assert(last.m_pElem > first.m_pElem);
+    assert(last.m_pElem >= first.m_pElem);
 
     pointer fptr = first.m_pElem;
     pointer lptr = last.m_pElem;

src/coreclr/jit/morphblock.cpp

@@ -410,22 +410,19 @@ void MorphInitBlockHelper::TryInitFieldByField()
         return;
     }
 
-    const int64_t initPattern = (initVal->AsIntCon()->IconValue() & 0xFF) * 0x0101010101010101LL;
+    const uint8_t initPattern = (uint8_t)(initVal->AsIntCon()->IconValue() & 0xFF);
 
     if (initPattern != 0)
     {
         for (unsigned i = 0; i < destLclVar->lvFieldCnt; ++i)
         {
             LclVarDsc* fieldDesc = m_comp->lvaGetDesc(destLclVar->lvFieldLclStart + i);
 
-            if (varTypeIsSIMD(fieldDesc) || varTypeIsGC(fieldDesc))
+            if (varTypeIsGC(fieldDesc))
             {
-                // Cannot initialize GC or SIMD types with a non-zero constant.
-                // The former is completely bogus. The later restriction could be
-                // lifted by supporting non-zero SIMD constants or by generating
-                // field initialization code that converts an integer constant to
-                // the appropriate SIMD value. Unlikely to be very useful, though.
-                JITDUMP(" dest contains GC and/or SIMD fields and source constant is not 0.\n");
+                // Cannot initialize GC types with a non-zero constant. The
+                // former is completely bogus.
+                JITDUMP(" dest contains GC fields and source constant is not 0.\n");
                 return;
             }
         }
@@ -448,58 +445,7 @@ void MorphInitBlockHelper::TryInitFieldByField()
         LclVarDsc* fieldDesc = m_comp->lvaGetDesc(fieldLclNum);
         var_types  fieldType = fieldDesc->TypeGet();
 
-        GenTree* src;
-        switch (fieldType)
-        {
-            case TYP_BOOL:
-            case TYP_BYTE:
-            case TYP_UBYTE:
-            case TYP_SHORT:
-            case TYP_USHORT:
-                // Promoted fields are expected to be "normalize on load". If that changes then
-                // we may need to adjust this code to widen the constant correctly.
-                assert(fieldDesc->lvNormalizeOnLoad());
-                FALLTHROUGH;
-            case TYP_INT:
-            {
-                int64_t mask = (int64_t(1) << (genTypeSize(fieldType) * 8)) - 1;
-                src          = m_comp->gtNewIconNode(static_cast<int32_t>(initPattern & mask));
-                break;
-            }
-            case TYP_LONG:
-                src = m_comp->gtNewLconNode(initPattern);
-                break;
-            case TYP_FLOAT:
-                float floatPattern;
-                memcpy(&floatPattern, &initPattern, sizeof(floatPattern));
-                src = m_comp->gtNewDconNode(floatPattern, TYP_FLOAT);
-                break;
-            case TYP_DOUBLE:
-                double doublePattern;
-                memcpy(&doublePattern, &initPattern, sizeof(doublePattern));
-                src = m_comp->gtNewDconNode(doublePattern);
-                break;
-            case TYP_REF:
-            case TYP_BYREF:
-#ifdef FEATURE_SIMD
-            case TYP_SIMD8:
-            case TYP_SIMD12:
-            case TYP_SIMD16:
-#if defined(TARGET_XARCH)
-            case TYP_SIMD32:
-            case TYP_SIMD64:
-#endif // TARGET_XARCH
-#endif // FEATURE_SIMD
-            {
-                assert(initPattern == 0);
-                src = m_comp->gtNewZeroConNode(fieldType);
-                break;
-            }
-
-            default:
-                unreached();
-        }
-
+        GenTree* src   = m_comp->gtNewConWithPattern(fieldType, initPattern);
         GenTree* store = m_comp->gtNewTempAssign(fieldLclNum, src);
 
         if (m_comp->optLocalAssertionProp)

src/coreclr/jit/promotion.cpp

@@ -627,7 +627,7 @@ bool Replacement::Overlaps(unsigned otherStart, unsigned otherSize) const
 //       LCL_VAR int V01
 //
 // Parameters:
-//   compiler     - Compiler instance
+//   compiler - Compiler instance
 //   structLclNum - Struct local
 //   replacement  - Information about the replacement
 //
@@ -651,7 +651,7 @@ GenTree* Promotion::CreateWriteBack(Compiler* compiler, unsigned structLclNum, c
 //       LCL_FLD int V00 [+4]
 //
 // Parameters:
-//   compiler     - Compiler instance
+//   compiler - Compiler instance
 //   structLclNum - Struct local
 //   replacement  - Information about the replacement
 //

src/coreclr/jit/promotion.h

@@ -20,8 +20,6 @@ struct Replacement
     // a basic block, i.e. all predecessors would have read the replacement
     // back before transferring control if necessary.
     bool NeedsReadBack = false;
-    // Arbitrary flag bit used e.g. by decomposition. Assumed to be false.
-    bool Handled = false;
 #ifdef DEBUG
     const char* Description;
 #endif
@@ -46,6 +44,8 @@ class Promotion
     friend class LocalUses;
     friend class LocalsUseVisitor;
     friend class ReplaceVisitor;
+    friend class DecompositionPlan;
+    friend class StructSegments;
 
     void InsertInitialReadBack(unsigned lclNum, const jitstd::vector<Replacement>& replacements, Statement** prevStmt);
     void ExplicitlyZeroInitReplacementLocals(unsigned                           lclNum,
@@ -107,6 +107,7 @@ class Promotion
 };
 
 class DecompositionStatementList;
+class DecompositionPlan;
 
 class ReplaceVisitor : public GenTreeVisitor<ReplaceVisitor>
 {
@@ -150,17 +151,15 @@ class ReplaceVisitor : public GenTreeVisitor<ReplaceVisitor>
                                  Replacement**        firstReplacement,
                                  Replacement**        endReplacement = nullptr);
     void EliminateCommasInBlockOp(GenTreeOp* asg, DecompositionStatementList* result);
-    void UpdateEarlyRefCount(GenTree* candidate);
-    void IncrementRefCount(unsigned lclNum);
-    void InitFieldByField(Replacement*                firstRep,
-                          Replacement*                endRep,
-                          unsigned char               initVal,
-                          DecompositionStatementList* result);
-    void CopyIntoFields(Replacement*                firstRep,
-                        Replacement*                endRep,
-                        GenTreeLclVarCommon*        dst,
-                        GenTree*                    src,
-                        DecompositionStatementList* result);
+    void InitFields(GenTreeLclVarCommon* dst, Replacement* firstRep, Replacement* endRep, DecompositionPlan* plan);
+    void CopyBetweenFields(GenTree*                    dst,
+                           Replacement*                dstFirstRep,
+                           Replacement*                dstEndRep,
+                           GenTree*                    src,
+                           Replacement*                srcFirstRep,
+                           Replacement*                srcEndRep,
+                           DecompositionStatementList* statements,
+                           DecompositionPlan*          plan);
 };
 
 #endif