JIT: Support bitwise field insertions for call arguments (dotnet#115977)

Add support using bitwise operations to reconstruct registers passed
into calls from multiple promoted fields. Remove the IR invariant that
`FIELD_LIST` args must always map cleanly to registers; instead, any
`FIELD_LIST` is allowed for register-only arguments before lowering, and
lowering takes care to normalize them into a handled shape.

`fgTryMorphStructArg` is changed to take advantage of this by now
producing `FIELD_LIST` even when a promoted arg does not match the
target ABI. Support in physical promotion will be added in a follow up.

Split arguments are not handled and retain the old IR invariant of
requiring registers and stack slots to make cleanly from `FIELD_LIST`.

win-x64 examples:
```csharp
static void Foo(int x)
{
    Use<int?>(x);
    Use<int?>(5);
    Use<int?>(null);
}
```

```diff
 G_M7200_IG02:  ;; offset=0x0004
-       mov      byte  ptr [rsp+0x20], 1
-       mov      dword ptr [rsp+0x24], ecx
-       mov      rcx, qword ptr [rsp+0x20]
+       mov      ecx, ecx
+       shl      rcx, 32
+       or       rcx, 1
        call     [Program:Bar(System.Nullable`1[int])]
-       mov      dword ptr [rsp+0x24], 5
-       mov      rcx, qword ptr [rsp+0x20]
+       mov      rcx, 0x500000001
        call     [Program:Bar(System.Nullable`1[int])]
-       mov      byte  ptr [rsp+0x20], 0
        xor      ecx, ecx
-       mov      dword ptr [rsp+0x24], ecx
-       mov      rcx, qword ptr [rsp+0x20]
-						;; size=55 bbWeight=1 PerfScore 14.25
+						;; size=34 bbWeight=1 PerfScore 7.50

 G_M7200_IG03:
        add      rsp, 40
        tail.jmp [Program:Bar(System.Nullable`1[int])]
 						;; size=10 bbWeight=1 PerfScore 2.25
```

```csharp
static void Foo(int x, float y)
{
    Use((x, y));
}
```

```diff
 G_M42652_IG01:  ;; offset=0x0000
-       push     rax
-						;; size=1 bbWeight=1 PerfScore 1.00
+						;; size=0 bbWeight=1 PerfScore 0.00

 G_M42652_IG02:
-       mov      dword ptr [rsp], ecx
-       vmovss   dword ptr [rsp+0x04], xmm1
-       mov      rcx, qword ptr [rsp]
+       vmovd    eax, xmm1
+       shl      rax, 32
+       mov      ecx, ecx
+       or       rcx, rax
 						;; size=13 bbWeight=1 PerfScore 3.00

 G_M42652_IG03:
-       add      rsp, 8
        tail.jmp [Program:Use[System.ValueTuple`2[int,float]](System.ValueTuple`2[int,float])]
```

A win-arm64 example:
```csharp
static void Foo(int[] arr)
{
    Use(arr.AsMemory());
}
```

```diff
 G_M33990_IG01:
-            stp     fp, lr, [sp, #-0x20]!
+            stp     fp, lr, [sp, #-0x10]!
             mov     fp, sp
-            str     xzr, [fp, #0x10]	// [V03 tmp2]
-						;; size=12 bbWeight=1 PerfScore 2.50
+						;; size=8 bbWeight=1 PerfScore 1.50

 G_M33990_IG02:
             cbz     x0, G_M33990_IG04
 						;; size=4 bbWeight=1 PerfScore 1.00

 G_M33990_IG03:
-            str     x0, [fp, #0x10]	// [V07 tmp6]
-            str     wzr, [fp, #0x18]	// [V08 tmp7]
-            ldr     w0, [x0, #0x08]
-            str     w0, [fp, #0x1C]	// [V09 tmp8]
+            ldr     w1, [x0, #0x08]
             b       G_M33990_IG05
-						;; size=20 bbWeight=0.50 PerfScore 3.50
+						;; size=8 bbWeight=0.50 PerfScore 2.00

 G_M33990_IG04:
-            str     xzr, [fp, #0x10]	// [V07 tmp6]
-            str     xzr, [fp, #0x18]
-						;; size=8 bbWeight=0.50 PerfScore 1.00
+            mov     x0, xzr
+            mov     w1, wzr
+						;; size=8 bbWeight=0.50 PerfScore 0.50

 G_M33990_IG05:
-            ldp     x0, x1, [fp, #0x10]	// [V03 tmp2], [V03 tmp2+0x08]
-            movz    x2, #0xD920      // code for Program:Use[System.Memory`1[int]](System.Memory`1[int])
-            movk    x2, #0x4590 LSL #16
+            mov     w1, w1
+            lsl     x1, x1, #32
+            movz    x2, #0xD950      // code for Program:Use[System.Memory`1[int]](System.Memory`1[int])
+            movk    x2, #0x4592 LSL #16
             movk    x2, #0x7FFE LSL #32
             ldr     x2, [x2]
-						;; size=20 bbWeight=1 PerfScore 7.50
+						;; size=24 bbWeight=1 PerfScore 6.00

 G_M33990_IG06:
-            ldp     fp, lr, [sp], #0x20
+            ldp     fp, lr, [sp], #0x10
             br      x2
 						;; size=8 bbWeight=1 PerfScore 2.00
-; Total bytes of code: 72
+; Total bytes of code: 60
```

Author: jakobbotsch

src/coreclr/jit/abi.cpp

@@ -154,6 +154,29 @@ var_types ABIPassingSegment::GetRegisterType() const
     }
 }
 
+//-----------------------------------------------------------------------------
+// GetRegisterType:
+//  Return the smallest type larger or equal to Size that most naturally
+//  represents the register this segment is passed in, taking into account the
+//  GC info of the specified layout.
+//
+// Return Value:
+//   A type that matches ABIPassingSegment::Size and the register.
+//
+var_types ABIPassingSegment::GetRegisterType(ClassLayout* layout) const
+{
+    if (genIsValidIntReg(GetRegister()))
+    {
+        assert(Offset < layout->GetSize());
+        if (((Offset % TARGET_POINTER_SIZE) == 0) && (Size == TARGET_POINTER_SIZE))
+        {
+            return layout->GetGCPtrType(Offset / TARGET_POINTER_SIZE);
+        }
+    }
+
+    return GetRegisterType();
+}
+
 //-----------------------------------------------------------------------------
 // InRegister:
 //   Create an ABIPassingSegment representing that a segment is passed in a

src/coreclr/jit/abi.h

@@ -41,6 +41,7 @@ class ABIPassingSegment
     unsigned GetStackSize() const;
 
     var_types GetRegisterType() const;
+    var_types GetRegisterType(ClassLayout* layout) const;
 
     static ABIPassingSegment InRegister(regNumber reg, unsigned offset, unsigned size);
     static ABIPassingSegment OnStack(unsigned stackOffset, unsigned offset, unsigned size);

src/coreclr/jit/compiler.h

@@ -11635,6 +11635,7 @@ class Compiler
 #endif // defined(UNIX_AMD64_ABI)
 
     bool fgTryMorphStructArg(CallArg* arg);
+    bool FieldsMatchAbi(LclVarDsc* varDsc, const ABIPassingInformation& abiInfo);
 
     bool killGCRefs(GenTree* tree);
 

src/coreclr/jit/gentree.h

@@ -2735,16 +2735,10 @@ struct GenTreeFieldList : public GenTree
 
     class UseList
     {
-        Use* m_head;
-        Use* m_tail;
+        Use* m_head = nullptr;
+        Use* m_tail = nullptr;
 
     public:
-        UseList()
-            : m_head(nullptr)
-            , m_tail(nullptr)
-        {
-        }
-
         Use* GetHead() const
         {
             return m_head;
@@ -2802,6 +2796,12 @@ struct GenTreeFieldList : public GenTree
             }
         }
 
+        void Clear()
+        {
+            m_head = nullptr;
+            m_tail = nullptr;
+        }
+
         bool IsSorted() const
         {
             unsigned offset = 0;

src/coreclr/jit/lower.cpp

@@ -1725,20 +1725,22 @@ void Lowering::LowerArg(GenTreeCall* call, CallArg* callArg)
     {
         if (abiInfo.HasAnyRegisterSegment())
         {
-#if FEATURE_MULTIREG_ARGS
-            if ((abiInfo.NumSegments > 1) && arg->OperIs(GT_FIELD_LIST))
+            if (arg->OperIs(GT_FIELD_LIST) || (abiInfo.NumSegments > 1))
             {
-                unsigned int regIndex = 0;
-                for (GenTreeFieldList::Use& use : arg->AsFieldList()->Uses())
+                if (!arg->OperIs(GT_FIELD_LIST))
                 {
-                    const ABIPassingSegment& segment = abiInfo.Segment(regIndex);
-                    InsertPutArgReg(&use.NodeRef(), segment);
-
-                    regIndex++;
+                    // Primitive arg, but the ABI requires it to be split into
+                    // registers. Insert the field list here.
+                    GenTreeFieldList* fieldList = comp->gtNewFieldList();
+                    fieldList->AddFieldLIR(comp, arg, 0, genActualType(arg->TypeGet()));
+                    BlockRange().InsertAfter(arg, fieldList);
+                    arg = *ppArg = fieldList;
                 }
+
+                LowerArgFieldList(callArg, arg->AsFieldList());
+                arg = *ppArg;
             }
             else
-#endif // FEATURE_MULTIREG_ARGS
             {
                 assert(abiInfo.HasExactlyOneRegisterSegment());
                 InsertPutArgReg(ppArg, abiInfo.Segment(0));
@@ -4809,6 +4811,18 @@ void Lowering::LowerRet(GenTreeOp* ret)
     ContainCheckRet(ret);
 }
 
+struct LowerFieldListRegisterInfo
+{
+    unsigned  Offset;
+    var_types RegType;
+
+    LowerFieldListRegisterInfo(unsigned offset, var_types regType)
+        : Offset(offset)
+        , RegType(regType)
+    {
+    }
+};
+
 //----------------------------------------------------------------------------------------------
 // LowerRetFieldList:
 //   Lower a returned FIELD_LIST node.
@@ -4822,21 +4836,18 @@ void Lowering::LowerRetFieldList(GenTreeOp* ret, GenTreeFieldList* fieldList)
     const ReturnTypeDesc& retDesc = comp->compRetTypeDesc;
     unsigned              numRegs = retDesc.GetReturnRegCount();
 
-    bool isCompatible = IsFieldListCompatibleWithReturn(fieldList);
+    auto getRegInfo = [=, &retDesc](unsigned regIndex) {
+        unsigned  offset  = retDesc.GetReturnFieldOffset(regIndex);
+        var_types regType = genActualType(retDesc.GetReturnRegType(regIndex));
+        return LowerFieldListRegisterInfo(offset, regType);
+    };
+
+    bool isCompatible = IsFieldListCompatibleWithRegisters(fieldList, numRegs, getRegInfo);
     if (!isCompatible)
     {
-        JITDUMP("Spilling field list [%06u] to stack\n", Compiler::dspTreeID(fieldList));
-        unsigned   lclNum = comp->lvaGrabTemp(true DEBUGARG("Spilled local for return value"));
+        unsigned lclNum =
+            StoreFieldListToNewLocal(comp->typGetObjLayout(comp->info.compMethodInfo->args.retTypeClass), fieldList);
         LclVarDsc* varDsc = comp->lvaGetDesc(lclNum);
-        comp->lvaSetStruct(lclNum, comp->info.compMethodInfo->args.retTypeClass, false);
-        comp->lvaSetVarDoNotEnregister(lclNum DEBUGARG(DoNotEnregisterReason::BlockOpRet));
-
-        for (GenTreeFieldList::Use& use : fieldList->Uses())
-        {
-            GenTree* store = comp->gtNewStoreLclFldNode(lclNum, use.GetType(), use.GetOffset(), use.GetNode());
-            BlockRange().InsertAfter(use.GetNode(), store);
-            LowerNode(store);
-        }
 
         GenTree* retValue = comp->gtNewLclvNode(lclNum, varDsc->TypeGet());
         ret->SetReturnValue(retValue);
@@ -4859,7 +4870,89 @@ void Lowering::LowerRetFieldList(GenTreeOp* ret, GenTreeFieldList* fieldList)
         return;
     }
 
-    LowerFieldListToFieldListOfRegisters(fieldList);
+    LowerFieldListToFieldListOfRegisters(fieldList, numRegs, getRegInfo);
+}
+
+//----------------------------------------------------------------------------------------------
+// StoreFieldListToNewLocal:
+//   Create a new local with the specified layout and store the specified
+//   fields of the specified FIELD_LIST into it.
+//
+// Arguments:
+//   layout    - Layout of the new local
+//   fieldList - Fields to store to it
+//
+// Returns:
+//   Var number of new local.
+//
+unsigned Lowering::StoreFieldListToNewLocal(ClassLayout* layout, GenTreeFieldList* fieldList)
+{
+    JITDUMP("Spilling field list [%06u] to stack\n", Compiler::dspTreeID(fieldList));
+    unsigned   lclNum = comp->lvaGrabTemp(true DEBUGARG("Spilled local for field list"));
+    LclVarDsc* varDsc = comp->lvaGetDesc(lclNum);
+    comp->lvaSetStruct(lclNum, layout, false);
+    comp->lvaSetVarDoNotEnregister(lclNum DEBUGARG(DoNotEnregisterReason::LocalField));
+
+    for (GenTreeFieldList::Use& use : fieldList->Uses())
+    {
+        GenTree* store = comp->gtNewStoreLclFldNode(lclNum, use.GetType(), use.GetOffset(), use.GetNode());
+        BlockRange().InsertAfter(use.GetNode(), store);
+        LowerNode(store);
+    }
+
+    return lclNum;
+}
+
+//----------------------------------------------------------------------------------------------
+// LowerArgFieldList:
+//   Lower an argument FIELD_LIST node.
+//
+// Arguments:
+//   arg       - The argument
+//   fieldList - The FIELD_LIST node
+//
+void Lowering::LowerArgFieldList(CallArg* arg, GenTreeFieldList* fieldList)
+{
+    assert(!arg->AbiInfo.HasAnyStackSegment());
+
+    auto getRegInfo = [=](unsigned regIndex) {
+        const ABIPassingSegment& seg = arg->AbiInfo.Segment(regIndex);
+        return LowerFieldListRegisterInfo(seg.Offset, seg.GetRegisterType());
+    };
+
+    bool isCompatible = IsFieldListCompatibleWithRegisters(fieldList, arg->AbiInfo.NumSegments, getRegInfo);
+    if (!isCompatible)
+    {
+        ClassLayout* layout = comp->typGetObjLayout(arg->GetSignatureClassHandle());
+        unsigned     lclNum = StoreFieldListToNewLocal(layout, fieldList);
+        fieldList->Uses().Clear();
+        for (const ABIPassingSegment& seg : arg->AbiInfo.Segments())
+        {
+            GenTreeLclFld* fld = comp->gtNewLclFldNode(lclNum, seg.GetRegisterType(layout), seg.Offset);
+            fieldList->AddFieldLIR(comp, fld, seg.Offset, fld->TypeGet());
+            BlockRange().InsertBefore(fieldList, fld);
+        }
+    }
+    else
+    {
+        LowerFieldListToFieldListOfRegisters(fieldList, arg->AbiInfo.NumSegments, getRegInfo);
+    }
+
+    GenTreeFieldList::Use* field = fieldList->Uses().GetHead();
+    for (const ABIPassingSegment& seg : arg->AbiInfo.Segments())
+    {
+        assert((field != nullptr) && "Ran out of fields while inserting PUTARG_REG");
+        InsertPutArgReg(&field->NodeRef(), seg);
+        field = field->GetNext();
+    }
+
+    assert((field == nullptr) && "Missed fields while inserting PUTARG_REG");
+
+    arg->NodeRef() = fieldList->SoleFieldOrThis();
+    if (arg->GetNode() != fieldList)
+    {
+        BlockRange().Remove(fieldList);
+    }
 }
 
 //----------------------------------------------------------------------------------------------
@@ -4874,21 +4967,29 @@ void Lowering::LowerRetFieldList(GenTreeOp* ret, GenTreeFieldList* fieldList)
 //   True if the fields of the FIELD_LIST are all direct insertions into the
 //   return registers.
 //
-bool Lowering::IsFieldListCompatibleWithReturn(GenTreeFieldList* fieldList)
+template <typename GetRegisterInfoFunc>
+bool Lowering::IsFieldListCompatibleWithRegisters(GenTreeFieldList*   fieldList,
+                                                  unsigned            numRegs,
+                                                  GetRegisterInfoFunc getRegInfo)
 {
-    JITDUMP("Checking if field list [%06u] is compatible with return ABI: ", Compiler::dspTreeID(fieldList));
-    const ReturnTypeDesc& retDesc    = comp->compRetTypeDesc;
-    unsigned              numRetRegs = retDesc.GetReturnRegCount();
+    JITDUMP("Checking if field list [%06u] is compatible with registers: ", Compiler::dspTreeID(fieldList));
 
     GenTreeFieldList::Use* use = fieldList->Uses().GetHead();
-    for (unsigned i = 0; i < numRetRegs; i++)
+    for (unsigned i = 0; i < numRegs; i++)
     {
-        unsigned  regStart = retDesc.GetReturnFieldOffset(i);
-        var_types regType  = retDesc.GetReturnRegType(i);
-        unsigned  regEnd   = regStart + genTypeSize(regType);
+        LowerFieldListRegisterInfo regInfo  = getRegInfo(i);
+        unsigned                   regStart = regInfo.Offset;
+        var_types                  regType  = regInfo.RegType;
+        unsigned                   regEnd   = regStart + genTypeSize(regType);
+
+        if ((i == numRegs - 1) && !varTypeUsesFloatReg(regType))
+        {
+            // Allow tail end to pass undefined bits into the register
+            regEnd = regStart + REGSIZE_BYTES;
+        }
 
         // TODO-CQ: Could just create a 0 for this.
-        if (use == nullptr)
+        if ((use == nullptr) || (use->GetOffset() >= regEnd))
         {
             JITDUMP("it is not; register %u has no corresponding field\n", i);
             return false;
@@ -4949,19 +5050,26 @@ bool Lowering::IsFieldListCompatibleWithReturn(GenTreeFieldList* fieldList)
 // Arguments:
 //     fieldList - The field list
 //
-void Lowering::LowerFieldListToFieldListOfRegisters(GenTreeFieldList* fieldList)
+template <typename GetRegisterInfoFunc>
+void Lowering::LowerFieldListToFieldListOfRegisters(GenTreeFieldList*   fieldList,
+                                                    unsigned            numRegs,
+                                                    GetRegisterInfoFunc getRegInfo)
 {
-    const ReturnTypeDesc& retDesc = comp->compRetTypeDesc;
-    unsigned              numRegs = retDesc.GetReturnRegCount();
-
     GenTreeFieldList::Use* use = fieldList->Uses().GetHead();
     assert(fieldList->Uses().IsSorted());
 
     for (unsigned i = 0; i < numRegs; i++)
     {
-        unsigned  regStart = retDesc.GetReturnFieldOffset(i);
-        var_types regType  = genActualType(retDesc.GetReturnRegType(i));
-        unsigned  regEnd   = regStart + genTypeSize(regType);
+        LowerFieldListRegisterInfo regInfo  = getRegInfo(i);
+        unsigned                   regStart = regInfo.Offset;
+        var_types                  regType  = regInfo.RegType;
+        unsigned                   regEnd   = regStart + genTypeSize(regType);
+
+        if ((i == numRegs - 1) && !varTypeUsesFloatReg(regType))
+        {
+            // Allow tail end to pass undefined bits into the register
+            regEnd = regStart + REGSIZE_BYTES;
+        }
 
         GenTreeFieldList::Use* regEntry = use;
 
@@ -5001,7 +5109,7 @@ void Lowering::LowerFieldListToFieldListOfRegisters(GenTreeFieldList* fieldList)
             }
 
             // If this is a float -> int insertion, then we need the bitcast now.
-            if (varTypeUsesFloatReg(value) && varTypeUsesIntReg(regType))
+            if (varTypeUsesFloatReg(value) && varTypeUsesIntReg(regInfo.RegType))
             {
                 assert((genTypeSize(value) == 4) || (genTypeSize(value) == 8));
                 var_types castType = genTypeSize(value) == 4 ? TYP_INT : TYP_LONG;

src/coreclr/jit/lower.h

@@ -187,10 +187,14 @@ class Lowering final : public Phase
     GenTree*   LowerAsyncContinuation(GenTree* asyncCont);
     void       LowerReturnSuspend(GenTree* retSuspend);
     void       LowerRetFieldList(GenTreeOp* ret, GenTreeFieldList* fieldList);
-    bool       IsFieldListCompatibleWithReturn(GenTreeFieldList* fieldList);
-    void       LowerFieldListToFieldListOfRegisters(GenTreeFieldList* fieldList);
-    void       LowerCallStruct(GenTreeCall* call);
-    void       LowerStoreSingleRegCallStruct(GenTreeBlk* store);
+    unsigned   StoreFieldListToNewLocal(ClassLayout* layout, GenTreeFieldList* fieldList);
+    void       LowerArgFieldList(CallArg* arg, GenTreeFieldList* fieldList);
+    template <typename GetRegisterInfoFunc>
+    bool IsFieldListCompatibleWithRegisters(GenTreeFieldList* fieldList, unsigned numRegs, GetRegisterInfoFunc func);
+    template <typename GetRegisterInfoFunc>
+    void LowerFieldListToFieldListOfRegisters(GenTreeFieldList* fieldList, unsigned numRegs, GetRegisterInfoFunc func);
+    void LowerCallStruct(GenTreeCall* call);
+    void LowerStoreSingleRegCallStruct(GenTreeBlk* store);
 #if !defined(WINDOWS_AMD64_ABI)
     GenTreeLclVar* SpillStructCallResult(GenTreeCall* call) const;
 #endif // WINDOWS_AMD64_ABI