Introduce ASG rationalization phase

Author: SingleAccretion

src/coreclr/jit/compiler.cpp

@@ -1919,6 +1919,7 @@ void Compiler::compInit(ArenaAllocator*       pAlloc,
     compLocallocUsed             = false;
     compLocallocOptimized        = false;
     compQmarkRationalized        = false;
+    compAssignmentRationalized   = false;
     compQmarkUsed                = false;
     compFloatingPointUsed        = false;
 
@@ -5064,6 +5065,8 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     //
     DoPhase(this, PHASE_DETERMINE_FIRST_COLD_BLOCK, &Compiler::fgDetermineFirstColdBlock);
 
+    DoPhase(this, PHASE_RATIONALIZE_ASSIGNMENTS, &Compiler::fgRationalizeAssignments);
+
 #ifdef DEBUG
     // Stash the current estimate of the function's size if necessary.
     if (verbose)

src/coreclr/jit/compiler.h

@@ -4830,6 +4830,9 @@ class Compiler
     void fgExpandQmarkStmt(BasicBlock* block, Statement* stmt);
     void fgExpandQmarkNodes();
 
+    PhaseStatus fgRationalizeAssignments();
+    GenTree* fgRationalizeAssignment(GenTreeOp* assignment);
+
     // Do "simple lowering."  This functionality is (conceptually) part of "general"
     // lowering that is distributed between fgMorph and the lowering phase of LSRA.
     PhaseStatus fgSimpleLowering();
@@ -9232,6 +9235,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
     bool compLocallocOptimized;        // Does the method have an optimized localloc
     bool compQmarkUsed;                // Does the method use GT_QMARK/GT_COLON
     bool compQmarkRationalized;        // Is it allowed to use a GT_QMARK/GT_COLON node.
+    bool compAssignmentRationalized;   // Have the ASG nodes been turned into their store equivalents?
     bool compHasBackwardJump;          // Does the method (or some inlinee) have a lexically backwards jump?
     bool compHasBackwardJumpInHandler; // Does the method have a lexically backwards jump in a handler?
     bool compSwitchedToOptimized;      // Codegen initially was Tier0 but jit switched to FullOpts

src/coreclr/jit/compphases.h

@@ -91,6 +91,7 @@ CompPhaseNameMacro(PHASE_IF_CONVERSION,              "If conversion",
 CompPhaseNameMacro(PHASE_VN_BASED_DEAD_STORE_REMOVAL,"VN-based dead store removal",    false, -1, false)
 CompPhaseNameMacro(PHASE_OPT_UPDATE_FLOW_GRAPH,      "Update flow graph opt pass",     false, -1, false)
 CompPhaseNameMacro(PHASE_COMPUTE_EDGE_WEIGHTS2,      "Compute edge weights (2, false)",false, -1, false)
+CompPhaseNameMacro(PHASE_RATIONALIZE_ASSIGNMENTS,    "Rationalize assignments",        false, -1, false)
 CompPhaseNameMacro(PHASE_STRESS_SPLIT_TREE,          "Stress gtSplitTree",             false, -1, false)
 CompPhaseNameMacro(PHASE_EXPAND_RTLOOKUPS,           "Expand runtime lookups",         false, -1, true)
 CompPhaseNameMacro(PHASE_EXPAND_STATIC_INIT,         "Expand static init",             false, -1, true)

src/coreclr/jit/flowgraph.cpp

@@ -2841,6 +2841,175 @@ PhaseStatus Compiler::fgFindOperOrder()
     return PhaseStatus::MODIFIED_EVERYTHING;
 }
 
+//------------------------------------------------------------------------
+// fgRationalizeAssignments: Rewrite assignment nodes into stores.
+//
+// TODO-ASG: delete.
+//
+PhaseStatus Compiler::fgRationalizeAssignments()
+{
+    class AssignmentRationalizationVisitor : public GenTreeVisitor<AssignmentRationalizationVisitor>
+    {
+    public:
+        enum
+        {
+            DoPreOrder = true
+        };
+
+        AssignmentRationalizationVisitor(Compiler* compiler) : GenTreeVisitor(compiler)
+        {
+        }
+
+        fgWalkResult PreOrderVisit(GenTree** use, GenTree* user)
+        {
+            GenTree* node = *use;
+
+            // GTF_ASG is sometimes not propagated from setup arg assignments so we have to check for GTF_CALL too.
+            if ((node->gtFlags & (GTF_ASG | GTF_CALL)) == 0)
+            {
+                return fgWalkResult::WALK_SKIP_SUBTREES;
+            }
+
+            if (node->OperIs(GT_ASG))
+            {
+                GenTreeFlags lhsRhsFlags = node->gtGetOp1()->gtFlags | node->gtGetOp2()->gtFlags;
+                *use                     = m_compiler->fgRationalizeAssignment(node->AsOp());
+
+                // TP: return early quickly for simple assignments.
+                if ((lhsRhsFlags & (GTF_ASG | GTF_CALL)) == 0)
+                {
+                    return fgWalkResult::WALK_SKIP_SUBTREES;
+                }
+            }
+
+            return fgWalkResult::WALK_CONTINUE;
+        }
+    };
+
+    AssignmentRationalizationVisitor visitor(this);
+    for (BasicBlock* block : Blocks())
+    {
+        for (Statement* stmt : block->Statements())
+        {
+            GenTree** use = stmt->GetRootNodePointer();
+            if (visitor.PreOrderVisit(use, nullptr) == fgWalkResult::WALK_CONTINUE)
+            {
+                visitor.WalkTree(use, nullptr);
+            }
+        }
+    }
+
+    compAssignmentRationalized = true;
+
+    return PhaseStatus::MODIFIED_EVERYTHING;
+}
+
+//------------------------------------------------------------------------
+// fgRationalizeAssignment: Rewrite GT_ASG into a store node.
+//
+// Arguments:
+//    assignment - The assignment node to rewrite
+//
+// Return Value:
+//    Assignment's location, turned into the appropriate store node.
+//
+GenTree* Compiler::fgRationalizeAssignment(GenTreeOp* assignment)
+{
+    assert(assignment->OperGet() == GT_ASG);
+
+    bool     isReverseOp = assignment->IsReverseOp();
+    GenTree* location    = assignment->gtGetOp1();
+    GenTree* value       = assignment->gtGetOp2();
+    if (location->OperIsLocal())
+    {
+        assert((location->gtFlags & GTF_VAR_DEF) != 0);
+    }
+    else if (value->OperIs(GT_LCL_VAR))
+    {
+        assert((value->gtFlags & GTF_VAR_DEF) == 0);
+    }
+
+    if (assignment->OperIsInitBlkOp())
+    {
+        // No SIMD types are allowed for InitBlks (including zero-inits).
+        assert(assignment->TypeIs(TYP_STRUCT) && location->TypeIs(TYP_STRUCT));
+    }
+
+    genTreeOps storeOp;
+    switch (location->OperGet())
+    {
+        case GT_LCL_VAR:
+            storeOp = GT_STORE_LCL_VAR;
+            break;
+        case GT_LCL_FLD:
+            storeOp = GT_STORE_LCL_FLD;
+            break;
+        case GT_BLK:
+            storeOp = GT_STORE_BLK;
+            break;
+        case GT_IND:
+            storeOp = GT_STOREIND;
+            break;
+        default:
+            unreached();
+    }
+
+    JITDUMP("Rewriting GT_ASG(%s, X) to %s(X)\n", GenTree::OpName(location->OperGet()), GenTree::OpName(storeOp));
+
+    GenTree* store = location;
+    store->SetOperRaw(storeOp);
+    store->Data() = value;
+    store->gtFlags |= GTF_ASG;
+    store->AddAllEffectsFlags(value);
+    store->AddAllEffectsFlags(assignment->gtFlags & GTF_GLOB_REF); // TODO-ASG: zero-diff quirk, delete.
+    if (isReverseOp && !store->OperIsLocalStore())
+    {
+        store->SetReverseOp();
+    }
+
+    if (storeOp == GT_STOREIND)
+    {
+        store->AsStoreInd()->SetRMWStatusDefault();
+    }
+
+    // [..., LHS, ..., RHS, ASG] -> [..., ..., RHS, LHS<STORE>] (normal)
+    // [..., RHS, ..., LHS, ASG] -> [..., RHS, ..., LHS<STORE>] (reversed)
+    if (assignment->gtPrev != nullptr)
+    {
+        assert(fgNodeThreading == NodeThreading::AllTrees);
+        if (isReverseOp)
+        {
+            GenTree* nextNode = assignment->gtNext;
+            store->gtNext     = nextNode;
+            if (nextNode != nullptr)
+            {
+                nextNode->gtPrev = store;
+            }
+        }
+        else
+        {
+            if (store->gtPrev != nullptr)
+            {
+                store->gtPrev->gtNext = store->gtNext;
+            }
+            store->gtNext->gtPrev = store->gtPrev;
+
+            store->gtPrev         = assignment->gtPrev;
+            store->gtNext         = assignment->gtNext;
+            store->gtPrev->gtNext = store;
+            if (store->gtNext != nullptr)
+            {
+                store->gtNext->gtPrev = store;
+            }
+        }
+    }
+
+    DISPNODE(store);
+    JITDUMP("\n");
+
+    return store;
+}
+
 //------------------------------------------------------------------------
 // fgSimpleLowering: do full walk of all IR, lowering selected operations
 // and computing lvaOutgoingArgSpaceSize.

src/coreclr/jit/gentree.cpp

@@ -15972,23 +15972,18 @@ GenTree* Compiler::gtNewTempAssign(
     GenTree* asg;
     GenTree* dest = gtNewLclvNode(tmp, dstTyp);
 
-    if (val->IsInitVal())
-    {
-        asg = gtNewAssignNode(dest, val);
-    }
-    else if (varTypeIsStruct(varDsc))
+    if (varTypeIsStruct(varDsc) && !val->IsInitVal())
     {
         asg = impAssignStruct(dest, val, CHECK_SPILL_NONE, pAfterStmt, di, block);
     }
     else
     {
-        assert(!varTypeIsStruct(valTyp));
         asg = gtNewAssignNode(dest, val);
     }
 
-    if (compRationalIRForm)
+    if (compAssignmentRationalized)
     {
-        Rationalizer::RewriteAssignmentIntoStoreLcl(asg->AsOp());
+        asg = fgRationalizeAssignment(asg->AsOp());
     }
 
     return asg;
@@ -17242,9 +17237,15 @@ bool GenTree::IsPhiNode()
 
 bool GenTree::IsPhiDefn()
 {
-    bool res = OperIs(GT_ASG) && AsOp()->gtOp2->OperIs(GT_PHI);
-    assert(!res || AsOp()->gtOp1->OperIs(GT_LCL_VAR));
-    return res;
+    if (OperIs(GT_ASG))
+    {
+        return AsOp()->gtOp2->OperIs(GT_PHI);
+    }
+    if (OperIs(GT_STORE_LCL_VAR))
+    {
+        return AsLclVar()->Data()->OperIs(GT_PHI);
+    }
+    return false;
 }
 
 bool GenTree::IsLclVarAddr() const