Handle last uses that become contained

2 files changed, 37 insertions, 1 deletions

diff --git a/src/jit/lsra.h b/src/jit/lsra.h
index d3d3f99b00..22e3ead7ec 100644
--- a/src/jit/lsra.h
+++ b/src/jit/lsra.h
@@ -1683,6 +1683,28 @@ private:
     bool isRMWRegOper(GenTree* tree);
     void BuildMul(GenTree* tree);
     void SetContainsAVXFlags(bool isFloatingPointType = true, unsigned sizeOfSIMDVector = 0);
+    // Move the last use bit, if any, from 'fromTree' to 'toTree'; 'fromTree' must be contained.
+    void CheckAndMoveRMWLastUse(GenTree* fromTree, GenTree* toTree)
+    {
+        // If 'fromTree' is not a last-use lclVar, there's nothing to do.
+        if ((fromTree == nullptr) || !fromTree->OperIs(GT_LCL_VAR) || ((fromTree->gtFlags & GTF_VAR_DEATH) == 0))
+        {
+            return;
+        }
+        // If 'fromTree' was a lclVar, it must be contained and 'toTree' must match.
+        if (!fromTree->isContained() || (toTree == nullptr) || !toTree->OperIs(GT_LCL_VAR) ||
+            (toTree->AsLclVarCommon()->gtLclNum != toTree->AsLclVarCommon()->gtLclNum))
+        {
+            assert(!"Unmatched RMW indirections");
+            return;
+        }
+        // This is probably not necessary, but keeps things consistent.
+        fromTree->gtFlags &= ~GTF_VAR_DEATH;
+        if (toTree != nullptr) // Just to be conservative
+        {
+            toTree->gtFlags |= GTF_VAR_DEATH;
+        }
+    }
 #endif // defined(_TARGET_XARCH_)
 
 #ifdef FEATURE_SIMD
diff --git a/src/jit/lsraxarch.cpp b/src/jit/lsraxarch.cpp
index 09cc9ca96e..f72403d0e3 100644
--- a/src/jit/lsraxarch.cpp
+++ b/src/jit/lsraxarch.cpp
@@ -2515,7 +2515,8 @@ void LinearScan::BuildIndir(GenTreeIndir* indirTree)
             // Because 'source' is contained, we haven't yet determined its special register requirements, if any.
             // As it happens, the Shift or Rotate cases are the only ones with special requirements.
             assert(source->isContained() && source->OperIsRMWMemOp());
-            GenTree* nonMemSource = nullptr;
+            GenTree*      nonMemSource = nullptr;
+            GenTreeIndir* otherIndir   = nullptr;
 
             if (source->OperIsShiftOrRotate())
             {
@@ -2527,6 +2528,7 @@ void LinearScan::BuildIndir(GenTreeIndir* indirTree)
             }
             if (indirTree->AsStoreInd()->IsRMWDstOp1())
             {
+                otherIndir = source->gtGetOp1()->AsIndir();
                 if (source->OperIsBinary())
                 {
                     nonMemSource = source->gtOp.gtOp2;
@@ -2534,8 +2536,20 @@ void LinearScan::BuildIndir(GenTreeIndir* indirTree)
             }
             else if (indirTree->AsStoreInd()->IsRMWDstOp2())
             {
+                otherIndir   = source->gtGetOp2()->AsIndir();
                 nonMemSource = source->gtOp.gtOp1;
             }
+            if (otherIndir != nullptr)
+            {
+                // Any lclVars in the addressing mode of this indirection are contained.
+                // If they are marked as lastUse, transfer the last use flag to the store indir.
+                GenTree* base    = otherIndir->Base();
+                GenTree* dstBase = indirTree->Base();
+                CheckAndMoveRMWLastUse(base, dstBase);
+                GenTree* index    = otherIndir->Index();
+                GenTree* dstIndex = indirTree->Index();
+                CheckAndMoveRMWLastUse(index, dstIndex);
+            }
             if (nonMemSource != nullptr)
             {
                 assert(!nonMemSource->isContained() || (!nonMemSource->isMemoryOp() && !nonMemSource->IsLocal()));