JIT: Clean up and optimize StackLevelSetter (#112561)

The stack levels computed are not actually used for anything outside
x86, so avoid computing them in that case. Also, the max stack level
does not actually get used by the backend, even for x86, so no need to
save it in `Compiler`.

Author: jakobbotsch

src/coreclr/jit/block.cpp

@@ -1432,8 +1432,8 @@ bool BasicBlock::endsWithJmpMethod(Compiler* comp) const
 //
 bool BasicBlock::endsWithTailCallOrJmp(Compiler* comp, bool fastTailCallsOnly /*=false*/) const
 {
-    GenTree* tailCall                       = nullptr;
-    bool     tailCallsConvertibleToLoopOnly = false;
+    GenTreeCall* tailCall                       = nullptr;
+    bool         tailCallsConvertibleToLoopOnly = false;
     return endsWithJmpMethod(comp) ||
            endsWithTailCall(comp, fastTailCallsOnly, tailCallsConvertibleToLoopOnly, &tailCall);
 }
@@ -1454,10 +1454,10 @@ bool BasicBlock::endsWithTailCallOrJmp(Compiler* comp, bool fastTailCallsOnly /*
 // Notes:
 //    At most one of fastTailCallsOnly and tailCallsConvertibleToLoopOnly flags can be true.
 //
-bool BasicBlock::endsWithTailCall(Compiler* comp,
-                                  bool      fastTailCallsOnly,
-                                  bool      tailCallsConvertibleToLoopOnly,
-                                  GenTree** tailCall) const
+bool BasicBlock::endsWithTailCall(Compiler*     comp,
+                                  bool          fastTailCallsOnly,
+                                  bool          tailCallsConvertibleToLoopOnly,
+                                  GenTreeCall** tailCall) const
 {
     assert(!fastTailCallsOnly || !tailCallsConvertibleToLoopOnly);
     *tailCall   = nullptr;
@@ -1524,7 +1524,7 @@ bool BasicBlock::endsWithTailCall(Compiler* comp,
 // Return Value:
 //    true if the block ends with a tail call convertible to loop.
 //
-bool BasicBlock::endsWithTailCallConvertibleToLoop(Compiler* comp, GenTree** tailCall) const
+bool BasicBlock::endsWithTailCallConvertibleToLoop(Compiler* comp, GenTreeCall** tailCall) const
 {
     bool fastTailCallsOnly              = false;
     bool tailCallsConvertibleToLoopOnly = true;

src/coreclr/jit/block.h

@@ -1766,14 +1766,14 @@ struct BasicBlock : private LIR::Range
 
     bool endsWithJmpMethod(Compiler* comp) const;
 
-    bool endsWithTailCall(Compiler* comp,
-                          bool      fastTailCallsOnly,
-                          bool      tailCallsConvertibleToLoopOnly,
-                          GenTree** tailCall) const;
+    bool endsWithTailCall(Compiler*     comp,
+                          bool          fastTailCallsOnly,
+                          bool          tailCallsConvertibleToLoopOnly,
+                          GenTreeCall** tailCall) const;
 
     bool endsWithTailCallOrJmp(Compiler* comp, bool fastTailCallsOnly = false) const;
 
-    bool endsWithTailCallConvertibleToLoop(Compiler* comp, GenTree** tailCall) const;
+    bool endsWithTailCallConvertibleToLoop(Compiler* comp, GenTreeCall** tailCall) const;
 
     // Returns the first statement in the statement list of "this" that is
     // not an SSA definition (a lcl = phi(...) store).

src/coreclr/jit/codegenxarch.cpp

@@ -9607,9 +9607,6 @@ void CodeGen::genProfilingEnterCallback(regNumber initReg, bool* pInitRegZeroed)
                       0,           // argSize. Again, we have to lie about it
                       EA_UNKNOWN); // retSize
 
-    // Check that we have place for the push.
-    assert(compiler->fgGetPtrArgCntMax() >= 1);
-
 #if defined(UNIX_X86_ABI)
     // Restoring alignment manually. This is similar to CodeGen::genRemoveAlignmentAfterCall
     GetEmitter()->emitIns_R_I(INS_add, EA_4BYTE, REG_SPBASE, 0x10);
@@ -9688,9 +9685,6 @@ void CodeGen::genProfilingLeaveCallback(unsigned helper)
 #endif
     genEmitHelperCall(helper, argSize, EA_UNKNOWN /* retSize */);
 
-    // Check that we have place for the push.
-    assert(compiler->fgGetPtrArgCntMax() >= 1);
-
 #if defined(UNIX_X86_ABI)
     // Restoring alignment manually. This is similar to CodeGen::genRemoveAlignmentAfterCall
     GetEmitter()->emitIns_R_I(INS_add, EA_4BYTE, REG_SPBASE, 0x10);

src/coreclr/jit/compiler.h

@@ -6703,33 +6703,6 @@ class Compiler
 
     //------------------------- Morphing --------------------------------------
 
-    unsigned fgPtrArgCntMax = 0;
-
-public:
-    //------------------------------------------------------------------------
-    // fgGetPtrArgCntMax: Return the maximum number of pointer-sized stack arguments that calls inside this method
-    // can push on the stack. This value is calculated during morph.
-    //
-    // Return Value:
-    //    Returns fgPtrArgCntMax, that is a private field.
-    //
-    unsigned fgGetPtrArgCntMax() const
-    {
-        return fgPtrArgCntMax;
-    }
-
-    //------------------------------------------------------------------------
-    // fgSetPtrArgCntMax: Set the maximum number of pointer-sized stack arguments that calls inside this method
-    // can push on the stack. This function is used during StackLevelSetter to fix incorrect morph calculations.
-    //
-    void fgSetPtrArgCntMax(unsigned argCntMax)
-    {
-        fgPtrArgCntMax = argCntMax;
-    }
-
-    bool compCanEncodePtrArgCntMax();
-
-private:
     hashBv*               fgAvailableOutgoingArgTemps;
     ArrayStack<unsigned>* fgUsedSharedTemps = nullptr;
 

src/coreclr/jit/compiler.hpp

@@ -3024,20 +3024,6 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 */
 
-inline bool Compiler::compCanEncodePtrArgCntMax()
-{
-#ifdef JIT32_GCENCODER
-    // DDB 204533:
-    // The GC encoding for fully interruptible methods does not
-    // support more than 1023 pushed arguments, so we have to
-    // use a partially interruptible GC info/encoding.
-    //
-    return (fgPtrArgCntMax < MAX_PTRARG_OFS);
-#else // JIT32_GCENCODER
-    return true;
-#endif
-}
-
 /*****************************************************************************
  *
  *  Call the given function pointer for all nodes in the tree. The 'visitor'

src/coreclr/jit/flowgraph.cpp

@@ -4022,7 +4022,7 @@ PhaseStatus Compiler::fgSetBlockOrder()
     BasicBlock::s_nMaxTrees = 0;
 #endif
 
-    if (compCanEncodePtrArgCntMax() && fgHasCycleWithoutGCSafePoint())
+    if (fgHasCycleWithoutGCSafePoint())
     {
         JITDUMP("Marking method as fully interruptible\n");
         SetInterruptible(true);

src/coreclr/jit/lower.cpp

@@ -10929,5 +10929,38 @@ void Lowering::FinalizeOutgoingArgSpace()
     comp->lvaOutgoingArgSpaceSize = m_outgoingArgSpaceSize;
     comp->lvaOutgoingArgSpaceSize.MarkAsReadOnly();
     comp->lvaGetDesc(comp->lvaOutgoingArgSpaceVar)->GrowBlockLayout(comp->typGetBlkLayout(m_outgoingArgSpaceSize));
+
+    SetFramePointerFromArgSpaceSize();
 #endif
 }
+
+//----------------------------------------------------------------------------------------------
+// Lowering::SetFramePointerFromArgSpaceSize:
+//   Set the frame pointer from the arg space size. This is a quirk because
+//   StackLevelSetter used to do this even outside x86.
+//
+void Lowering::SetFramePointerFromArgSpaceSize()
+{
+    unsigned stackLevelSpace = m_outgoingArgSpaceSize;
+
+    if (comp->compTailCallUsed)
+    {
+        // StackLevelSetter also used to count tailcalls.
+        for (BasicBlock* block : comp->Blocks())
+        {
+            GenTreeCall* tailCall;
+            if (block->endsWithTailCall(comp, true, false, &tailCall))
+            {
+                stackLevelSpace = max(stackLevelSpace, tailCall->gtArgs.OutgoingArgsStackSize());
+            }
+        }
+    }
+
+    unsigned stackLevel =
+        (max(stackLevelSpace, (unsigned)MIN_ARG_AREA_FOR_CALL) - MIN_ARG_AREA_FOR_CALL) / TARGET_POINTER_SIZE;
+
+    if (stackLevel >= 4)
+    {
+        comp->codeGen->setFramePointerRequired(true);
+    }
+}

src/coreclr/jit/lower.h

@@ -45,6 +45,7 @@ class Lowering final : public Phase
     }
 
     void FinalizeOutgoingArgSpace();
+    void SetFramePointerFromArgSpaceSize();
 
 private:
     // LowerRange handles new code that is introduced by or after Lowering.

src/coreclr/jit/morph.cpp

@@ -12796,10 +12796,10 @@ void Compiler::fgMorphStmts(BasicBlock* block)
     }
 
 #if FEATURE_FASTTAILCALL
-    GenTree* recursiveTailCall = nullptr;
+    GenTreeCall* recursiveTailCall = nullptr;
     if (block->endsWithTailCallConvertibleToLoop(this, &recursiveTailCall))
     {
-        fgMorphRecursiveFastTailCallIntoLoop(block, recursiveTailCall->AsCall());
+        fgMorphRecursiveFastTailCallIntoLoop(block, recursiveTailCall);
     }
 #endif
 

src/coreclr/jit/stacklevelsetter.cpp

@@ -24,28 +24,15 @@ StackLevelSetter::StackLevelSetter(Compiler* compiler)
 }
 
 //------------------------------------------------------------------------
-// DoPhase: Calculate stack slots numbers for outgoing args.
+// DoPhase: Calculate stack slots numbers for outgoing args and compute
+// requirements of throw helper blocks.
 //
 // Returns:
 //   PhaseStatus indicating what, if anything, was changed.
 //
-// Notes:
-//   For non-x86 platforms it calculates the max number of slots
-//   that calls inside this method can push on the stack.
-//   This value is used for sanity checks in the emitter.
-//
-//   Stack slots are pointer-sized: 4 bytes for 32-bit platforms, 8 bytes for 64-bit platforms.
-//
-//   For x86 it also sets throw-helper blocks incoming stack depth and set
-//   framePointerRequired when it is necessary. These values are used to pop
-//   pushed args when an exception occurs.
-//
 PhaseStatus StackLevelSetter::DoPhase()
 {
-    for (BasicBlock* const block : comp->Blocks())
-    {
-        ProcessBlock(block);
-    }
+    ProcessBlocks();
 
 #if !FEATURE_FIXED_OUT_ARGS
     if (framePointerRequired)
@@ -56,7 +43,6 @@ PhaseStatus StackLevelSetter::DoPhase()
 
     CheckAdditionalArgs();
 
-    comp->fgSetPtrArgCntMax(maxStackLevel);
     CheckArgCnt();
 
     // When optimizing, check if there are any unused throw helper blocks,
@@ -109,7 +95,28 @@ PhaseStatus StackLevelSetter::DoPhase()
 }
 
 //------------------------------------------------------------------------
-// ProcessBlock: Do stack level calculations for one block.
+// ProcessBlocks: Process all the blocks if necessary.
+//
+void StackLevelSetter::ProcessBlocks()
+{
+#ifndef TARGET_X86
+    // Outside x86 we do not need to compute pushed/popped stack slots.
+    // However, we do optimize throw-helpers and need to process the blocks for
+    // that, but only when optimizing.
+    if (!throwHelperBlocksUsed || comp->opts.OptimizationDisabled())
+    {
+        return;
+    }
+#endif
+
+    for (BasicBlock* const block : comp->Blocks())
+    {
+        ProcessBlock(block);
+    }
+}
+
+//------------------------------------------------------------------------
+// ProcessBlock: Do stack level and throw helper determinations for one block.
 //
 // Notes:
 //   Block starts and ends with an empty outgoing stack.
@@ -125,10 +132,13 @@ PhaseStatus StackLevelSetter::DoPhase()
 void StackLevelSetter::ProcessBlock(BasicBlock* block)
 {
     assert(currentStackLevel == 0);
+
     LIR::ReadOnlyRange& range = LIR::AsRange(block);
     for (auto i = range.rbegin(); i != range.rend(); ++i)
     {
         GenTree* node = *i;
+
+#ifdef TARGET_X86
         if (node->OperIsPutArgStkOrSplit())
         {
             GenTreePutArgStk* putArg   = node->AsPutArgStk();
@@ -145,6 +155,7 @@ void StackLevelSetter::ProcessBlock(BasicBlock* block)
             call->gtArgs.SetStkSizeBytes(usedStackSlotsCount * TARGET_POINTER_SIZE);
 #endif // UNIX_X86_ABI
         }
+#endif
 
         if (!throwHelperBlocksUsed)
         {
@@ -418,7 +429,12 @@ void StackLevelSetter::SubStackLevel(unsigned value)
 //
 void StackLevelSetter::CheckArgCnt()
 {
-    if (!comp->compCanEncodePtrArgCntMax())
+#ifdef JIT32_GCENCODER
+    // The GC encoding for fully interruptible methods does not
+    // support more than 1023 pushed arguments, so we have to
+    // use a partially interruptible GC info/encoding.
+    //
+    if (maxStackLevel >= MAX_PTRARG_OFS)
     {
 #ifdef DEBUG
         if (comp->verbose)
@@ -429,6 +445,7 @@ void StackLevelSetter::CheckArgCnt()
 #endif
         comp->SetInterruptible(false);
     }
+
     if (maxStackLevel >= sizeof(unsigned))
     {
 #ifdef DEBUG
@@ -439,6 +456,7 @@ void StackLevelSetter::CheckArgCnt()
 #endif
         comp->codeGen->setFramePointerRequired(true);
     }
+#endif
 }
 
 //------------------------------------------------------------------------