Stackalloc localloc

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

@@ -1508,6 +1508,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
@@ -1539,9 +1540,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;
         }
@@ -1592,7 +1593,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)
             {
@@ -1641,10 +1642,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;
         }
@@ -1655,7 +1656,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

@@ -2870,6 +2870,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())
@@ -2893,7 +2895,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;
@@ -2918,7 +2920,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
@@ -2939,7 +2941,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
@@ -3013,7 +3015,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

@@ -13233,6 +13233,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
@@ -4568,6 +4570,7 @@ enum class WellKnownArg : unsigned
     SwiftSelf,
     X86TailCallSpecialArg,
     StackArrayLocal,
+    StackArrayElemSize,
 };
 
 #ifdef DEBUG

src/coreclr/jit/helperexpansion.cpp

@@ -2832,17 +2832,28 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
             return false;
     }
 
-    // If this is a local array, the new helper will have an arg for the array's address
+    // If this is a local array, the new helper will have an arg for the array's address or an arg
+    // for the array element size
     //
     CallArg* const stackLocalAddressArg = call->gtArgs.FindWellKnownArg(WellKnownArg::StackArrayLocal);
+    CallArg* const elemSizeArg          = call->gtArgs.FindWellKnownArg(WellKnownArg::StackArrayElemSize);
 
-    if (stackLocalAddressArg == nullptr)
+    if ((stackLocalAddressArg == nullptr) && (elemSizeArg == nullptr))
     {
         return false;
     }
 
-    JITDUMP("Expanding new array helper for stack allocated array at [%06d] in " FMT_BB ":\n", dspTreeID(call),
-            block->bbNum);
+    // If we have an elem size arg, this is intended to be a localloc
+    //
+    // Note we may have figured out the array length after we did the
+    // escape analysis (that is, lengthArg might be a constant), so we
+    // could possibly change this from a localloc to a fixed alloc,
+    // if we could show that was sound.
+    //
+    bool const isLocAlloc = (elemSizeArg != nullptr);
+
+    JITDUMP("Expanding new array helper for stack allocated array at [%06d] %sin " FMT_BB ":\n", dspTreeID(call),
+            isLocAlloc ? " into localloc " : "", block->bbNum);
     DISPTREE(call);
     JITDUMP("\n");
 
@@ -2859,7 +2870,56 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
         }
     }
 
-    GenTree* const stackLocalAddress = stackLocalAddressArg->GetNode();
+    GenTree* const lengthArg         = call->gtArgs.GetArgByIndex(lengthArgIndex)->GetNode();
+    GenTree*       stackLocalAddress = nullptr;
+
+    // Todo -- clone and leave option to make a helper call under some runtime check
+    // for sufficient stack.
+    //
+    if (isLocAlloc)
+    {
+        assert(elemSizeArg != nullptr);
+        assert(stackLocalAddressArg == nullptr);
+        GenTree* const elemSize = elemSizeArg->GetNode();
+        assert(elemSize->IsCnsIntOrI());
+
+        unsigned const locallocTemp   = lvaGrabTemp(true DEBUGARG("localloc stack address"));
+        lvaTable[locallocTemp].lvType = TYP_I_IMPL;
+
+        GenTree* const arrayLength  = gtCloneExpr(lengthArg);
+        GenTree* const baseSize     = gtNewIconNode(OFFSETOF__CORINFO_Array__data, TYP_I_IMPL);
+        GenTree* const payloadSize  = gtNewOperNode(GT_MUL, TYP_I_IMPL, elemSize, arrayLength);
+        GenTree* const totalSize    = gtNewOperNode(GT_ADD, TYP_I_IMPL, baseSize, payloadSize);
+        GenTree* const locallocNode = gtNewOperNode(GT_LCLHEAP, TYP_I_IMPL, totalSize);
+
+        // Allocation might fail. Codegen must zero the allocation
+        //
+        locallocNode->gtFlags |= (GTF_EXCEPT | GTF_LCLHEAP_MUSTINIT);
+
+        GenTree* const   locallocStore = gtNewStoreLclVarNode(locallocTemp, locallocNode);
+        Statement* const locallocStmt  = fgNewStmtFromTree(locallocStore);
+
+        gtUpdateStmtSideEffects(locallocStmt);
+        fgInsertStmtBefore(block, stmt, locallocStmt);
+
+        // Array address is the result of the localloc
+        //
+        stackLocalAddress = gtNewLclVarNode(locallocTemp);
+        compLocallocUsed  = true;
+
+        // We now require a frame pointer
+        //
+        codeGen->setFramePointerRequired(true);
+    }
+    else
+    {
+        assert(elemSizeArg == nullptr);
+        assert(stackLocalAddressArg != nullptr);
+
+        // Array address is the block local we created earlier
+        //
+        stackLocalAddress = stackLocalAddressArg->GetNode();
+    }
 
     // Initialize the array method table pointer.
     //
@@ -2873,7 +2933,6 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
 
     // Initialize the array length.
     //
-    GenTree* const   lengthArg     = call->gtArgs.GetArgByIndex(lengthArgIndex)->GetNode();
     GenTree* const   lengthArgInt  = fgOptimizeCast(gtNewCastNode(TYP_INT, lengthArg, false, TYP_INT));
     GenTree* const   lengthAddress = gtNewOperNode(GT_ADD, TYP_I_IMPL, gtCloneExpr(stackLocalAddress),
                                                    gtNewIconNode(OFFSETOF__CORINFO_Array__length, TYP_I_IMPL));

src/coreclr/jit/jitconfigvalues.h

@@ -674,6 +674,8 @@ RELEASE_CONFIG_INTEGER(JitObjectStackAllocationConditionalEscape, "JitObjectStac
 CONFIG_STRING(JitObjectStackAllocationConditionalEscapeRange, "JitObjectStackAllocationConditionalEscapeRange")
 RELEASE_CONFIG_INTEGER(JitObjectStackAllocationArray, "JitObjectStackAllocationArray", 1)
 RELEASE_CONFIG_INTEGER(JitObjectStackAllocationSize, "JitObjectStackAllocationSize", 528)
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocationLocalloc, "JitObjectStackAllocationLocalloc", 1);
+RELEASE_CONFIG_INTEGER(JitObjectStackAllocationInLoop, "JitObjectStackAllocationInLoop", 1);
 
 RELEASE_CONFIG_INTEGER(JitEECallTimingInfo, "JitEECallTimingInfo", 0)
 

src/coreclr/jit/lsraarm.cpp

@@ -68,7 +68,7 @@ int LinearScan::BuildLclHeap(GenTree* tree)
             {
                 internalIntCount = 0;
             }
-            else if (!compiler->info.compInitMem)
+            else if (!(compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT)))
             {
                 // No need to initialize allocated stack space.
                 if (sizeVal < compiler->eeGetPageSize())

src/coreclr/jit/lsraarm64.cpp

@@ -1183,7 +1183,7 @@ int LinearScan::BuildNode(GenTree* tree)
                     {
                         // Need no internal registers
                     }
-                    else if (!compiler->info.compInitMem)
+                    else if (!(compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT)))
                     {
                         // No need to initialize allocated stack space.
                         if (sizeVal < compiler->eeGetPageSize())
@@ -1202,7 +1202,7 @@ int LinearScan::BuildNode(GenTree* tree)
             else
             {
                 srcCount = 1;
-                if (!compiler->info.compInitMem)
+                if (!(compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_MUSTINIT)))
                 {
                     buildInternalIntRegisterDefForNode(tree);
                     buildInternalIntRegisterDefForNode(tree);