Stackalloc localloc conditional2

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeHelpers.CoreCLR.cs

@@ -328,6 +328,50 @@ public static void EnsureSufficientExecutionStack()
         [MethodImpl(MethodImplOptions.InternalCall)]
         public static extern bool TryEnsureSufficientExecutionStack();
 
+        internal static unsafe bool CanStackAllocate(nuint size, nuint* allocatedInFrame)
+        {
+            nuint stackBase;
+            nuint stackLimit;
+            GetStackBounds(&stackBase, &stackLimit);
+
+            nuint currentStackAddress = (nuint)(&stackBase);
+            if ((stackLimit == 0) ||
+                (stackBase <= stackLimit) ||
+                (currentStackAddress <= stackLimit) ||
+                (currentStackAddress >= stackBase))
+            {
+                // Unknown or unexpected stack bounds: use the heap.
+                return false;
+            }
+
+            nuint stackSize = stackBase - stackLimit;
+            nuint stackUsed = stackBase - currentStackAddress;
+            if (stackUsed >= (stackSize >> 1))
+            {
+                // The stack is already at least half consumed.
+                return false;
+            }
+
+            nuint remainingStack = currentStackAddress - stackLimit;
+            // 1/128th of remaining stack; about 8 KB when 1 MB remains.
+            nuint stackAllocationLimit = remainingStack >> 7;
+            nuint newAllocatedInFrame = *allocatedInFrame + size;
+            if ((newAllocatedInFrame < *allocatedInFrame) ||
+                (size > stackAllocationLimit) ||
+                (newAllocatedInFrame > stackAllocationLimit))
+            {
+                // Overflow, request too large, or this frame has used too much stack.
+                return false;
+            }
+
+            *allocatedInFrame = newAllocatedInFrame;
+
+            return true;
+        }
+
+        [MethodImpl(MethodImplOptions.InternalCall)]
+        private static extern unsafe void GetStackBounds(nuint* stackBase, nuint* stackLimit);
+
         public static object GetUninitializedObject(
             // This API doesn't call any constructors, but the type needs to be seen as constructed.
             // A type is seen as constructed if a constructor is kept.

src/coreclr/inc/corinfo.h

@@ -485,6 +485,7 @@ enum CorInfoHelpFunc
     CORINFO_HELP_TAILCALL,              // Perform a tail call
 
     CORINFO_HELP_GETCURRENTMANAGEDTHREADID,
+    CORINFO_HELP_CAN_STACK_ALLOCATE, // decide whether a conditional localloc is safe
 
     CORINFO_HELP_INIT_PINVOKE_FRAME,   // initialize an inlined PInvoke Frame for the JIT-compiler
 

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* fcb1b400-696c-4425-a8a7-bb082430a217 */
-    0xfcb1b400,
-    0x696c,
-    0x4425,
-    {0xa8, 0xa7, 0xbb, 0x08, 0x24, 0x30, 0xa2, 0x17}
+constexpr GUID JITEEVersionIdentifier = { /* f967451e-74d1-477d-8470-776ff4dc31eb */
+    0xf967451e,
+    0x74d1,
+    0x477d,
+    {0x84, 0x70, 0x77, 0x6f, 0xf4, 0xdc, 0x31, 0xeb}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/inc/jithelpers.h

@@ -225,6 +225,7 @@
 #endif
 
     DYNAMICJITHELPER(CORINFO_HELP_GETCURRENTMANAGEDTHREADID,  NULL, METHOD__ENVIRONMENT__CURRENT_MANAGED_THREAD_ID)
+    DYNAMICJITHELPER(CORINFO_HELP_CAN_STACK_ALLOCATE,  NULL, METHOD__RUNTIME_HELPERS__CAN_STACK_ALLOCATE)
 
     JITHELPER(CORINFO_HELP_INIT_PINVOKE_FRAME,  JIT_InitPInvokeFrame,  METHOD__NIL)
 

src/coreclr/jit/codegenarm.cpp

@@ -391,6 +391,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 = m_compiler->gtMustZeroLocalloc(tree);
+
     // Result of localloc will be returned in regCnt.
     // Also it used as temporary register in code generation
     // for storing allocation size
@@ -476,7 +478,7 @@ void CodeGen::genLclHeap(GenTree* tree)
 
             goto ALLOC_DONE;
         }
-        else if (!m_compiler->info.compInitMem && (amount < m_compiler->eeGetPageSize())) // must be < not <=
+        else if (!initMem && (amount < m_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
@@ -500,7 +502,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     }
 
     // Allocation
-    if (m_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

@@ -3105,15 +3105,18 @@ 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 needsZeroing = m_compiler->info.compInitMem;
+    bool initMem = m_compiler->gtMustZeroLocalloc(tree);
 
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->isContainedIntOrIImmed())
     {
         // The size node being a contained constant means that Lower has taken care of
         // zeroing the memory if compInitMem is true.
-        needsZeroing = false;
+        if (m_compiler->info.compInitMem)
+        {
+            initMem = false;
+        }
 
         // If amount is zero then return null in targetReg
         amount = size->AsIntCon()->IconValue();
@@ -3137,7 +3140,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 (needsZeroing)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -3224,10 +3227,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 (needsZeroing)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -3239,7 +3242,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (needsZeroing)
+    if (initMem)
     {
         BasicBlock* loop = genCreateTempLabel();
 

src/coreclr/jit/codegenloongarch64.cpp

@@ -1455,6 +1455,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 = m_compiler->gtMustZeroLocalloc(tree);
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI())
@@ -1481,9 +1483,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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -1535,7 +1537,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         static_assert(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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             if (stpCount <= 4)
             {
@@ -1582,10 +1584,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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -1597,7 +1599,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (m_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

@@ -1446,6 +1446,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  = m_compiler->gtMustZeroLocalloc(tree);
     const target_size_t pageSize = m_compiler->eeGetPageSize();
 
     // According to RISC-V Privileged ISA page size is 4KiB
@@ -1477,9 +1478,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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             regCnt = targetReg;
         }
@@ -1531,7 +1532,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         static_assert(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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             if (stpCount <= 4)
             {
@@ -1580,10 +1581,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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             regCnt = targetReg;
         }
@@ -1594,7 +1595,7 @@ void CodeGen::genLclHeap(GenTree* tree)
         instGen_Set_Reg_To_Imm(((unsigned int)amount == amount) ? EA_4BYTE : EA_8BYTE, regCnt, amount);
     }
 
-    if (m_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/codegenwasm.cpp

@@ -3237,7 +3237,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     assert(m_compiler->compLocallocUsed);
     assert(isFramePointerUsed());
 
-    bool const     needsZeroing = m_compiler->info.compInitMem;
+    bool const     needsZeroing = m_compiler->gtMustZeroLocalloc(tree);
     GenTree* const size         = tree->AsOp()->gtOp1;
 
     // We reserve this amount of space below any allocation for

src/coreclr/jit/codegenxarch.cpp

@@ -2798,6 +2798,8 @@ void CodeGen::genLclHeap(GenTree* tree)
     target_size_t stackAdjustment     = 0;
     target_size_t locAllocStackOffset = 0;
 
+    bool const initMem = m_compiler->gtMustZeroLocalloc(tree);
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI() && size->isContained())
@@ -2821,7 +2823,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 (m_compiler->info.compInitMem)
+        if (initMem)
         {
             assert(internalRegisters.Count(tree) == 0);
             regCnt = targetReg;
@@ -2846,7 +2848,7 @@ void CodeGen::genLclHeap(GenTree* tree)
 
         inst_RV_IV(INS_add, regCnt, STACK_ALIGN - 1, emitActualTypeSize(type));
 
-        if (m_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
@@ -2867,7 +2869,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     }
 
     bool initMemOrLargeAlloc; // Declaration must be separate from initialization to avoid clang compiler error.
-    initMemOrLargeAlloc = m_compiler->info.compInitMem || (amount >= m_compiler->eeGetPageSize()); // must be >= not >
+    initMemOrLargeAlloc = initMem || (amount >= m_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
@@ -2941,7 +2943,7 @@ void CodeGen::genLclHeap(GenTree* tree)
     // We should not have any temp registers at this point.
     assert(internalRegisters.Count(tree) == 0);
 
-    if (m_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/compiler.h

@@ -3956,6 +3956,15 @@ class Compiler
     bool gtTreeHasLocalRead(GenTree* tree, unsigned lclNum);
     bool gtTreeHasLocalStore(GenTree* tree, unsigned lclNum);
 
+    // Returns true iff the LCLHEAP node "tree" must zero-initialize its
+    // allocation, either because the method requests init-mem semantics or
+    // because the node carries the GTF_LCLHEAP_MUSTINIT flag.
+    bool gtMustZeroLocalloc(GenTree* tree)
+    {
+        assert(tree->OperIs(GT_LCLHEAP));
+        return info.compInitMem || ((tree->gtFlags & GTF_LCLHEAP_MUSTINIT) != 0);
+    }
+
     void gtSetStmtInfo(Statement* stmt);
 
     // Returns "true" iff "node" has any of the side effects in "flags".
@@ -6634,7 +6643,10 @@ class Compiler
     bool fgExpandStaticInitForCall(BasicBlock** pBlock, Statement* stmt, GenTreeCall* call);
 
     PhaseStatus fgExpandStackArrayAllocations();
-    bool fgExpandStackArrayAllocation(BasicBlock* pBlock, Statement* stmt, GenTreeCall* call);
+    bool fgExpandStackArrayAllocation(BasicBlock*  pBlock,
+                                      Statement*   stmt,
+                                      GenTreeCall* call,
+                                      unsigned&    frameRunningTotalLclNum);
 
     PhaseStatus fgVNBasedIntrinsicExpansion();
     bool fgVNBasedIntrinsicExpansionForCall(BasicBlock** pBlock, Statement* stmt, GenTreeCall* call);

src/coreclr/jit/gentree.cpp

@@ -14585,6 +14585,8 @@ const char* Compiler::gtGetWellKnownArgNameForArgMsg(WellKnownArg arg)
             return "tail call";
         case WellKnownArg::StackArrayLocal:
             return "&lcl arr";
+        case WellKnownArg::StackArrayElemSize:
+            return "arr elemsz";
         case WellKnownArg::RuntimeMethodHandle:
             return "meth hnd";
         case WellKnownArg::AsyncExecutionContext:

src/coreclr/jit/gentree.h

@@ -538,6 +538,8 @@ enum GenTreeFlags : unsigned
 
     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
@@ -4813,6 +4815,7 @@ enum class WellKnownArg : unsigned
     SwiftSelf,
     X86TailCallSpecialArg,
     StackArrayLocal,
+    StackArrayElemSize,
     RuntimeMethodHandle,
     AsyncExecutionContext,
     AsyncSynchronizationContext,