Implement the two-pass EH dispatch

src/coreclr/jit/jitconfigvalues.h

@@ -659,6 +659,10 @@ CONFIG_INTEGER(JitDispIns, W("JitDispIns"), 0)
 #endif // defined(TARGET_LOONGARCH64)
 #endif // DEBUG
 
+#ifdef TARGET_WASM
+CONFIG_INTEGER(JitUseDynamicStackForLclHeap, W("JitUseDynamicStackForLclHeap"), 0)
+#endif // TARGET_WASM
+
 CONFIG_INTEGER(JitEnregStructLocals, W("JitEnregStructLocals"), 1) // Allow to enregister locals with struct type.
 
 #undef CONFIG_INTEGER

src/coreclr/jit/llvm.cpp

@@ -166,9 +166,10 @@ bool Llvm::callRequiresShadowStackSave(const GenTreeCall* call) const
 bool Llvm::helperCallRequiresShadowStackSave(CorInfoHelpAnyFunc helperFunc) const
 {
     // Save/restore is needed if the helper doesn't have a shadow stack argument, unless we know it won't call
-    // back into managed code. TODO-LLVM-CQ: mark (make, if required) more helpers "HFIF_NO_RPI_OR_GC".
-    const HelperFuncInfo& helperInfo = getHelperFuncInfo(helperFunc);
-    return !helperInfo.HasFlags(HFIF_SS_ARG) && !helperInfo.HasFlags(HFIF_NO_RPI_OR_GC);
+    // back into managed code or has special semantics. TODO-LLVM-CQ: mark (make, if required) more helpers
+    // "HFIF_NO_RPI_OR_GC".
+    unsigned helperFlags = getHelperFuncInfo(helperFunc).Flags;
+    return (helperFlags & (HFIF_SS_ARG | HFIF_NO_RPI_OR_GC | HFIF_NO_SS_SAVE)) == HFIF_NONE;
 }
 
 bool Llvm::callHasShadowStackArg(const GenTreeCall* call) const
@@ -563,10 +564,11 @@ bool Llvm::helperCallHasManagedCallingConvention(CorInfoHelpAnyFunc helperFunc)
 
         { FUNC(CORINFO_HELP_LLVM_EH_DISPATCHER_CATCH) CORINFO_TYPE_INT, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
         { FUNC(CORINFO_HELP_LLVM_EH_DISPATCHER_FILTER) CORINFO_TYPE_INT, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
-        { FUNC(CORINFO_HELP_LLVM_EH_DISPATCHER_FAULT) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
+        { FUNC(CORINFO_HELP_LLVM_EH_DISPATCHER_FAULT) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_NO_SS_SAVE },
         { FUNC(CORINFO_HELP_LLVM_EH_DISPATCHER_MUTUALLY_PROTECTING) CORINFO_TYPE_INT, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
         { FUNC(CORINFO_HELP_LLVM_EH_UNHANDLED_EXCEPTION) CORINFO_TYPE_VOID, { CORINFO_TYPE_CLASS }, HFIF_SS_ARG },
-
+        { FUNC(CORINFO_HELP_LLVM_DYNAMIC_STACK_ALLOC) CORINFO_TYPE_PTR, { CORINFO_TYPE_INT, CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC },
+        { FUNC(CORINFO_HELP_LLVM_DYNAMIC_STACK_RELEASE) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC },
     };
     // clang-format on
 

src/coreclr/jit/llvm.h

@@ -70,6 +70,8 @@ enum CorInfoHelpLlvmFunc
     CORINFO_HELP_LLVM_EH_DISPATCHER_FAULT,
     CORINFO_HELP_LLVM_EH_DISPATCHER_MUTUALLY_PROTECTING,
     CORINFO_HELP_LLVM_EH_UNHANDLED_EXCEPTION,
+    CORINFO_HELP_LLVM_DYNAMIC_STACK_ALLOC,
+    CORINFO_HELP_LLVM_DYNAMIC_STACK_RELEASE,
     CORINFO_HELP_ANY_COUNT
 };
 
@@ -108,6 +110,7 @@ enum HelperFuncInfoFlags
     HFIF_SS_ARG = 1, // The helper has shadow stack arg.
     HFIF_VAR_ARG = 1 << 1, // The helper has a variable number of args and must be treated specially.
     HFIF_NO_RPI_OR_GC = 1 << 2, // The helper will not call (back) into managed code or trigger GC.
+    HFIF_NO_SS_SAVE = 1 << 3, // This a special helper that does not need shadow stack save.
 };
 
 struct HelperFuncInfo
@@ -206,6 +209,7 @@ class Llvm
 
     unsigned m_unhandledExceptionHandlerIndex = EHblkDsc::NO_ENCLOSING_INDEX;
     Value* m_rootFunctionShadowStackValue = nullptr;
+    bool m_lclHeapUsed = false; // Same as "compLocallocUsed", but calculated in lowering.
 
     // Codegen emit context.
     unsigned m_currentLlvmFunctionIndex = ROOT_FUNC_IDX;
@@ -328,7 +332,7 @@ class Llvm
     void lowerSpillTempsLiveAcrossSafePoints();
     void lowerLocals();
     void populateLlvmArgNums();
-    void assignShadowStackOffsets(std::vector<LclVarDsc*>& shadowStackLocals, unsigned shadowStackParamCount);
+    void assignShadowStackOffsets(std::vector<unsigned>& shadowStackLocals, unsigned shadowStackParamCount);
     void initializeLocalInProlog(unsigned lclNum, GenTree* value);
 
     void insertProlog();
@@ -375,6 +379,8 @@ class Llvm
     unsigned getOriginalShadowFrameSize() const;
     unsigned getCatchArgOffset() const;
 
+    bool doUseDynamicStackForLclHeap();
+
     // ================================================================================================================
     // |                                                   Codegen                                                    |
     // ================================================================================================================

src/coreclr/jit/llvmcodegen.cpp

@@ -572,18 +572,18 @@ void Llvm::generateEHDispatch()
             Value* handlerValue = isReachable(ehDsc->ebdHndBeg) ? getLlvmFunctionForIndex(ehDsc->ebdFuncIndex)
                                                                 : llvm::Constant::getNullValue(getPtrLlvmType());
 
-            if (ehDsc->ebdHandlerType == EH_HANDLER_FILTER)
-            {
-                Value* filterValue = getLlvmFunctionForIndex(ehDsc->ebdFuncIndex - 1);
-                dispatchDestValue = emitHelperCall(CORINFO_HELP_LLVM_EH_DISPATCHER_FILTER, {funcletShadowStackValue,
-                                                   dispatchDataRefValue, handlerValue, filterValue});
-            }
-            else if (ehDsc->ebdHandlerType == EH_HANDLER_CATCH)
+            if (ehDsc->ebdHandlerType == EH_HANDLER_CATCH)
             {
                 Value* typeSymbolRefValue = getOrCreateSymbol(getSymbolHandleForClassToken(ehDsc->ebdTyp));
                 dispatchDestValue = emitHelperCall(CORINFO_HELP_LLVM_EH_DISPATCHER_CATCH, {funcletShadowStackValue,
                                                    dispatchDataRefValue, handlerValue, typeSymbolRefValue});
             }
+            else if (ehDsc->ebdHandlerType == EH_HANDLER_FILTER)
+            {
+                Value* filterValue = getLlvmFunctionForIndex(ehDsc->ebdFuncIndex - 1);
+                dispatchDestValue = emitHelperCall(CORINFO_HELP_LLVM_EH_DISPATCHER_FILTER, {funcletShadowStackValue,
+                                                   dispatchDataRefValue, handlerValue, filterValue});
+            }
             else
             {
                 dispatchDestValue = emitHelperCall(CORINFO_HELP_LLVM_EH_DISPATCHER_FAULT, {funcletShadowStackValue,
@@ -1732,30 +1732,52 @@ void Llvm::buildLclHeap(GenTreeUnOp* lclHeap)
     }
     else
     {
-        llvm::AllocaInst* allocaInst = _builder.CreateAlloca(Type::getInt8Ty(m_context->Context), sizeValue);
+        llvm::BasicBlock* beforeAllocLlvmBlock = nullptr;
+        llvm::BasicBlock* joinLlvmBlock = nullptr;
+        if (!sizeNode->IsIntegralConst())
+        {
+            beforeAllocLlvmBlock = _builder.GetInsertBlock();
+            llvm::BasicBlock* allocLlvmBlock = createInlineLlvmBlock();
+            joinLlvmBlock = createInlineLlvmBlock();
+
+            Value* zeroSizeValue = llvm::Constant::getNullValue(sizeValue->getType());
+            Value* isSizeZeroValue = _builder.CreateICmpEQ(sizeValue, zeroSizeValue);
+            _builder.CreateCondBr(isSizeZeroValue, joinLlvmBlock, allocLlvmBlock);
+            _builder.SetInsertPoint(allocLlvmBlock);
+        }
 
-        // LCLHEAP (aka IL's "localloc") is specified to return a pointer "...aligned so that any built-in data type
-        // can be stored there using the stind instructions", so we'll be a bit conservative and align it maximally.
-        llvm::Align allocaAlignment = llvm::Align(genTypeSize(TYP_DOUBLE));
-        allocaInst->setAlignment(allocaAlignment);
+        // LCLHEAP (aka IL's "localloc") is specified to return a pointer "...aligned so that any built-in
+        // data type can be stored there using the stind instructions"; that means 8 bytes for a double.
+        llvm::Align lclHeapAlignment = llvm::Align(genTypeSize(TYP_DOUBLE));
+
+        if (doUseDynamicStackForLclHeap())
+        {
+            lclHeapValue = emitHelperCall(CORINFO_HELP_LLVM_DYNAMIC_STACK_ALLOC, {sizeValue, getShadowStack()});
+        }
+        else
+        {
+            llvm::AllocaInst* allocaInst = _builder.CreateAlloca(Type::getInt8Ty(m_context->Context), sizeValue);
+            allocaInst->setAlignment(lclHeapAlignment);
+            lclHeapValue = allocaInst;
+        }
 
         // "If the localsinit flag on the method is true, the block of memory returned is initialized to 0".
         if (_compiler->info.compInitMem)
         {
-            _builder.CreateMemSet(allocaInst, _builder.getInt8(0), sizeValue, allocaAlignment);
+            _builder.CreateMemSet(lclHeapValue, _builder.getInt8(0), sizeValue, lclHeapAlignment);
         }
 
-        if (!sizeNode->IsIntegralConst()) // Build: %lclHeapValue = (%sizeValue != 0) ? "alloca" : "null".
+        if (joinLlvmBlock != nullptr)
         {
-            Value* zeroSizeValue = llvm::Constant::getNullValue(sizeValue->getType());
-            Value* isSizeNotZeroValue = _builder.CreateCmp(llvm::CmpInst::ICMP_NE, sizeValue, zeroSizeValue);
-            Value* nullValue = llvm::Constant::getNullValue(getPtrLlvmType());
+            llvm::BasicBlock* allocLlvmBlock = _builder.GetInsertBlock();
+            _builder.CreateBr(joinLlvmBlock);
 
-            lclHeapValue = _builder.CreateSelect(isSizeNotZeroValue, allocaInst, nullValue);
-        }
-        else
-        {
-            lclHeapValue = allocaInst;
+            _builder.SetInsertPoint(joinLlvmBlock);
+            llvm::PHINode* lclHeapPhi = _builder.CreatePHI(lclHeapValue->getType(), 2);
+            lclHeapPhi->addIncoming(lclHeapValue, allocLlvmBlock);
+            lclHeapPhi->addIncoming(llvm::Constant::getNullValue(getPtrLlvmType()), beforeAllocLlvmBlock);
+
+            lclHeapValue = lclHeapPhi;
         }
     }
 
@@ -2188,9 +2210,16 @@ void Llvm::buildReturn(GenTree* node)
 {
     assert(node->OperIs(GT_RETURN, GT_RETFILT));
 
-    if (node->OperIs(GT_RETURN) && _compiler->opts.IsReversePInvoke())
+    if (node->OperIs(GT_RETURN))
     {
-        emitHelperCall(CORINFO_HELP_LLVM_SET_SHADOW_STACK_TOP, getShadowStack());
+        if (m_lclHeapUsed && doUseDynamicStackForLclHeap())
+        {
+            emitHelperCall(CORINFO_HELP_LLVM_DYNAMIC_STACK_RELEASE, getShadowStack());
+        }
+        if (_compiler->opts.IsReversePInvoke())
+        {
+            emitHelperCall(CORINFO_HELP_LLVM_SET_SHADOW_STACK_TOP, getShadowStack());
+        }
     }
 
     if (node->TypeIs(TYP_VOID))

src/coreclr/jit/llvmlower.cpp

@@ -238,6 +238,11 @@ void Llvm::lowerSpillTempsLiveAcrossSafePoints()
 
         for (GenTree* node : blockRange)
         {
+            if (node->OperIs(GT_LCLHEAP))
+            {
+                m_lclHeapUsed = true;
+            }
+
             if (node->isContained())
             {
                 assert(!isPotentialGcSafePoint(node));
@@ -341,7 +346,7 @@ void Llvm::lowerLocals()
 {
     populateLlvmArgNums();
 
-    std::vector<LclVarDsc*> shadowStackLocals;
+    std::vector<unsigned> shadowStackLocals;
     unsigned shadowStackParamCount = 0;
 
     for (unsigned lclNum = 0; lclNum < _compiler->lvaCount; lclNum++)
@@ -413,7 +418,7 @@ void Llvm::lowerLocals()
             if (varDsc->lvIsParam && !isLlvmParam)
             {
                 shadowStackParamCount++;
-                shadowStackLocals.push_back(varDsc);
+                shadowStackLocals.push_back(lclNum);
                 continue;
             }
 
@@ -450,14 +455,28 @@ void Llvm::lowerLocals()
                 }
             }
 
-            shadowStackLocals.push_back(varDsc);
+            shadowStackLocals.push_back(lclNum);
         }
         else
         {
             INDEBUG(varDsc->lvOnFrame = false); // For more accurate frame layout dumping.
         }
     }
 
+    if ((shadowStackLocals.size() == 0) && m_lclHeapUsed && doUseDynamicStackForLclHeap())
+    {
+        // The dynamic stack is tied to the shadow one. If we have an empty shadow frame with a non-empty dynamic one,
+        // an ambiguity in what state must be released on return arises - our caller might have an empty shadow frame
+        // as well, but of course we don't want to release its dynamic state accidentally. To solve this, pad out the
+        // shadow frame in methods that use the dynamic stack if it is empty. The need to do this should be pretty rare
+        // so it is ok to waste a shadow stack slow here.
+        unsigned paddingLclNum = _compiler->lvaGrabTempWithImplicitUse(true DEBUGARG("SS padding for the dynamic stack"));
+        _compiler->lvaGetDesc(paddingLclNum)->lvType = TYP_REF;
+        initializeLocalInProlog(paddingLclNum, _compiler->gtNewIconNode(0, TYP_REF));
+
+        shadowStackLocals.push_back(paddingLclNum);
+    }
+
     assignShadowStackOffsets(shadowStackLocals, shadowStackParamCount);
 }
 
@@ -528,12 +547,17 @@ void Llvm::populateLlvmArgNums()
     _llvmArgCount = nextLlvmArgNum;
 }
 
-void Llvm::assignShadowStackOffsets(std::vector<LclVarDsc*>& shadowStackLocals, unsigned shadowStackParamCount)
+void Llvm::assignShadowStackOffsets(std::vector<unsigned>& shadowStackLocals, unsigned shadowStackParamCount)
 {
     if (_compiler->opts.OptimizationEnabled())
     {
         std::sort(shadowStackLocals.begin() + shadowStackParamCount, shadowStackLocals.end(),
-                  [](const LclVarDsc* lhs, const LclVarDsc* rhs) { return lhs->lvRefCntWtd() > rhs->lvRefCntWtd(); });
+                  [compiler = _compiler](unsigned lhsLclNum, unsigned rhsLclNum)
+        {
+            LclVarDsc* lhsVarDsc = compiler->lvaGetDesc(lhsLclNum);
+            LclVarDsc* rhsVarDsc = compiler->lvaGetDesc(rhsLclNum);
+            return lhsVarDsc->lvRefCntWtd() > rhsVarDsc->lvRefCntWtd();
+        });
     }
 
     unsigned offset = 0;
@@ -567,7 +591,7 @@ void Llvm::assignShadowStackOffsets(std::vector<LclVarDsc*>& shadowStackLocals,
     unsigned assignedShadowStackParamCount = 0;
     for (unsigned i = 0; i < shadowStackLocals.size(); i++)
     {
-        LclVarDsc* varDsc = shadowStackLocals.at(i);
+        LclVarDsc* varDsc = _compiler->lvaGetDesc(shadowStackLocals.at(i));
 
         if (varDsc->lvIsParam && (varDsc->lvLlvmArgNum == BAD_LLVM_ARG_NUM))
         {
@@ -584,7 +608,7 @@ void Llvm::assignShadowStackOffsets(std::vector<LclVarDsc*>& shadowStackLocals,
 
     for (unsigned i = 0; i < shadowStackLocals.size(); i++)
     {
-        LclVarDsc* varDsc = shadowStackLocals.at(i);
+        LclVarDsc* varDsc = _compiler->lvaGetDesc(shadowStackLocals.at(i));
 
         if (!isShadowFrameLocal(varDsc))
         {
@@ -677,7 +701,6 @@ void Llvm::lowerBlock(BasicBlock* block)
     }
 
     INDEBUG(CurrentRange().CheckLIR(_compiler, /* checkUnusedValues */ true));
-
 }
 
 void Llvm::lowerNode(GenTree* node)
@@ -1675,3 +1698,13 @@ unsigned Llvm::getCatchArgOffset() const
 {
     return 0;
 }
+
+bool Llvm::doUseDynamicStackForLclHeap()
+{
+    // TODO-LLVM: add a stress mode.
+    assert(m_lclHeapUsed);
+
+    // We assume LCLHEAPs in methods with EH escape into handlers and so
+    // have to use a special EH-aware allocator instead of the native stack.
+    return _compiler->ehAnyFunclets() || JitConfig.JitUseDynamicStackForLclHeap();
+}