Decouple call store operands from local ret buf optimization (#68469)

* Move GenTreeCall::GetLclRetBufArgNode ->
  Compiler::gtCallGetDefinedRetBufLclAddr and change it to not try to
  look into stores
* Assert that the node we are returning actually defines a local that
  has lvHiddenBufferStructArg set
* Assert that call morphing does not break our recognition of the defined
  local when optimizing
* Update Compiler::DefinesLocalAddr to check for GT_LCL_FLD_ADDR

Author: jakobbotsch

src/coreclr/jit/compiler.h

@@ -2790,6 +2790,8 @@ class Compiler
     // Check if this tree is a gc static base helper call
     bool gtIsStaticGCBaseHelperCall(GenTree* tree);
 
+    GenTree* gtCallGetDefinedRetBufLclAddr(GenTreeCall* call);
+
 //-------------------------------------------------------------------------
 // Functions to display the trees
 

src/coreclr/jit/gentree.cpp

@@ -6273,7 +6273,7 @@ bool GenTree::OperRequiresAsgFlag()
     {
         // If the call has return buffer argument, it produced a definition and hence
         // should be marked with assignment.
-        return AsCall()->GetLclRetBufArgNode() != nullptr;
+        return AsCall()->IsOptimizingRetBufAsLocal();
     }
     return false;
 }
@@ -16138,7 +16138,7 @@ bool GenTree::DefinesLocal(Compiler* comp, GenTreeLclVarCommon** pLclVarTree, bo
     }
     else if (OperIs(GT_CALL))
     {
-        GenTree* retBufArg = AsCall()->GetLclRetBufArgNode();
+        GenTree* retBufArg = comp->gtCallGetDefinedRetBufLclAddr(AsCall());
         if (retBufArg == nullptr)
         {
             return false;
@@ -16180,7 +16180,7 @@ bool GenTree::DefinesLocal(Compiler* comp, GenTreeLclVarCommon** pLclVarTree, bo
 // Returns true if this GenTree defines a result which is based on the address of a local.
 bool GenTree::DefinesLocalAddr(Compiler* comp, unsigned width, GenTreeLclVarCommon** pLclVarTree, bool* pIsEntire)
 {
-    if (OperGet() == GT_ADDR || OperGet() == GT_LCL_VAR_ADDR)
+    if (OperIs(GT_ADDR, GT_LCL_VAR_ADDR, GT_LCL_FLD_ADDR))
     {
         GenTree* addrArg = this;
         if (OperGet() == GT_ADDR)
@@ -17864,6 +17864,57 @@ bool Compiler::gtIsStaticGCBaseHelperCall(GenTree* tree)
     return false;
 }
 
+//------------------------------------------------------------------------
+// gtCallGetDefinedRetBufLclAddr:
+//   Get the tree corresponding to the address of the retbuf taht this call defines.
+//
+// Parameters:
+//   call - The call node
+//
+// Returns:
+//   A tree representing the address of a local.
+//
+// Remarks:
+//   This function should not be used until after morph when local address
+//   nodes have been normalized. However, before that IsOptimizingRetBufAsLocal
+//   can be used to at least check if the call has a retbuf that we are
+//   optimizing.
+//
+GenTree* Compiler::gtCallGetDefinedRetBufLclAddr(GenTreeCall* call)
+{
+    if (!call->IsOptimizingRetBufAsLocal())
+    {
+        return nullptr;
+    }
+
+    CallArg* retBufArg = call->gtArgs.GetRetBufferArg();
+    assert(retBufArg != nullptr);
+
+    GenTree* node = retBufArg->GetNode();
+    switch (node->OperGet())
+    {
+        // Get the value from putarg wrapper nodes
+        case GT_PUTARG_REG:
+        case GT_PUTARG_STK:
+            node = node->AsOp()->gtGetOp1();
+            break;
+
+        default:
+            break;
+    }
+
+    // This may be called very late to check validity of LIR.
+    node = node->gtSkipReloadOrCopy();
+
+#ifdef DEBUG
+    unsigned             size = typGetObjLayout(call->gtRetClsHnd)->GetSize();
+    GenTreeLclVarCommon* lcl;
+    assert(node->DefinesLocalAddr(this, size, &lcl, nullptr) && lvaGetDesc(lcl)->lvHiddenBufferStructArg);
+#endif
+
+    return node;
+}
+
 //------------------------------------------------------------------------
 // ParseArrayAddress: Rehydrate the array and index expression from ARR_ADDR.
 //

src/coreclr/jit/gentree.h

@@ -5035,6 +5035,11 @@ struct GenTreeCall final : public GenTree
         return (gtCallMoreFlags & GTF_CALL_M_EXPANDED_EARLY) != 0;
     }
 
+    bool IsOptimizingRetBufAsLocal()
+    {
+        return (gtCallMoreFlags & GTF_CALL_M_RETBUFFARG_LCLOPT) != 0;
+    }
+
     //-----------------------------------------------------------------------------------------
     // GetIndirectionCellArgKind: Get the kind of indirection cell used by this call.
     //
@@ -5181,39 +5186,6 @@ struct GenTreeCall final : public GenTree
     {
     }
 #endif
-
-    GenTree* GetLclRetBufArgNode()
-    {
-        if (!gtArgs.HasRetBuffer() || ((gtCallMoreFlags & GTF_CALL_M_RETBUFFARG_LCLOPT) == 0))
-        {
-            return nullptr;
-        }
-
-        CallArg* retBufArg        = gtArgs.GetRetBufferArg();
-        GenTree* lclRetBufArgNode = retBufArg->GetEarlyNode();
-        if (lclRetBufArgNode == nullptr)
-        {
-            lclRetBufArgNode = retBufArg->GetLateNode();
-        }
-
-        switch (lclRetBufArgNode->OperGet())
-        {
-            // Get the true value from setup args
-            case GT_ASG:
-                return lclRetBufArgNode->AsOp()->gtGetOp2();
-            case GT_STORE_LCL_VAR:
-                return lclRetBufArgNode->AsUnOp()->gtGetOp1();
-
-            // Get the value from putarg wrapper nodes
-            case GT_PUTARG_REG:
-            case GT_PUTARG_STK:
-                return lclRetBufArgNode->AsOp()->gtGetOp1();
-
-            // Otherwise the node should be in the Use*
-            default:
-                return lclRetBufArgNode;
-        }
-    }
 };
 
 struct GenTreeCmpXchg : public GenTree

src/coreclr/jit/lclmorph.cpp

@@ -723,8 +723,8 @@ class LocalAddressVisitor final : public GenTreeVisitor<LocalAddressVisitor>
 
         // TODO-ADDR: For now use LCL_VAR_ADDR and LCL_FLD_ADDR only as call arguments and assignment sources.
         // Other usages require more changes. For example, a tree like OBJ(ADD(ADDR(LCL_VAR), 4))
-        // could be changed to OBJ(LCL_FLD_ADDR) but then DefinesLocalAddr does not recognize
-        // LCL_FLD_ADDR (even though it does recognize LCL_VAR_ADDR).
+        // could be changed to OBJ(LCL_FLD_ADDR) but historically DefinesLocalAddr did not recognize
+        // LCL_FLD_ADDR (and there may be other things now as well).
         if (user->OperIs(GT_CALL, GT_ASG) && !hasHiddenStructArg)
         {
             MorphLocalAddress(val);

src/coreclr/jit/morph.cpp

@@ -1782,6 +1782,10 @@ void CallArgs::EvalArgsToTemps(Compiler* comp, GenTreeCall* call)
         lateTail  = &arg.LateNextRef();
     }
 
+    // Make sure we did not do anything here that stops us from being able to
+    // find the local ret buf if we are optimizing it.
+    noway_assert(!call->IsOptimizingRetBufAsLocal() || (comp->gtCallGetDefinedRetBufLclAddr(call) != nullptr));
+
 #ifdef DEBUG
     if (comp->verbose)
     {

src/coreclr/jit/optimizer.cpp

@@ -5533,7 +5533,7 @@ Compiler::fgWalkResult Compiler::optIsVarAssgCB(GenTree** pTree, fgWalkData* dat
         {
             desc->ivaMaskCall = optCallInterf(tree->AsCall());
 
-            dest = tree->AsCall()->GetLclRetBufArgNode();
+            dest = data->compiler->gtCallGetDefinedRetBufLclAddr(tree->AsCall());
             if (dest == nullptr)
             {
                 return WALK_CONTINUE;

src/coreclr/jit/sideeffects.cpp

@@ -141,15 +141,9 @@ AliasSet::NodeInfo::NodeInfo(Compiler* compiler, GenTree* node)
     if (node->IsCall())
     {
         // For calls having return buffer, update the local number that is written after this call.
-        GenTree* retBufArgNode = node->AsCall()->GetLclRetBufArgNode();
+        GenTree* retBufArgNode = compiler->gtCallGetDefinedRetBufLclAddr(node->AsCall());
         if (retBufArgNode != nullptr)
         {
-            // If a copy/reload is inserted by LSRA, retrieve the returnBuffer
-            if (retBufArgNode->IsCopyOrReload())
-            {
-                retBufArgNode = retBufArgNode->AsCopyOrReload()->gtGetOp1();
-            }
-
             m_flags |= ALIAS_WRITES_LCL_VAR;
             m_lclNum  = retBufArgNode->AsLclVarCommon()->GetLclNum();
             m_lclOffs = retBufArgNode->AsLclVarCommon()->GetLclOffs();