JIT: Check for call interference when placing PUTARG nodes (#103301)

This adds a legalization pass for call argument `PUTARG` nodes after they have been created. We currently insert `PUTARG_*` nodes right after argument definitions, which is usually a good placement because morph ensured that this was ok. However, due to LIR transformations this placement might not be legal if there are any calls after the node. This PR moves the PUTARG_* nodes ahead of those interfering calls when they exist. For example, for the case in #102919 we now end up with: ``` Call [000278] has 7 PUTARG nodes that interfere with [002686]; will move them after it Relocating [002700] after [002686] Relocating [002699] after [002686] Relocating [002698] after [002686] Relocating [002697] after [002686] Relocating [002696] after [002686] Relocating [002695] after [002686] Relocating [002694] after [002686] Final result after legalization: N001 ( 1, 1) [000263] ----------- t263 = LCL_VAR ref V65 tmp36 u:4 $441 N002 ( 1, 1) [000264] ----------- t264 = LCL_VAR ref V76 tmp47 u:4 $446 N003 ( 1, 1) [000265] ----------- t265 = LCL_VAR ref V07 loc3 u:4 $448 N004 ( 1, 1) [000266] ----------- t266 = LCL_VAR int V08 loc4 u:2 $371 N005 ( 3, 2) [000267] ----------- t267 = LCL_VAR ref V09 loc5 u:6 $44c N006 ( 1, 1) [000268] ----------- t268 = LCL_VAR ref V10 loc6 u:6 $44d N007 ( 3, 2) [000269] ----------- t269 = LCL_VAR int V11 loc7 u:4 $780 N008 ( 3, 2) [001733] ----------- t1733 = LCL_VAR ref V125 tmp96 u:2 (last use) $755 N009 ( 1, 1) [001740] -c--------- t1740 = CNS_INT long 16 $1c3 ┌──▌ t1733 ref ├──▌ t1740 long N010 ( 3, 3) [001741] -------N--- t1741 = ▌ ADD byref $3d4 ┌──▌ t1741 byref [002690] DA--------- ▌ STORE_LCL_VAR byref V162 rat3 [002691] ----------- t2691 = LCL_VAR byref V162 rat3 [002692] ----------- t2692 = LCL_VAR byref V162 rat3 ┌──▌ t2692 byref [002693] ---X------- ▌ NULLCHECK byte [002681] D---------- t2681 = LCL_ADDR byref V141 tmp112 [+0] ┌──▌ t2681 byref [002687] ----------- t2687 = ▌ PUTARG_REG byref REG rcx ┌──▌ t2691 byref [002688] ----------- t2688 = ▌ PUTARG_REG byref REG rdx [002682] ----------- t2682 = CNS_INT long 128 ┌──▌ t2682 long [002689] ----------- t2689 = ▌ PUTARG_REG long REG r8 ┌──▌ t2687 byref arg0 in rcx ├──▌ t2688 byref arg1 in rdx ├──▌ t2689 long arg2 in r8 N004 ( 17, 11) [002686] --CXG------ ▌ CALL help void CORINFO_HELP_BULK_WRITEBARRIER ┌──▌ t263 ref [002694] ----------- ▌ PUTARG_STK [+0x20] void ┌──▌ t264 ref [002695] ----------- ▌ PUTARG_STK [+0x28] void ┌──▌ t265 ref [002696] ----------- ▌ PUTARG_STK [+0x30] void ┌──▌ t266 int [002697] ----------- ▌ PUTARG_STK [+0x38] void ┌──▌ t267 ref [002698] ----------- ▌ PUTARG_STK [+0x40] void ┌──▌ t268 ref [002699] ----------- ▌ PUTARG_STK [+0x48] void ┌──▌ t269 int [002700] ----------- ▌ PUTARG_STK [+0x50] void N013 ( 19, 14) [001744] ----------- NOP void N014 ( 3, 2) [000275] ----------- t275 = LCL_VAR ref V123 tmp94 u:2 $754 ┌──▌ t275 ref [002701] ----------- ▌ PUTARG_STK [+0x60] void N015 ( 3, 3) [000276] ----------- t276 = LCL_ADDR long V22 loc18 [+0] $3d1 ┌──▌ t276 long [002702] ----------- ▌ PUTARG_STK [+0x68] void N016 ( 1, 1) [000277] ----------- t277 = LCL_VAR ref V03 arg3 u:1 $103 ┌──▌ t277 ref [002703] ----------- ▌ PUTARG_STK [+0x70] void N017 ( 3, 3) [001745] ----------- t1745 = LCL_ADDR long V141 tmp112 [+0] $3d8 ┌──▌ t1745 long [002704] ----------- ▌ PUTARG_STK [+0x58] void N018 ( 1, 1) [000261] ----------- t261 = LCL_VAR ref V01 arg1 u:1 $101 ┌──▌ t261 ref [002705] ----------- t2705 = ▌ PUTARG_REG ref REG r8 N019 ( 1, 1) [000262] ----------- t262 = LCL_VAR ref V02 arg2 u:1 $102 ┌──▌ t262 ref [002706] ----------- t2706 = ▌ PUTARG_REG ref REG r9 N020 ( 1, 1) [000259] ----------- t259 = LCL_VAR ref V00 this u:1 $100 ┌──▌ t259 ref [002707] ----------- t2707 = ▌ PUTARG_REG ref REG rcx N021 ( 1, 1) [001124] ----------- t1124 = CNS_INT int 1 $c1 ┌──▌ t1124 int [002708] ----------- t2708 = ▌ PUTARG_REG int REG rdx ┌──▌ t2705 ref arg2 in r8 ├──▌ t2706 ref arg3 in r9 ├──▌ t2707 ref this in rcx ├──▌ t2708 int arg1 in rdx N022 ( 91, 46) [000278] -ACXGO----- t278 = ▌ CALL ref Microsoft.CodeAnalysis.CSharp.Binder:BindConstructorInitializerCoreContinued(ubyte,Microsoft.CodeAnalysis.CSharp.Syntax.ArgumentListSyntax,Microsoft.CodeAnalysis.CSharp.Symbols.MethodSymbol,Microsoft.CodeAnalysis.CSharp.AnalyzedArguments,Microsoft.CodeAnalysis.CSharp.Symbols.TypeSymbol,Microsoft.CodeAnalysis.CSharp.Symbols.NamedTypeSymbol,ubyte,Microsoft.CodeAnalysis.CSharp.CSharpSyntaxNode,Microsoft.CodeAnalysis.Location,ubyte,Microsoft.CodeAnalysis.CSharp.MemberResolutionResult`1[Microsoft.CodeAnalysis.CSharp.Symbols.MethodSymbol],System.Collections.Immutable.ImmutableArray`1[Microsoft.CodeAnalysis.CSharp.Symbols.MethodSymbol],byref,Microsoft.CodeAnalysis.CSharp.BindingDiagnosticBag):Microsoft.CodeAnalysis.CSharp.BoundExpression:this $135 ``` without this change the `PUTARG_STK` nodes end up before the `BULKWRITEBARRIER` call. The algorithm works by first marking all `PUTARG` operand nodes of the call, and then walking backwards from the call in linear order until we have visited all marked nodes. If we see any call before visiting all marked nodes then we have found an interfering call, and the remaining `PUTARG_` nodes need to be moved after it. Fix #103300 Fix #102919 Fix #104042
dotnet · Jun 26, 2024 · e498286 · e498286
1 parent 58e1a7e
commit e498286
Show file tree

Hide file tree

Showing 2 changed files with 167 additions and 26 deletions.
diff --git a/src/coreclr/jit/lower.cpp b/src/coreclr/jit/lower.cpp
@@ -1495,6 +1495,14 @@ bool Lowering::TryLowerSwitchToBitTest(FlowEdge*   jumpTable[],
     return true;
 }
 
+//------------------------------------------------------------------------
+// ReplaceArgWithPutArgOrBitcast: Insert a PUTARG_* node in the right location
+// and replace the call operand with that node.
+//
+// Arguments:
+//    argSlot         - slot in call of argument
+//    putArgOrBitcast - the node that is being inserted
+//
 void Lowering::ReplaceArgWithPutArgOrBitcast(GenTree** argSlot, GenTree* putArgOrBitcast)
 {
     assert(argSlot != nullptr);
@@ -1981,6 +1989,113 @@ void Lowering::LowerArgsForCall(GenTreeCall* call)
     {
         LowerArg(call, &arg, true);
     }
+
+    LegalizeArgPlacement(call);
+}
+
+//------------------------------------------------------------------------
+// LegalizeArgPlacement: Move arg placement nodes (PUTARG_*) into a legal
+// ordering after they have been created.
+//
+// Arguments:
+//   call - GenTreeCall node that has had PUTARG_* nodes created for arguments.
+//
+// Remarks:
+//   PUTARG_* nodes are created and inserted right after the definitions of the
+//   argument values. However, there are constraints on how the PUTARG nodes
+//   can appear:
+//
+//   - No other GT_CALL nodes are allowed between a PUTARG_REG/PUTARG_SPLIT
+//   node and the call. For FEATURE_FIXED_OUT_ARGS this condition is also true
+//   for PUTARG_STK.
+//   - For !FEATURE_FIXED_OUT_ARGS, the PUTARG_STK nodes must come in push
+//   order.
+//
+//   Morph has mostly already solved this problem, but transformations on LIR
+//   can make the ordering we end up with here illegal. This function legalizes
+//   the placement while trying to minimize the distance between an argument
+//   definition and its corresponding placement node.
+//
+void Lowering::LegalizeArgPlacement(GenTreeCall* call)
+{
+    size_t numMarked = MarkCallPutArgAndFieldListNodes(call);
+
+    // We currently do not try to resort the PUTARG_STK nodes, but rather just
+    // assert here that they are ordered.
+#if defined(DEBUG) && !FEATURE_FIXED_OUT_ARGS
+    unsigned nextPushOffset = UINT_MAX;
+#endif
+
+    GenTree* cur = call->gtPrev;
+    while (numMarked > 0)
+    {
+        assert(cur != nullptr);
+
+        if ((cur->gtLIRFlags & LIR::Flags::Mark) != 0)
+        {
+            numMarked--;
+            cur->gtLIRFlags &= ~LIR::Flags::Mark;
+
+#if defined(DEBUG) && !FEATURE_FIXED_OUT_ARGS
+            if (cur->OperIs(GT_PUTARG_STK))
+            {
+                // For !FEATURE_FIXED_OUT_ARGS (only x86) byte offsets are
+                // subtracted from the top of the stack frame; so last pushed
+                // arg has highest offset.
+                assert(nextPushOffset > cur->AsPutArgStk()->getArgOffset());
+                nextPushOffset = cur->AsPutArgStk()->getArgOffset();
+            }
+#endif
+        }
+
+        if (cur->IsCall())
+        {
+            break;
+        }
+
+        cur = cur->gtPrev;
+    }
+
+    if (numMarked == 0)
+    {
+        // Already legal; common case
+        return;
+    }
+
+    JITDUMP("Call [%06u] has %zu PUTARG nodes that interfere with [%06u]; will move them after it\n",
+            Compiler::dspTreeID(call), numMarked, Compiler::dspTreeID(cur));
+
+    // We found interference; remaining PUTARG nodes need to be moved after
+    // this point.
+    GenTree* insertionPoint = cur;
+
+    while (numMarked > 0)
+    {
+        assert(cur != nullptr);
+
+        GenTree* prev = cur->gtPrev;
+        if ((cur->gtLIRFlags & LIR::Flags::Mark) != 0)
+        {
+            numMarked--;
+            cur->gtLIRFlags &= ~LIR::Flags::Mark;
+
+            // For FEATURE_FIXED_OUT_ARGS: all PUTARG nodes must be moved after the interfering call
+            // For !FEATURE_FIXED_OUT_ARGS: only PUTARG_REG nodes must be moved after the interfering call
+            if (FEATURE_FIXED_OUT_ARGS || cur->OperIs(GT_FIELD_LIST, GT_PUTARG_REG))
+            {
+                JITDUMP("Relocating [%06u] after [%06u]\n", Compiler::dspTreeID(cur),
+                        Compiler::dspTreeID(insertionPoint));
+
+                BlockRange().Remove(cur);
+                BlockRange().InsertAfter(insertionPoint, cur);
+            }
+        }
+
+        cur = prev;
+    }
+
+    JITDUMP("Final result after legalization:\n");
+    DISPTREERANGE(BlockRange(), call);
 }
 
 // helper that create a node representing a relocatable physical address computation
@@ -1997,6 +2112,7 @@ GenTree* Lowering::AddrGen(void* addr)
     return AddrGen((ssize_t)addr);
 }
 
+//------------------------------------------------------------------------
 // LowerCallMemset: Replaces the following memset-like special intrinsics:
 //
 //    SpanHelpers.Fill<T>(ref dstRef, CNS_SIZE, CNS_VALUE)
@@ -2780,19 +2896,7 @@ void Lowering::InsertProfTailCallHook(GenTreeCall* call, GenTree* insertionPoint
 //
 GenTree* Lowering::FindEarliestPutArg(GenTreeCall* call)
 {
-    size_t numMarkedNodes = 0;
-    for (CallArg& arg : call->gtArgs.Args())
-    {
-        if (arg.GetEarlyNode() != nullptr)
-        {
-            numMarkedNodes += MarkPutArgNodes(arg.GetEarlyNode());
-        }
-
-        if (arg.GetLateNode() != nullptr)
-        {
-            numMarkedNodes += MarkPutArgNodes(arg.GetLateNode());
-        }
-    }
+    size_t numMarkedNodes = MarkCallPutArgAndFieldListNodes(call);
 
     if (numMarkedNodes <= 0)
     {
@@ -2818,33 +2922,68 @@ GenTree* Lowering::FindEarliestPutArg(GenTreeCall* call)
 }
 
 //------------------------------------------------------------------------
-// MarkPutArgNodes: Mark all direct operand PUTARG nodes with a LIR mark.
+// MarkCallPutArgNodes: Mark all operand FIELD_LIST and PUTARG nodes
+// corresponding to a call.
 //
 // Arguments:
-//    node - the node (either a field list or PUTARG node)
+//   call - the call
 //
 // Returns:
-//    The number of marks added.
+//   The number of nodes marked.
+//
+// Remarks:
+//   FIELD_LIST operands are marked too, and their PUTARG operands are in turn
+//   marked as well.
 //
-size_t Lowering::MarkPutArgNodes(GenTree* node)
+size_t Lowering::MarkCallPutArgAndFieldListNodes(GenTreeCall* call)
+{
+    size_t numMarkedNodes = 0;
+    for (CallArg& arg : call->gtArgs.Args())
+    {
+        if (arg.GetEarlyNode() != nullptr)
+        {
+            numMarkedNodes += MarkPutArgAndFieldListNodes(arg.GetEarlyNode());
+        }
+
+        if (arg.GetLateNode() != nullptr)
+        {
+            numMarkedNodes += MarkPutArgAndFieldListNodes(arg.GetLateNode());
+        }
+    }
+
+    return numMarkedNodes;
+}
+
+//------------------------------------------------------------------------
+// MarkPutArgAndFieldListNodes: Mark all operand FIELD_LIST and PUTARG nodes
+// with a LIR mark.
+//
+// Arguments:
+//   node - the node (either a FIELD_LIST or PUTARG operand)
+//
+// Returns:
+//   The number of marks added.
+//
+// Remarks:
+//   FIELD_LIST operands are marked too, and their PUTARG operands are in turn
+//   marked as well.
+//
+size_t Lowering::MarkPutArgAndFieldListNodes(GenTree* node)
 {
     assert(node->OperIsPutArg() || node->OperIsFieldList());
 
-    size_t result = 0;
+    assert((node->gtLIRFlags & LIR::Flags::Mark) == 0);
+    node->gtLIRFlags |= LIR::Flags::Mark;
+
+    size_t result = 1;
     if (node->OperIsFieldList())
     {
         for (GenTreeFieldList::Use& operand : node->AsFieldList()->Uses())
         {
             assert(operand.GetNode()->OperIsPutArg());
-            result += MarkPutArgNodes(operand.GetNode());
+            result += MarkPutArgAndFieldListNodes(operand.GetNode());
         }
     }
-    else
-    {
-        assert((node->gtLIRFlags & LIR::Flags::Mark) == 0);
-        node->gtLIRFlags |= LIR::Flags::Mark;
-        result++;
-    }
 
     return result;
 }

diff --git a/src/coreclr/jit/lower.h b/src/coreclr/jit/lower.h
@@ -175,7 +175,8 @@ class Lowering final : public Phase
                                       GenTreeCall* callNode);
     void     InsertProfTailCallHook(GenTreeCall* callNode, GenTree* insertionPoint);
     GenTree* FindEarliestPutArg(GenTreeCall* call);
-    size_t   MarkPutArgNodes(GenTree* node);
+    size_t   MarkCallPutArgAndFieldListNodes(GenTreeCall* call);
+    size_t   MarkPutArgAndFieldListNodes(GenTree* node);
     GenTree* LowerVirtualVtableCall(GenTreeCall* call);
     GenTree* LowerVirtualStubCall(GenTreeCall* call);
     void     LowerArgsForCall(GenTreeCall* call);
@@ -186,6 +187,7 @@ class Lowering final : public Phase
     GenTree* LowerFloatArg(GenTree** pArg, CallArg* callArg);
     GenTree* LowerFloatArgReg(GenTree* arg, regNumber regNum);
 #endif
+    void LegalizeArgPlacement(GenTreeCall* call);
 
     void     InsertPInvokeCallProlog(GenTreeCall* call);
     void     InsertPInvokeCallEpilog(GenTreeCall* call);