diff --git a/src/coreclr/jit/codegens390x.cpp b/src/coreclr/jit/codegens390x.cpp index 9bb1a4b05e3def..41bdffee4fa829 100644 --- a/src/coreclr/jit/codegens390x.cpp +++ b/src/coreclr/jit/codegens390x.cpp @@ -571,9 +571,9 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode) #endif // PROFILING_SUPPORTED break; -// case GT_LCLHEAP: -// genLclHeap(treeNode); -// break; + case GT_LCLHEAP: + genLclHeap(treeNode); + break; case GT_CNS_INT: case GT_CNS_DBL: @@ -5183,7 +5183,8 @@ void CodeGen::genFnEpilog(BasicBlock* block) } #else // TARGET_ARM64 int totalFrameSize = genTotalFrameSize(); - GetEmitter()->emitIns_R_R_R_I(INS_lmg, EA_PTRSIZE, REG_FP, REG_SP, REG_SPBASE, totalFrameSize + 88); + regNumber lmgBase = compiler->compLocallocUsed ? REG_FPBASE : REG_SPBASE; + GetEmitter()->emitIns_R_R_R_I(INS_lmg, EA_PTRSIZE, REG_FP, REG_SP, lmgBase, totalFrameSize + 88); inst_RV(INS_ret, REG_LR, TYP_I_IMPL); compiler->unwindReturn(REG_LR); #endif // TARGET_ARM64 @@ -5864,6 +5865,120 @@ void CodeGen::genZeroInitFrame(unsigned frameSize, regNumber initReg, bool* pIni } } //----------------------------------------------------------------------------- +// genLclHeap: Dynamic memory allocation on stack frame +// +// Moves r15 down by (aligned size + RSA + outgoing arg area) and +// returns the address just above the save/outgoing area, so calls made after the +// localloc cannot overwrite the buffer. The back-chain is copied to the new r15. +// +void CodeGen::genLclHeap(GenTree* tree) +{ + assert(tree->OperIs(GT_LCLHEAP)); + assert(compiler->compLocallocUsed); + + GenTree* size = tree->AsOp()->gtOp1; + emitter* emit = GetEmitter(); + regNumber targetReg = tree->GetRegNum(); + + noway_assert((genActualType(size->gtType) == TYP_INT) || (genActualType(size->gtType) == TYP_I_IMPL)); + noway_assert(isFramePointerUsed()); + noway_assert(genStackLevel == 0); + + const ssize_t frameBelow = S390X_REG_SAVE_AREA_SIZE + (ssize_t)compiler->lvaOutgoingArgSpaceSize; + + if (size->IsCnsIntOrI()) + { + assert(size->isContained()); + size_t amount = size->AsIntCon()->gtIconVal; + + if (amount == 0) + { + emit->emitIns_R_I(INS_lgfi, EA_8BYTE, targetReg, 0); + genProduceReg(tree); + return; + } + + // Round size up to a multiple of STACK_ALIGN (8) so SP stays 8-byte aligned. + ssize_t alignedSize = (ssize_t)AlignUp(amount, STACK_ALIGN); + + if (alignedSize <= 0x7FFFF) + { + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, REG_SPBASE, REG_SPBASE, -alignedSize); + } + else + { + // REG_R0 is scratch holding the negated size. + regNumber rTemp = REG_R0; + emit->emitIns_R_I(INS_lgfi, EA_8BYTE, rTemp, -(int)alignedSize); + emit->emitIns_R_R(INS_agr, EA_PTRSIZE, REG_SPBASE, rTemp); + } + + // REG_FPBASE (r11) is stored as the back chain at the new SP. + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, REG_FPBASE, REG_SPBASE, 0); + + if (compiler->info.compInitMem) + { + regNumber rZero = REG_R0; + emit->emitIns_R_R(INS_xgr, EA_8BYTE, rZero, rZero); + + for (ssize_t offset = 0; offset < alignedSize; offset += 8) + { + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, rZero, REG_SPBASE, frameBelow + offset); + } + } + + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); + } + else + { + // sizeReg is the read-only input, alignedReg the rounded size, tmpReg the mask then write pointer. + regNumber sizeReg = genConsumeReg(size); + regNumber alignedReg = internalRegisters.Extract(tree); + regNumber tmpReg = internalRegisters.Extract(tree); + BasicBlock* endLabel = genCreateTempLabel(); + + // Zero-extending a 32-bit value so the upper 32 bits are clean. + emit->emitIns_R_R(size->TypeIs(TYP_INT) ? INS_llgfr : INS_lgr, EA_8BYTE, alignedReg, sizeReg); + + emit->emitIns_R_I(INS_lgfi, EA_8BYTE, targetReg, 0); + emit->emitIns_R_I(INS_cgfi, EA_8BYTE, alignedReg, 0); + inst_JMP(EJ_eq, endLabel); + + // alignedReg = (size + 7) & -8. + emit->emitIns_R_I(INS_agfi, EA_8BYTE, alignedReg, 7); + emit->emitIns_R_I(INS_lgfi, EA_8BYTE, tmpReg, -8); + emit->emitIns_R_R(INS_ngr, EA_8BYTE, alignedReg, tmpReg); + + emit->emitIns_R_R(INS_sgr, EA_8BYTE, REG_SPBASE, alignedReg); + + // REG_FPBASE (r11) is stored as the back chain at the new SP. + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, REG_FPBASE, REG_SPBASE, 0); + + if (compiler->info.compInitMem) + { + // tmpReg is the write pointer, alignedReg the remaining byte count. + BasicBlock* loop = genCreateTempLabel(); + regNumber rZero = REG_R0; + emit->emitIns_R_R(INS_xgr, EA_8BYTE, rZero, rZero); + + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, tmpReg, REG_SPBASE, frameBelow); + + genDefineTempLabel(loop); + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, rZero, tmpReg, 0); + emit->emitIns_R_I(INS_aghi, EA_8BYTE, tmpReg, 8); + emit->emitIns_R_I(INS_aghi, EA_8BYTE, alignedReg, -8); + emit->emitIns_R_I(INS_cgfi, EA_8BYTE, alignedReg, 0); + inst_JMP(EJ_ne, loop); + } + + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); + + genDefineTempLabel(endLabel); + } + + genProduceReg(tree); +} +//----------------------------------------------------------------------------- // genAllocLclFrame: Allocate the local stack frame for s390x architecture // // Arguments: diff --git a/src/coreclr/jit/emits390x.cpp b/src/coreclr/jit/emits390x.cpp index 0f6497a975efdf..88c3bcd1bd76b3 100644 --- a/src/coreclr/jit/emits390x.cpp +++ b/src/coreclr/jit/emits390x.cpp @@ -7534,7 +7534,7 @@ void emitter::emitIns_R_S(instruction ins, emitAttr attr, regNumber reg1, int va int disp = base + offs; ssize_t imm = disp; - regNumber reg2 = REG_SPBASE;//encodingSPtoZR(FPbased ? REG_FPBASE : REG_SPBASE); + regNumber reg2 = FPbased ? REG_FPBASE : REG_SPBASE; // TODO-ARM64-CQ: use unscaled loads? /* Figure out the encoding format of the instruction */ @@ -7815,7 +7815,7 @@ void emitter::emitIns_S_R(instruction ins, emitAttr attr, regNumber reg1, int va ssize_t imm = disp; // TODO-ARM64-CQ: with compLocallocUsed, should we use REG_SAVED_LOCALLOC_SP instead? - regNumber reg2 = REG_SPBASE;//encodingSPtoZR(FPbased ? REG_FPBASE : REG_SPBASE); + regNumber reg2 = FPbased ? REG_FPBASE : REG_SPBASE; // TODO-ARM64-CQ: use unscaled loads? /* Figure out the encoding format of the instruction */ diff --git a/src/coreclr/jit/lsras390x.cpp b/src/coreclr/jit/lsras390x.cpp index 1d98ecbabda434..0ed796fdc53db6 100644 --- a/src/coreclr/jit/lsras390x.cpp +++ b/src/coreclr/jit/lsras390x.cpp @@ -1979,11 +1979,9 @@ int LinearScan::BuildNode(GenTree* tree) else { srcCount = 1; - if (!compiler->info.compInitMem) - { - buildInternalIntRegisterDefForNode(tree); - buildInternalIntRegisterDefForNode(tree); - } + // Allocate 2 registers in either case + buildInternalIntRegisterDefForNode(tree); + buildInternalIntRegisterDefForNode(tree); } if (!size->isContained()) diff --git a/src/s390tests/tests/Allocation.cs b/src/s390tests/tests/Allocation.cs new file mode 100644 index 00000000000000..277227e45c5ab5 --- /dev/null +++ b/src/s390tests/tests/Allocation.cs @@ -0,0 +1,44 @@ +using System; +using System.Runtime.CompilerServices; + +public class S390xInstructionTest +{ + public static int Main() + { + int result = s390xHw(); + Console.WriteLine(result == 0 ? "PASS" : "FAIL: " + result); + return result == 0 ? 0 : 1; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + static int Consume(Span s) + { + return s.Length; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + static int FlushTest(int n) + { + Span buf = stackalloc byte[n]; + buf[0] = 0xAA; + buf[n - 1] = 0xBB; + return buf.Length; // Returns n + } + + public static int s390xHw() + { + int r; + + // 1. Test FlushTest buffer length returns + r = FlushTest(5); + if (r != 5) return 1; + + r = FlushTest(40); + if (r != 40) return 2; + + r = FlushTest(400); + if (r != 400) return 3; + + return 0; // PASS + } +} diff --git a/src/s390tests/tests/Allocation_2.cs b/src/s390tests/tests/Allocation_2.cs new file mode 100644 index 00000000000000..43170af07314e4 --- /dev/null +++ b/src/s390tests/tests/Allocation_2.cs @@ -0,0 +1,161 @@ +using System; +using System.Runtime.CompilerServices; + +#pragma warning disable CA2014 // Suppress warning for stackalloc inside loop + +public class S390xInstructionTest +{ + public static int Main() + { + int result = s390xHw(); + Console.WriteLine(result == 0 ? "PASS" : "FAIL: " + result); + return result == 0 ? 0 : 1; + } + + public static int s390xHw() + { + // 1. Test Constant Size Allocation (<= 64B) + if (!TestConstantAllocation()) + return 101; + + // 2. Test Dynamic Size Allocation (5B, 127B, 1024B) + if (!TestDynamicAllocation(5) || !TestDynamicAllocation(127) || !TestDynamicAllocation(1024)) + return 102; + + // 3. Test Zero-Size Allocation (0B) + if (!TestZeroSizeAllocation()) + return 103; + + // 4. Test Multiple Allocations in a Single Method + if (!TestMultipleAllocations()) + return 104; + + // 5. Test localloc in a Loop with Child Calls + if (!TestLocallocInLoop()) + return 105; + + // All tests passed + return 0; + } + + // ========================================================================= + // Test 1: Constant Size Allocation (Unrolled & Zero-Initialized) + // ========================================================================= + [MethodImpl(MethodImplOptions.NoInlining)] + private static unsafe bool TestConstantAllocation() + { + byte* buf = stackalloc byte[16]; + + // Verify zero-init + for (int i = 0; i < 16; i++) + { + if (buf[i] != 0) return false; + buf[i] = (byte)(i + 1); + } + + // Verify written data + for (int i = 0; i < 16; i++) + { + if (buf[i] != (byte)(i + 1)) return false; + } + + return true; + } + + // ========================================================================= + // Test 2: Dynamic (Variable) Size Allocation with Loop Zeroing + // ========================================================================= + [MethodImpl(MethodImplOptions.NoInlining)] + private static unsafe bool TestDynamicAllocation(int size) + { + byte* buf = stackalloc byte[size]; + + // Verify zero-init + for (int i = 0; i < size; i++) + { + if (buf[i] != 0) return false; + } + + // Write pattern + for (int i = 0; i < size; i++) + { + buf[i] = (byte)(i & 0xFF); + } + + // Verify pattern + for (int i = 0; i < size; i++) + { + if (buf[i] != (byte)(i & 0xFF)) return false; + } + + return true; + } + + // ========================================================================= + // Test 3: Zero-Size Allocation (Bailout path returning NULL) + // ========================================================================= + [MethodImpl(MethodImplOptions.NoInlining)] + private static unsafe bool TestZeroSizeAllocation() + { + int zero = 0; + byte* pDynamicZero = stackalloc byte[zero]; + byte* pConstZero = stackalloc byte[0]; + + return (pDynamicZero == null) && (pConstZero == null); + } + + // ========================================================================= + // Test 4: Multiple localloc Calls in a Single Method + // ========================================================================= + [MethodImpl(MethodImplOptions.NoInlining)] + private static unsafe bool TestMultipleAllocations() + { + int localBefore = 0x12345678; + + byte* buf1 = stackalloc byte[32]; + byte* buf2 = stackalloc byte[64]; + byte* buf3 = stackalloc byte[128]; + + int localAfter = 0x76543210; + + for (int i = 0; i < 32; i++) buf1[i] = 0xAA; + for (int i = 0; i < 64; i++) buf2[i] = 0xBB; + for (int i = 0; i < 128; i++) buf3[i] = 0xCC; + + for (int i = 0; i < 32; i++) if (buf1[i] != 0xAA) return false; + for (int i = 0; i < 64; i++) if (buf2[i] != 0xBB) return false; + for (int i = 0; i < 128; i++) if (buf3[i] != 0xCC) return false; + + if (localBefore != 0x12345678 || localAfter != 0x76543210) return false; + + return true; + } + + // ========================================================================= + // Test 5: localloc in a Loop with Child Calls (RSA / Outgoing Arg Safety) + // ========================================================================= + [MethodImpl(MethodImplOptions.NoInlining)] + private static int HelperFunction(int a, int b, int c, int d, int e, int f) + { + return a + b + c + d + e + f; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + private static unsafe bool TestLocallocInLoop() + { + for (int iter = 1; iter <= 10; iter++) + { + int size = iter * 16; + byte* buf = stackalloc byte[size]; + + buf[0] = (byte)iter; + buf[size - 1] = (byte)(iter * 2); + + int sum = HelperFunction(buf[0], buf[size - 1], 1, 2, 3, 4); + if (sum != (iter + (iter * 2) + 10)) return false; + + if (buf[0] != (byte)iter || buf[size - 1] != (byte)(iter * 2)) return false; + } + return true; + } +}