arm64 osx: support byte sizes from lowering to codegen. (#43024)

* arm64 osx: support byte sizes from lowering to codegen.

* a better message about GT_PUTARG_SPLIT.

* Fix arm

* Response review.

* format GenTreePutArgStk(genTreeOps.

* response Egor's review.

src/coreclr/src/jit/codegenarm64.cpp

@@ -242,6 +242,7 @@ void CodeGen::genPrologSaveRegPair(regNumber reg1,
         // stp REG, REG + 1, [SP, #offset]
         // 64-bit STP offset range: -512 to 504, multiple of 8.
         assert(spOffset <= 504);
+        assert((spOffset % 8) == 0);
         GetEmitter()->emitIns_R_R_R_I(INS_stp, EA_PTRSIZE, reg1, reg2, REG_SPBASE, spOffset);
 
 #if defined(TARGET_UNIX)

src/coreclr/src/jit/codegenarmarch.cpp

@@ -665,9 +665,13 @@ void CodeGen::genIntrinsic(GenTree* treeNode)
 void CodeGen::genPutArgStk(GenTreePutArgStk* treeNode)
 {
     assert(treeNode->OperIs(GT_PUTARG_STK));
-    GenTree*  source     = treeNode->gtOp1;
+    GenTree* source = treeNode->gtOp1;
+#if !defined(OSX_ARM64_ABI)
     var_types targetType = genActualType(source->TypeGet());
-    emitter*  emit       = GetEmitter();
+#else
+    var_types targetType = source->TypeGet();
+#endif
+    emitter* emit = GetEmitter();
 
     // This is the varNum for our store operations,
     // typically this is the varNum for the Outgoing arg space
@@ -678,12 +682,12 @@ void CodeGen::genPutArgStk(GenTreePutArgStk* treeNode)
     // Get argument offset to use with 'varNumOut'
     // Here we cross check that argument offset hasn't changed from lowering to codegen since
     // we are storing arg slot number in GT_PUTARG_STK node in lowering phase.
-    unsigned argOffsetOut = treeNode->gtSlotNum * TARGET_POINTER_SIZE;
+    unsigned argOffsetOut = treeNode->getArgOffset();
 
 #ifdef DEBUG
     fgArgTabEntry* curArgTabEntry = compiler->gtArgEntryByNode(treeNode->gtCall, treeNode);
-    assert(curArgTabEntry);
-    assert(argOffsetOut == (curArgTabEntry->slotNum * TARGET_POINTER_SIZE));
+    assert(curArgTabEntry != nullptr);
+    DEBUG_ARG_SLOTS_ASSERT(argOffsetOut == (curArgTabEntry->slotNum * TARGET_POINTER_SIZE));
 #endif // DEBUG
 
     // Whether to setup stk arg in incoming or out-going arg area?
@@ -730,6 +734,21 @@ void CodeGen::genPutArgStk(GenTreePutArgStk* treeNode)
             return;
         }
 
+#if defined(OSX_ARM64_ABI)
+        switch (treeNode->GetStackByteSize())
+        {
+            case 1:
+                targetType = TYP_BYTE;
+                break;
+            case 2:
+                targetType = TYP_SHORT;
+                break;
+            default:
+                assert(treeNode->GetStackByteSize() >= 4);
+                break;
+        }
+#endif
+
         instruction storeIns  = ins_Store(targetType);
         emitAttr    storeAttr = emitTypeSize(targetType);
 
@@ -1161,7 +1180,7 @@ void CodeGen::genPutArgSplit(GenTreePutArgSplit* treeNode)
     emitter* emit         = GetEmitter();
     unsigned varNumOut    = compiler->lvaOutgoingArgSpaceVar;
     unsigned argOffsetMax = compiler->lvaOutgoingArgSpaceSize;
-    unsigned argOffsetOut = treeNode->gtSlotNum * TARGET_POINTER_SIZE;
+    unsigned argOffsetOut = treeNode->getArgOffset();
 
     if (source->OperGet() == GT_FIELD_LIST)
     {
@@ -1292,13 +1311,12 @@ void CodeGen::genPutArgSplit(GenTreePutArgSplit* treeNode)
             assert(!compiler->IsHfa(source->AsObj()->GetLayout()->GetClassHandle()));
         }
 
-        int          structSize = treeNode->getArgSize();
-        ClassLayout* layout     = source->AsObj()->GetLayout();
+        ClassLayout* layout = source->AsObj()->GetLayout();
 
         // Put on stack first
         unsigned nextIndex     = treeNode->gtNumRegs;
         unsigned structOffset  = nextIndex * TARGET_POINTER_SIZE;
-        int      remainingSize = structSize - structOffset;
+        int      remainingSize = treeNode->GetStackByteSize();
 
         // remainingSize is always multiple of TARGET_POINTER_SIZE
         assert(remainingSize % TARGET_POINTER_SIZE == 0);

src/coreclr/src/jit/codegencommon.cpp

@@ -2355,7 +2355,7 @@ void CodeGen::genEmitMachineCode()
     }
 #endif
 
-#if EMIT_TRACK_STACK_DEPTH && defined(DEBUG) && !defined(OSX_ARM64_ABI)
+#if EMIT_TRACK_STACK_DEPTH && defined(DEBUG_ARG_SLOTS)
     // Check our max stack level. Needed for fgAddCodeRef().
     // We need to relax the assert as our estimation won't include code-gen
     // stack changes (which we know don't affect fgAddCodeRef()).

src/coreclr/src/jit/codegenlinear.cpp

@@ -1704,8 +1704,6 @@ void CodeGen::genConsumePutStructArgStk(GenTreePutArgStk* putArgNode,
     assert((src->gtOper == GT_OBJ) || ((src->gtOper == GT_IND && varTypeIsSIMD(src))));
     GenTree* srcAddr = src->gtGetOp1();
 
-    unsigned int size = putArgNode->getArgSize();
-
     assert(dstReg != REG_NA);
     assert(srcReg != REG_NA);
 
@@ -1757,6 +1755,7 @@ void CodeGen::genConsumePutStructArgStk(GenTreePutArgStk* putArgNode,
 
     if (sizeReg != REG_NA)
     {
+        unsigned size = putArgNode->GetStackByteSize();
         inst_RV_IV(INS_mov, sizeReg, size, EA_PTRSIZE);
     }
 }

src/coreclr/src/jit/codegenxarch.cpp

@@ -3156,7 +3156,7 @@ void CodeGen::genStructPutArgUnroll(GenTreePutArgStk* putArgNode)
     // in genPutStructArgStk().
     noway_assert(src->TypeGet() == TYP_STRUCT);
 
-    unsigned size = putArgNode->getArgSize();
+    unsigned size = putArgNode->GetStackByteSize();
     assert(size <= CPBLK_UNROLL_LIMIT);
 
     emitter* emit         = GetEmitter();
@@ -5049,7 +5049,7 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
         if (arg->OperIs(GT_PUTARG_STK) && ((arg->gtFlags & GTF_LATE_ARG) == 0))
         {
             GenTree* source = arg->AsPutArgStk()->gtGetOp1();
-            unsigned size   = arg->AsPutArgStk()->getArgSize();
+            unsigned size   = arg->AsPutArgStk()->GetStackByteSize();
             stackArgBytes += size;
 #ifdef DEBUG
             fgArgTabEntry* curArgTabEntry = compiler->gtArgEntryByNode(call, arg);
@@ -7292,7 +7292,7 @@ void CodeGen::genRemoveAlignmentAfterCall(GenTreeCall* call, unsigned bias)
 //
 bool CodeGen::genAdjustStackForPutArgStk(GenTreePutArgStk* putArgStk)
 {
-    const unsigned argSize = putArgStk->getArgSize();
+    const unsigned argSize = putArgStk->GetStackByteSize();
     GenTree*       source  = putArgStk->gtGetOp1();
 
 #ifdef FEATURE_SIMD
@@ -7389,7 +7389,7 @@ void CodeGen::genPutArgStkFieldList(GenTreePutArgStk* putArgStk)
     // If we are pushing the arguments (i.e. we have not pre-adjusted the stack), then we are pushing them
     // in reverse order, so we start with the current field offset at the size of the struct arg (which must be
     // a multiple of the target pointer size).
-    unsigned  currentOffset   = (preAdjustedStack) ? 0 : putArgStk->getArgSize();
+    unsigned  currentOffset   = (preAdjustedStack) ? 0 : putArgStk->GetStackByteSize();
     unsigned  prevFieldOffset = currentOffset;
     regNumber intTmpReg       = REG_NA;
     regNumber simdTmpReg      = REG_NA;
@@ -7601,7 +7601,7 @@ void CodeGen::genPutArgStk(GenTreePutArgStk* putArgStk)
     // On a 32-bit target, all of the long arguments are handled with GT_FIELD_LISTs of TYP_INT.
     assert(targetType != TYP_LONG);
 
-    const unsigned argSize = putArgStk->getArgSize();
+    const unsigned argSize = putArgStk->GetStackByteSize();
     assert((argSize % TARGET_POINTER_SIZE) == 0);
 
     if (data->isContainedIntOrIImmed())
@@ -7653,12 +7653,12 @@ void CodeGen::genPutArgStk(GenTreePutArgStk* putArgStk)
         // Get argument offset on stack.
         // Here we cross check that argument offset hasn't changed from lowering to codegen since
         // we are storing arg slot number in GT_PUTARG_STK node in lowering phase.
-        int            argOffset      = putArgStk->getArgOffset();
+        unsigned       argOffset      = putArgStk->getArgOffset();
 
 #ifdef DEBUG
         fgArgTabEntry* curArgTabEntry = compiler->gtArgEntryByNode(putArgStk->gtCall, putArgStk);
-        assert(curArgTabEntry);
-        assert(argOffset == (int)curArgTabEntry->slotNum * TARGET_POINTER_SIZE);
+        assert(curArgTabEntry != nullptr);
+        assert(argOffset == curArgTabEntry->slotNum * TARGET_POINTER_SIZE);
 #endif
 
         if (data->isContainedIntOrIImmed())
@@ -7899,7 +7899,10 @@ void CodeGen::genPutStructArgStk(GenTreePutArgStk* putArgStk)
         assert(m_pushStkArg);
 
         GenTree*       srcAddr  = source->gtGetOp1();
-        const unsigned numSlots = putArgStk->gtNumSlots;
+        const unsigned byteSize = putArgStk->GetStackByteSize();
+        assert(byteSize % TARGET_POINTER_SIZE == 0);
+        const unsigned numSlots = byteSize / TARGET_POINTER_SIZE;
+        assert(putArgStk->gtNumSlots == numSlots);
 
         regNumber  srcRegNum    = srcAddr->GetRegNum();
         const bool srcAddrInReg = srcRegNum != REG_NA;
@@ -7920,15 +7923,15 @@ void CodeGen::genPutStructArgStk(GenTreePutArgStk* putArgStk)
 
         for (int i = numSlots - 1; i >= 0; --i)
         {
-            emitAttr       slotAttr = emitTypeSize(layout->GetGCPtrType(i));
-            const unsigned offset   = i * TARGET_POINTER_SIZE;
+            emitAttr       slotAttr   = emitTypeSize(layout->GetGCPtrType(i));
+            const unsigned byteOffset = i * TARGET_POINTER_SIZE;
             if (srcAddrInReg)
             {
-                GetEmitter()->emitIns_AR_R(INS_push, slotAttr, REG_NA, srcRegNum, offset);
+                GetEmitter()->emitIns_AR_R(INS_push, slotAttr, REG_NA, srcRegNum, byteOffset);
             }
             else
             {
-                GetEmitter()->emitIns_S(INS_push, slotAttr, srcLclNum, srcLclOffset + offset);
+                GetEmitter()->emitIns_S(INS_push, slotAttr, srcLclNum, srcLclOffset + byteOffset);
             }
             AddStackLevel(TARGET_POINTER_SIZE);
         }
@@ -7945,7 +7948,10 @@ void CodeGen::genPutStructArgStk(GenTreePutArgStk* putArgStk)
         unsigned       numGCSlotsCopied = 0;
 #endif // DEBUG
 
-        const unsigned numSlots = putArgStk->gtNumSlots;
+        const unsigned byteSize = putArgStk->GetStackByteSize();
+        assert(byteSize % TARGET_POINTER_SIZE == 0);
+        const unsigned numSlots = byteSize / TARGET_POINTER_SIZE;
+        assert(putArgStk->gtNumSlots == numSlots);
         for (unsigned i = 0; i < numSlots;)
         {
             if (!layout->IsGCPtr(i))

src/coreclr/src/jit/compiler.h

@@ -1630,7 +1630,9 @@ struct fgArgTabEntry
         return 0;
     }
 
-    // Get the number of bytes that this argument is occupying on the stack.
+    // Get the number of bytes that this argument is occupying on the stack,
+    // including padding up to the target pointer size for platforms
+    // where a stack argument can't take less.
     unsigned GetStackByteSize() const
     {
         if (!IsSplit() && numRegs > 0)
@@ -1642,7 +1644,10 @@ struct fgArgTabEntry
 
         assert(GetByteSize() > TARGET_POINTER_SIZE * numRegs);
         unsigned stackByteSize = GetByteSize() - TARGET_POINTER_SIZE * numRegs;
-        return GetByteSize() - TARGET_POINTER_SIZE * numRegs;
+#if !defined(OSX_ARM64_ABI)
+        stackByteSize = roundUp(stackByteSize, TARGET_POINTER_SIZE);
+#endif
+        return stackByteSize;
     }
 
     var_types GetHfaType() const
@@ -1800,7 +1805,7 @@ struct fgArgTabEntry
         return size;
     }
 
-#endif // DEBUG && !OSX_ARM64_ABI
+#endif // DEBUG_ARG_SLOTS
 
 private:
     unsigned m_byteOffset;

src/coreclr/src/jit/flowgraph.cpp

@@ -9807,6 +9807,8 @@ void Compiler::fgSimpleLowering()
         JITDUMP("Bumping outgoingArgSpaceSize to %u for localloc", outgoingArgSpaceSize);
     }
 
+    assert((outgoingArgSpaceSize % TARGET_POINTER_SIZE) == 0);
+
     // Publish the final value and mark it as read only so any update
     // attempt later will cause an assert.
     lvaOutgoingArgSpaceSize = outgoingArgSpaceSize;

src/coreclr/src/jit/gentree.cpp

@@ -1206,7 +1206,7 @@ bool GenTreeCall::Equals(GenTreeCall* c1, GenTreeCall* c2)
 }
 
 #if !defined(FEATURE_PUT_STRUCT_ARG_STK)
-unsigned GenTreePutArgStk::getArgSize()
+unsigned GenTreePutArgStk::GetStackByteSize() const
 {
     return genTypeSize(genActualType(gtOp1->gtType));
 }
@@ -11523,10 +11523,16 @@ void Compiler::gtDispTree(GenTree*     tree,
 #if FEATURE_PUT_STRUCT_ARG_STK
         else if (tree->OperGet() == GT_PUTARG_STK)
         {
-            printf(" (%d slots)", tree->AsPutArgStk()->gtNumSlots);
-            if (tree->AsPutArgStk()->gtPutArgStkKind != GenTreePutArgStk::Kind::Invalid)
+            const GenTreePutArgStk* putArg = tree->AsPutArgStk();
+#if !defined(DEBUG_ARG_SLOTS)
+            printf(" (%d stackByteSize), (%d byteOffset)", putArg->GetStackByteSize(), putArg->getArgOffset());
+#else
+            printf(" (%d slots), (%d stackByteSize), (%d slot), (%d byteOffset)", putArg->gtNumSlots,
+                   putArg->GetStackByteSize(), putArg->gtSlotNum, putArg->getArgOffset());
+#endif
+            if (putArg->gtPutArgStkKind != GenTreePutArgStk::Kind::Invalid)
             {
-                switch (tree->AsPutArgStk()->gtPutArgStkKind)
+                switch (putArg->gtPutArgStkKind)
                 {
                     case GenTreePutArgStk::Kind::RepInstr:
                         printf(" (RepInstr)");
@@ -11545,6 +11551,18 @@ void Compiler::gtDispTree(GenTree*     tree,
                 }
             }
         }
+#if FEATURE_ARG_SPLIT
+        else if (tree->OperGet() == GT_PUTARG_SPLIT)
+        {
+            const GenTreePutArgSplit* putArg = tree->AsPutArgSplit();
+#if !defined(DEBUG_ARG_SLOTS)
+            printf(" (%d stackByteSize), (%d numRegs)", putArg->GetStackByteSize(), putArg->gtNumRegs);
+#else
+            printf(" (%d slots), (%d stackByteSize), (%d numRegs)", putArg->gtNumSlots, putArg->GetStackByteSize(),
+                   putArg->gtNumRegs);
+#endif
+        }
+#endif // FEATURE_ARG_SPLIT
 #endif // FEATURE_PUT_STRUCT_ARG_STK
 
         if (tree->gtOper == GT_INTRINSIC)
@@ -12007,9 +12025,7 @@ void Compiler::gtGetArgMsg(GenTreeCall* call, GenTree* arg, unsigned argNum, cha
             }
 #endif // TARGET_ARM
 #if FEATURE_FIXED_OUT_ARGS
-
             sprintf_s(bufp, bufLength, "arg%d out+%02x%c", argNum, curArgTabEntry->slotNum * TARGET_POINTER_SIZE, 0);
-
 #else
             sprintf_s(bufp, bufLength, "arg%d on STK%c", argNum, 0);
 #endif
@@ -12052,8 +12068,7 @@ void Compiler::gtGetLateArgMsg(GenTreeCall* call, GenTree* argx, int lateArgInde
 #else
     if (argReg == REG_STK)
     {
-        sprintf_s(bufp, bufLength, "arg%d in out+%02x%c", curArgTabEntry->argNum,
-                  curArgTabEntry->slotNum * TARGET_POINTER_SIZE, 0);
+        sprintf_s(bufp, bufLength, "arg%d in out+%02x%c", curArgTabEntry->argNum, curArgTabEntry->GetByteOffset(), 0);
     }
     else
 #endif