Stackalloc localloc conditional2

src/coreclr/jit/codegen.h

@@ -125,6 +125,7 @@ class CodeGen final : public CodeGenInterface
 
     //-------------------------------------------------------------------------
 
+    bool     genLocallocUsed;  // true if we have used localloc in the method
     bool     genUseBlockInit;  // true if we plan to block-initialize the local stack frame
     unsigned genInitStkLclCnt; // The count of local variables that we need to zero init
 

src/coreclr/jit/codegenarm.cpp

@@ -385,6 +385,8 @@ void CodeGen::genLclHeap(GenTree* tree)
     GenTree* size = tree->AsOp()->gtOp1;
     noway_assert((genActualType(size->gtType) == TYP_INT) || (genActualType(size->gtType) == TYP_I_IMPL));
 
+    bool const initMem = compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT);
+
     // Result of localloc will be returned in regCnt.
     // Also it used as temporary register in code generation
     // for storing allocation size
@@ -470,7 +472,7 @@ void CodeGen::genLclHeap(GenTree* tree)
 
             goto ALLOC_DONE;
         }
-        else if (!compiler->info.compInitMem && (amount < compiler->eeGetPageSize())) // must be < not <=
+        else if (!initMem && (amount < compiler->eeGetPageSize())) // must be < not <=
         {
             // Since the size is less than a page, simply adjust the SP value.
             // The SP might already be in the guard page, must touch it BEFORE
@@ -494,7 +496,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     }
 
     // Allocation
-    if (compiler->info.compInitMem)
+    if (initMem)
     {
         // At this point 'regCnt' is set to the total number of bytes to localloc.
         // Since we have to zero out the allocated memory AND ensure that the stack pointer is always valid

src/coreclr/jit/codegenarm64.cpp

@@ -3155,6 +3155,8 @@ void CodeGen::genLclHeap(GenTree* tree)
     noway_assert(isFramePointerUsed()); // localloc requires Frame Pointer to be established since SP changes
     noway_assert(genStackLevel == 0);   // Can't have anything on the stack
 
+    bool const initMem = compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT);
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI())
@@ -3184,7 +3186,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         // Compute the size of the block to allocate and perform alignment.
         // If compInitMem=true, we can reuse targetReg as regcnt,
         // since we don't need any internal registers.
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -3232,7 +3234,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         static_assert_no_msg(STACK_ALIGN == storePairRegsWritesBytes);
         assert(amount % storePairRegsWritesBytes == 0); // stp stores two registers at a time
 
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             if (amount <= compiler->getUnrollThreshold(Compiler::UnrollKind::Memset))
             {
@@ -3303,10 +3305,10 @@ void CodeGen::genLclHeap(GenTree* tree)
         }
 
         // else, "mov regCnt, amount"
-        // If compInitMem=true, we can reuse targetReg as regcnt.
+        // If initMem=true, we can reuse targetReg as regcnt.
         // Since size is a constant, regCnt is not yet initialized.
         assert(regCnt == REG_NA);
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -3318,7 +3320,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (compiler->info.compInitMem)
+    if (initMem)
     {
         BasicBlock* loop = genCreateTempLabel();
 

src/coreclr/jit/codegenloongarch64.cpp

@@ -1600,6 +1600,8 @@ void CodeGen::genLclHeap(GenTree* tree)
     noway_assert(isFramePointerUsed()); // localloc requires Frame Pointer to be established since SP changes
     noway_assert(genStackLevel == 0);   // Can't have anything on the stack
 
+    bool const initMem = compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT);
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI())
@@ -1626,9 +1628,9 @@ void CodeGen::genLclHeap(GenTree* tree)
         emit->emitIns_J_cond_la(INS_beq, endLabel, targetReg, REG_R0);
 
         // Compute the size of the block to allocate and perform alignment.
-        // If compInitMem=true, we can reuse targetReg as regcnt,
+        // If initMem=true, we can reuse targetReg as regcnt,
         // since we don't need any internal registers.
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -1680,7 +1682,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         static_assert_no_msg(STACK_ALIGN == (REGSIZE_BYTES * 2));
         assert(amount % (REGSIZE_BYTES * 2) == 0); // stp stores two registers at a time
         size_t stpCount = amount / (REGSIZE_BYTES * 2);
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             if (stpCount <= 4)
             {
@@ -1727,10 +1729,10 @@ void CodeGen::genLclHeap(GenTree* tree)
         }
 
         // else, "mov regCnt, amount"
-        // If compInitMem=true, we can reuse targetReg as regcnt.
+        // If initMem=true, we can reuse targetReg as regcnt.
         // Since size is a constant, regCnt is not yet initialized.
         assert(regCnt == REG_NA);
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -1742,7 +1744,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (compiler->info.compInitMem)
+    if (initMem)
     {
         // At this point 'regCnt' is set to the total number of bytes to locAlloc.
         // Since we have to zero out the allocated memory AND ensure that the stack pointer is always valid

src/coreclr/jit/codegenriscv64.cpp

@@ -1523,6 +1523,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     noway_assert(isFramePointerUsed()); // localloc requires Frame Pointer to be established since SP changes
     noway_assert(genStackLevel == 0);   // Can't have anything on the stack
 
+    bool const          initMem  = compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT);
     const target_size_t pageSize = compiler->eeGetPageSize();
 
     // According to RISC-V Privileged ISA page size is 4KiB
@@ -1554,9 +1555,9 @@ void CodeGen::genLclHeap(GenTree* tree)
         emit->emitIns_J_cond_la(INS_beq, endLabel, targetReg, REG_R0);
 
         // Compute the size of the block to allocate and perform alignment.
-        // If compInitMem=true, we can reuse targetReg as regcnt,
+        // If initMem=true, we can reuse targetReg as regcnt,
         // since we don't need any internal registers.
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             regCnt = targetReg;
         }
@@ -1607,7 +1608,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         static_assert_no_msg(STACK_ALIGN == (REGSIZE_BYTES * 2));
         assert(amount % (REGSIZE_BYTES * 2) == 0); // stp stores two registers at a time
         size_t stpCount = amount / (REGSIZE_BYTES * 2);
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             if (stpCount <= 4)
             {
@@ -1656,10 +1657,10 @@ void CodeGen::genLclHeap(GenTree* tree)
         }
 
         // else, "mov regCnt, amount"
-        // If compInitMem=true, we can reuse targetReg as regcnt.
+        // If initMem=true, we can reuse targetReg as regcnt.
         // Since size is a constant, regCnt is not yet initialized.
         assert(regCnt == REG_NA);
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             regCnt = targetReg;
         }
@@ -1670,7 +1671,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (compiler->info.compInitMem)
+    if (initMem)
     {
         // At this point 'regCnt' is set to the total number of bytes to locAlloc.
         // Since we have to zero out the allocated memory AND ensure that the stack pointer is always valid

src/coreclr/jit/codegenxarch.cpp

@@ -2850,6 +2850,7 @@ void CodeGen::genLclHeap(GenTree* tree)
 {
     assert(tree->OperGet() == GT_LCLHEAP);
     assert(compiler->compLocallocUsed);
+    genLocallocUsed = true;
 
     GenTree* size = tree->AsOp()->gtOp1;
     noway_assert((genActualType(size->gtType) == TYP_INT) || (genActualType(size->gtType) == TYP_I_IMPL));
@@ -2871,6 +2872,8 @@ void CodeGen::genLclHeap(GenTree* tree)
     target_size_t stackAdjustment     = 0;
     target_size_t locAllocStackOffset = 0;
 
+    bool const initMem = compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT);
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI() && size->isContained())
@@ -2894,7 +2897,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         // Compute the size of the block to allocate and perform alignment.
         // If compInitMem=true, we can reuse targetReg as regcnt,
         // since we don't need any internal registers.
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -2919,7 +2922,7 @@ void CodeGen::genLclHeap(GenTree* tree)
 
         inst_RV_IV(INS_add, regCnt, STACK_ALIGN - 1, emitActualTypeSize(type));
 
-        if (compiler->info.compInitMem)
+        if (initMem)
         {
             // Convert the count from a count of bytes to a loop count. We will loop once per
             // stack alignment size, so each loop will zero 4 bytes on Windows/x86, and 16 bytes
@@ -2940,7 +2943,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     }
 
     bool initMemOrLargeAlloc; // Declaration must be separate from initialization to avoid clang compiler error.
-    initMemOrLargeAlloc = compiler->info.compInitMem || (amount >= compiler->eeGetPageSize()); // must be >= not >
+    initMemOrLargeAlloc = initMem || (amount >= compiler->eeGetPageSize()); // must be >= not >
 
 #if FEATURE_FIXED_OUT_ARGS
     // If we have an outgoing arg area then we must adjust the SP by popping off the
@@ -3014,7 +3017,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     // We should not have any temp registers at this point.
     assert(internalRegisters.Count(tree) == 0);
 
-    if (compiler->info.compInitMem)
+    if (initMem)
     {
         // At this point 'regCnt' is set to the number of loop iterations for this loop, if each
         // iteration zeros (and subtracts from the stack pointer) STACK_ALIGN bytes.

src/coreclr/jit/gentree.cpp

@@ -13093,6 +13093,9 @@ const char* Compiler::gtGetWellKnownArgNameForArgMsg(WellKnownArg arg)
             return "tail call";
         case WellKnownArg::StackArrayLocal:
             return "&lcl arr";
+        case WellKnownArg::StackArrayElemSize:
+            return "arr elemsz";
+
         default:
             return nullptr;
     }

src/coreclr/jit/gentree.h

@@ -550,6 +550,8 @@ enum GenTreeFlags : unsigned int
 
     GTF_ALLOCOBJ_EMPTY_STATIC = 0x80000000, // GT_ALLOCOBJ -- allocation site is part of an empty static pattern
 
+    GTF_LCLHEAP_MUSTINIT = 0x80000000, // GT_LCLHEAP -- allocation must be zeroed
+
 #ifdef FEATURE_HW_INTRINSICS
     GTF_HW_EM_OP                  = 0x10000000, // GT_HWINTRINSIC -- node is used as an operand to an embedded mask
     GTF_HW_USER_CALL              = 0x20000000, // GT_HWINTRINSIC -- node is implemented via a user call
@@ -4564,6 +4566,7 @@ enum class WellKnownArg : unsigned
     SwiftSelf,
     X86TailCallSpecialArg,
     StackArrayLocal,
+    StackArrayElemSize,
 };
 
 #ifdef DEBUG