From 7acb3221ce6ab9dccfa8b5499eaee19216b7da8f Mon Sep 17 00:00:00 2001 From: Krishna-13-cyber Date: Fri, 23 Oct 2026 23:20:29 +0530 Subject: [PATCH 1/4] Add support for dynamic memory allocation on stack (localloc support) --- src/coreclr/jit/codegens390x.cpp | 76 ++++++++++++++++++++++++++++++-- src/coreclr/jit/emits390x.cpp | 4 +- 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/src/coreclr/jit/codegens390x.cpp b/src/coreclr/jit/codegens390x.cpp index 9bb1a4b05e3def..fa0f32f678c338 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,73 @@ 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(); + + if (compiler->info.compInitMem) + { + // TODO: Have to handle zeroing of memory + } + + 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; + } + + ssize_t total = (ssize_t)AlignUp(amount, STACK_ALIGN) + frameBelow; + if (total > 0x7FFFF) + { + // TODO: Have to handle the overflow + } + + regNumber tmpReg = internalRegisters.GetSingle(tree); + + emit->emitIns_R_R_I(INS_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, REG_SPBASE, REG_SPBASE, -total); + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, tmpReg, REG_SPBASE, 0); + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); + } + else + { + regNumber sizeReg = genConsumeReg(size); + regNumber alignedReg = internalRegisters.Extract(tree); + regNumber tmpReg = internalRegisters.Extract(tree); + + // alignedReg = ((size + 7) & -8) + frameBelow + emit->emitIns_R_R(INS_lgr, EA_8BYTE, alignedReg, sizeReg); + emit->emitIns_R_R(INS_ngr, EA_8BYTE, alignedReg, tmpReg); + emit->emitIns_R_R_I(INS_agfi, EA_8BYTE, alignedReg, alignedReg, frameBelow); + + emit->emitIns_R_R_I(INS_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); + emit->emitIns_R_R(INS_sgr, EA_8BYTE, REG_SPBASE, alignedReg); + emit->emitIns_R_R_I(INS_stg, EA_8BYTE, tmpReg, REG_SPBASE, 0); + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); + } + + 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 */ From faacf0efc4ddba3cd580bef89980a575f1dc9afa Mon Sep 17 00:00:00 2001 From: Krishna-13-cyber Date: Wed, 4 Nov 2026 11:08:21 +0530 Subject: [PATCH 2/4] Fix genLclHeap to allocate only the aligned size --- src/coreclr/jit/codegens390x.cpp | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/coreclr/jit/codegens390x.cpp b/src/coreclr/jit/codegens390x.cpp index fa0f32f678c338..a07b56a33074f3 100644 --- a/src/coreclr/jit/codegens390x.cpp +++ b/src/coreclr/jit/codegens390x.cpp @@ -5899,8 +5899,8 @@ void CodeGen::genLclHeap(GenTree* tree) return; } - ssize_t total = (ssize_t)AlignUp(amount, STACK_ALIGN) + frameBelow; - if (total > 0x7FFFF) + ssize_t alignedSize = (ssize_t)AlignUp(amount, STACK_ALIGN); + if (alignedSize > 0x7FFFF) { // TODO: Have to handle the overflow } @@ -5908,7 +5908,7 @@ void CodeGen::genLclHeap(GenTree* tree) regNumber tmpReg = internalRegisters.GetSingle(tree); emit->emitIns_R_R_I(INS_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); - emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, REG_SPBASE, REG_SPBASE, -total); + emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, REG_SPBASE, REG_SPBASE, -alignedSize); emit->emitIns_R_R_I(INS_stg, EA_8BYTE, tmpReg, REG_SPBASE, 0); emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); } @@ -5918,10 +5918,11 @@ void CodeGen::genLclHeap(GenTree* tree) regNumber alignedReg = internalRegisters.Extract(tree); regNumber tmpReg = internalRegisters.Extract(tree); - // alignedReg = ((size + 7) & -8) + frameBelow + // alignedReg = ((size + 7) & -8) emit->emitIns_R_R(INS_lgr, EA_8BYTE, alignedReg, sizeReg); - emit->emitIns_R_R(INS_ngr, EA_8BYTE, alignedReg, tmpReg); - emit->emitIns_R_R_I(INS_agfi, EA_8BYTE, alignedReg, alignedReg, frameBelow); + 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_I(INS_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); emit->emitIns_R_R(INS_sgr, EA_8BYTE, REG_SPBASE, alignedReg); From 0f3d5caaae9f8755ca56d31fe95a2b18ea0d075e Mon Sep 17 00:00:00 2001 From: Krishna-13-cyber Date: Tue, 10 Nov 2026 23:14:10 +0530 Subject: [PATCH 3/4] Add zeroing for the newly allocated memory and checks for correctness --- src/coreclr/jit/codegens390x.cpp | 86 ++++++++++++++++++++++++-------- src/coreclr/jit/lsras390x.cpp | 8 ++- 2 files changed, 69 insertions(+), 25 deletions(-) diff --git a/src/coreclr/jit/codegens390x.cpp b/src/coreclr/jit/codegens390x.cpp index a07b56a33074f3..41bdffee4fa829 100644 --- a/src/coreclr/jit/codegens390x.cpp +++ b/src/coreclr/jit/codegens390x.cpp @@ -5880,10 +5880,9 @@ void CodeGen::genLclHeap(GenTree* tree) emitter* emit = GetEmitter(); regNumber targetReg = tree->GetRegNum(); - if (compiler->info.compInitMem) - { - // TODO: Have to handle zeroing of memory - } + 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; @@ -5899,35 +5898,82 @@ void CodeGen::genLclHeap(GenTree* tree) return; } - ssize_t alignedSize = (ssize_t)AlignUp(amount, STACK_ALIGN); - if (alignedSize > 0x7FFFF) + // 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 { - // TODO: Have to handle the overflow + // 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); } - regNumber tmpReg = internalRegisters.GetSingle(tree); + // 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_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); - emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, REG_SPBASE, REG_SPBASE, -alignedSize); - emit->emitIns_R_R_I(INS_stg, EA_8BYTE, tmpReg, REG_SPBASE, 0); emit->emitIns_R_R_I(INS_lay, EA_PTRSIZE, targetReg, REG_SPBASE, frameBelow); } else { - regNumber sizeReg = genConsumeReg(size); - regNumber alignedReg = internalRegisters.Extract(tree); - regNumber tmpReg = internalRegisters.Extract(tree); + // 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_R(INS_lgr, EA_8BYTE, alignedReg, sizeReg); + // 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_I(INS_lgfi, EA_8BYTE, tmpReg, -8); + emit->emitIns_R_R(INS_ngr, EA_8BYTE, alignedReg, tmpReg); - emit->emitIns_R_R_I(INS_lg, EA_8BYTE, tmpReg, REG_SPBASE, 0); emit->emitIns_R_R(INS_sgr, EA_8BYTE, REG_SPBASE, alignedReg); - emit->emitIns_R_R_I(INS_stg, EA_8BYTE, tmpReg, REG_SPBASE, 0); + + // 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); 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()) From c052ed71ae41b6aee7ff18a3ab4935dc3d57e7c3 Mon Sep 17 00:00:00 2001 From: Krishna-13-cyber Date: Tue, 10 Nov 2026 23:24:37 +0530 Subject: [PATCH 4/4] Add test-cases for memory allocations on stack(Stackalloc) --- src/s390tests/tests/Allocation.cs | 44 ++++++++ src/s390tests/tests/Allocation_2.cs | 161 ++++++++++++++++++++++++++++ 2 files changed, 205 insertions(+) create mode 100644 src/s390tests/tests/Allocation.cs create mode 100644 src/s390tests/tests/Allocation_2.cs 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; + } +}