Better version

Author: EgorBo

src/coreclr/jit/codegenxarch.cpp

@@ -2573,6 +2573,9 @@ void CodeGen::genLclHeap(GenTree* tree)
     target_size_t stackAdjustment     = 0;
     target_size_t locAllocStackOffset = 0;
 
+    // Zeroed via BLK that follows this LCLHEAP
+    const bool zeroedViaBlk = tree->gtFlags & GTF_LCLHEAP_ZEROED;
+
     // compute the amount of memory to allocate to properly STACK_ALIGN.
     size_t amount = 0;
     if (size->IsCnsIntOrI())
@@ -2588,6 +2591,12 @@ void CodeGen::genLclHeap(GenTree* tree)
             goto BAILOUT;
         }
 
+        if (zeroedViaBlk)
+        {
+            // it is expected to be already STACK_ALIGN aligned
+            assert((amount % STACK_ALIGN) == 0);
+        }
+
         // 'amount' is the total number of bytes to localloc to properly STACK_ALIGN
         amount = AlignUp(amount, STACK_ALIGN);
     }
@@ -2694,6 +2703,32 @@ void CodeGen::genLclHeap(GenTree* tree)
         assert((amount % STACK_ALIGN) == 0);
         assert((amount % REGSIZE_BYTES) == 0);
 
+        // Whether this LCLHEAP is explicitly zeroed via BLK or not
+        if (zeroedViaBlk)
+        {
+            assert(compiler->info.compInitMem); // why would we zero it with !compInitMem
+            const bool largePage = amount >= compiler->eeGetPageSize();
+            assert(regCnt == REG_NA);
+            if (largePage || (TARGET_POINTER_SIZE == 4))
+            {
+                regCnt = tree->GetSingleTempReg();
+                instGen_Set_Reg_To_Imm(EA_8BYTE, regCnt, amount);
+                // Negate this value before calling the function to adjust the stack (which adds to ESP).
+                inst_RV(INS_NEG, regCnt, TYP_I_IMPL);
+                genStackPointerDynamicAdjustmentWithProbe(regCnt);
+                // lastTouchDelta is dynamic, and can be up to a page. So if we have outgoing arg space,
+                // we're going to assume the worst and probe.
+            }
+            else
+            {
+                // Since the size is less than a page, and we don't need to zero init memory, simply adjust ESP.
+                // ESP might already be in the guard page, so we must touch it BEFORE the alloc, not after.
+                lastTouchDelta = genStackPointerConstantAdjustmentLoopWithProbe(-(ssize_t)amount,
+                                                                                /* trackSpAdjustments */ true);
+            }
+            goto ALLOC_DONE;
+        }
+
         // For small allocations we will generate up to six push 0 inline
         size_t cntRegSizedWords = amount / REGSIZE_BYTES;
         if (compiler->info.compInitMem && (cntRegSizedWords <= 6))

src/coreclr/jit/gentree.h

@@ -479,6 +479,8 @@ enum GenTreeFlags : unsigned int
     GTF_FLD_INITCLASS           = 0x20000000, // GT_FIELD/GT_FIELD_ADDR -- field access requires preceding class/static init helper
     GTF_FLD_TGT_HEAP            = 0x10000000, // GT_FIELD -- same as GTF_IND_TGT_HEAP
 
+    GTF_LCLHEAP_ZEROED          = 0x80000000, // GT_LCLHEAP -- allocation is explicitly zeroed
+
     GTF_INX_RNGCHK              = 0x80000000, // GT_INDEX_ADDR -- this array address should be range-checked
     GTF_INX_ADDR_NONNULL        = 0x40000000, // GT_INDEX_ADDR -- this array address is not null
 

src/coreclr/jit/importer.cpp

@@ -9894,6 +9894,34 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                         // May throw a stack overflow exception. Obviously, we don't want locallocs to be CSE'd.
                         op1->gtFlags |= (GTF_EXCEPT | GTF_DONT_CSE);
 
+#ifdef TARGET_XARCH
+                        // Emit GT_BLK to zero const-sized LCLHEAP on XARCH
+                        if (op2->IsIntegralConst() && info.compInitMem)
+                        {
+                            const ssize_t size = op2->AsIntCon()->IconValue();
+                            if ((size > 0) && (size <= UINT_MAX))
+                            {
+                                // Align LCLHEAP's size so the zeroing via BLK will be more efficient
+                                op2->AsIntCon()->gtIconVal = AlignUp((size_t)op2->AsIntCon()->IconValue(), STACK_ALIGN);
+                                op2->gtFlags |= GTF_DONT_CSE;
+
+                                // Mark as "explicitly zeroed"
+                                op1->gtFlags |= GTF_LCLHEAP_ZEROED;
+
+                                // Spill LCLHEAP to a local
+                                unsigned tmpNum = lvaGrabTemp(true DEBUGARG("spilling LCLHEAP"));
+                                impAssignTempGen(tmpNum, op1);
+                                op1 = gtNewLclvNode(tmpNum, op1->TypeGet());
+
+                                GenTree* blkTree = new (this, GT_BLK)
+                                    GenTreeBlk(GT_BLK, TYP_STRUCT, op1, typGetBlkLayout((unsigned)size));
+                                blkTree = gtNewBlkOpNode(blkTree, gtNewIconNode(0));
+                                impAppendTree(blkTree, CHECK_SPILL_NONE, impCurStmtDI);
+                                op1 = gtClone(op1);
+                            }
+                        }
+#endif
+
                         // Ensure we have stack security for this method.
                         setNeedsGSSecurityCookie();
 

src/coreclr/jit/lsraxarch.cpp

@@ -1818,7 +1818,8 @@ int LinearScan::BuildLclHeap(GenTree* tree)
             // we will generate 'push 0'.
             assert((sizeVal % REGSIZE_BYTES) == 0);
 
-            if (!compiler->info.compInitMem)
+            // explictly zeroed LCLHEAP also needs a regCnt in case of x86 or large page
+            if (!compiler->info.compInitMem || (tree->gtFlags & GTF_LCLHEAP_ZEROED))
             {
 #ifdef TARGET_X86
                 // x86 always needs regCnt.