Optimize PI and RPI transitions

This changes the managed ABI in the following ways:

1) The PI transition frames becomes simply the shadow stack top. This
   "frame" is zero-sized - we don't store the current thread in it,
   since the WASM TLS model allows us to elide it.

   The obvious benefit from this is that the PI path is now almost
   100% optimal: two stores and one load. We can get rid of the load
   in an ST build as well, but that's left for a future change.

2) The RPI transition frame is now always allocated at a zero offset
   and returned directly from the RPI helper. We thus elide the intermediate
   state where we already have the shadow stack, but haven't yet attached
   the thread. This brings us in line with other targets.

3) The sparse virtual unwind frame is now allocated right after
   the RPI frame, and "combined" RPI helpers introduced that both
   effect the transition and push the EH frame.

The RPI changes reduce the number of helper calls that any RPI method
needs to make from 3 to 1 (for epilogs - 2 to 1) in the sparse virtual
unwinding model, and reduce the number of intructions on the critical
path.

Benchmarks:
  Node base:
    Bench_PInvoke took               : 86 ms (8.64 ns / op)
    Bench_ReversePInvoke_Empty took  : 113 ms (11.30 ns / op)
    Bench_ReversePInvoke_WithEH took : 172 ms (17.24 ns / op)

  Node diff:
    Bench_PInvoke took               : 81 ms (8.06 ns / op)
    Bench_ReversePInvoke_Empty took  : 58 ms (5.81 ns / op)
    Bench_ReversePInvoke_WithEH took : 108 ms (10.79 ns / op)

  Wasmtime base:
    Bench_PInvoke took               : 99 ms (9.86 ns / op)
    Bench_ReversePInvoke_Empty took  : 73 ms (7.28 ns / op)
    Bench_ReversePInvoke_WithEH took : 77 ms (7.71 ns / op)

  Wasmtime diff:
    Bench_PInvoke took               : 82 ms (8.16 ns / op)
    Bench_ReversePInvoke_Empty took  : 31 ms (3.06 ns / op)
    Bench_ReversePInvoke_WithEH took : 50 ms (4.98 ns / op)

Author: SingleAccretion

src/coreclr/inc/corinfo.h

@@ -595,14 +595,14 @@ enum CorInfoHelpFunc
     CORINFO_HELP_VALIDATE_INDIRECT_CALL,    // CFG: Validate function pointer
     CORINFO_HELP_DISPATCH_INDIRECT_CALL,    // CFG: Validate and dispatch to pointer
 
-    CORINFO_HELP_LLVM_GET_OR_INIT_SHADOW_STACK_TOP,
     CORINFO_HELP_LLVM_EH_CATCH,
     CORINFO_HELP_LLVM_EH_POP_UNWOUND_VIRTUAL_FRAMES,
     CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME,
+    CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_ENTER_AND_PUSH_VIRTUAL_UNWIND_FRAME,
     CORINFO_HELP_LLVM_EH_POP_VIRTUAL_UNWIND_FRAME,
+    CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_EXIT_AND_POP_VIRTUAL_UNWIND_FRAME,
     CORINFO_HELP_LLVM_EH_UNHANDLED_EXCEPTION,
     CORINFO_HELP_LLVM_RESOLVE_INTERFACE_CALL_TARGET,
-    CORINFO_HELP_LLVM_GET_EXTERNAL_CALL_TARGET,
     CORINFO_HELP_LLVM_STRESS_GC,
 
     CORINFO_HELP_COUNT,

src/coreclr/inc/jithelpers.h

@@ -343,14 +343,14 @@
     JITHELPER(CORINFO_HELP_DISPATCH_INDIRECT_CALL, NULL, METHOD__NIL)
 #endif
 
-    JITHELPER(CORINFO_HELP_LLVM_GET_OR_INIT_SHADOW_STACK_TOP, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_EH_CATCH, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_EH_POP_UNWOUND_VIRTUAL_FRAMES, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME, NULL, METHOD__NIL)
+    JITHELPER(CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_ENTER_AND_PUSH_VIRTUAL_UNWIND_FRAME, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_EH_POP_VIRTUAL_UNWIND_FRAME, NULL, METHOD__NIL)
+    JITHELPER(CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_EXIT_AND_POP_VIRTUAL_UNWIND_FRAME, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_EH_UNHANDLED_EXCEPTION, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_RESOLVE_INTERFACE_CALL_TARGET, NULL, METHOD__NIL)
-    JITHELPER(CORINFO_HELP_LLVM_GET_EXTERNAL_CALL_TARGET, NULL, METHOD__NIL)
     JITHELPER(CORINFO_HELP_LLVM_STRESS_GC, JIT_StressGC, METHOD__NIL)
 #undef JITHELPER
 #undef DYNAMICJITHELPER

src/coreclr/jit/flowgraph.cpp

@@ -1740,6 +1740,7 @@ void Compiler::fgAddReversePInvokeEnterExit()
     LclVarDsc* varDsc = lvaGetDesc(lvaReversePInvokeFrameVar);
     lvaSetStruct(lvaReversePInvokeFrameVar, typGetBlkLayout(eeGetEEInfo()->sizeOfReversePInvokeFrame), false);
 
+#ifndef TARGET_WASM // WASM RPI helpers have special ABI and are inserted in lowering.
     // Add enter pinvoke exit callout at the start of prolog
 
     GenTree* pInvokeFrameVar = gtNewLclVarAddrNode(lvaReversePInvokeFrameVar);
@@ -1804,6 +1805,7 @@ void Compiler::fgAddReversePInvokeEnterExit()
         printf("\n");
     }
 #endif
+#endif // !TARGET_WASM
 }
 
 /*****************************************************************************

src/coreclr/jit/llvm.cpp

@@ -517,11 +517,11 @@ bool Llvm::helperCallMayVirtuallyUnwind(CorInfoHelpFunc helperFunc) const
         { FUNC(CORINFO_HELP_THROW_ENTRYPOINT_NOT_FOUND_EXCEPTION) },
 
         // [R]PI helpers, implemented in "Runtime\thread.cpp".
-        { FUNC(CORINFO_HELP_JIT_PINVOKE_BEGIN) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_SS_ARG | HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
-        { FUNC(CORINFO_HELP_JIT_PINVOKE_END) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
-        { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_SS_ARG },
+        { FUNC(CORINFO_HELP_JIT_PINVOKE_BEGIN) CORINFO_TYPE_VOID, { }, HFIF_SS_ARG | HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_JIT_PINVOKE_END) CORINFO_TYPE_VOID, { }, HFIF_SS_ARG | HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER) CORINFO_TYPE_PTR, { }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
         { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_ENTER_TRACK_TRANSITIONS) },
-        { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_EXIT) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_EXIT) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
         { FUNC(CORINFO_HELP_JIT_REVERSE_PINVOKE_EXIT_TRACK_TRANSITIONS) },
 
         // Implemented in "CoreLib\src\System\Runtime\TypeLoaderExports.cs".
@@ -544,14 +544,15 @@ bool Llvm::helperCallMayVirtuallyUnwind(CorInfoHelpFunc helperFunc) const
         { FUNC(CORINFO_HELP_VALIDATE_INDIRECT_CALL) },
         { FUNC(CORINFO_HELP_DISPATCH_INDIRECT_CALL) },
 
-        { FUNC(CORINFO_HELP_LLVM_GET_OR_INIT_SHADOW_STACK_TOP) CORINFO_TYPE_PTR, { }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+            // TODO-LLVM-QPI: incorrect SS top args for RPI return helper.
         { FUNC(CORINFO_HELP_LLVM_EH_CATCH) CORINFO_TYPE_CLASS, { CORINFO_TYPE_NATIVEUINT }, HFIF_SS_ARG },
         { FUNC(CORINFO_HELP_LLVM_EH_POP_UNWOUND_VIRTUAL_FRAMES) CORINFO_TYPE_VOID, { }, HFIF_SS_ARG },
-        { FUNC(CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_NATIVEUINT }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND},
-        { FUNC(CORINFO_HELP_LLVM_EH_POP_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_VOID, { }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND},
+        { FUNC(CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR, CORINFO_TYPE_PTR, CORINFO_TYPE_NATIVEUINT }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_ENTER_AND_PUSH_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_PTR, { CORINFO_TYPE_PTR, CORINFO_TYPE_NATIVEUINT }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_LLVM_EH_POP_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_VOID, { }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
+        { FUNC(CORINFO_HELP_LLVM_EH_REVERSE_PINVOKE_EXIT_AND_POP_VIRTUAL_UNWIND_FRAME) CORINFO_TYPE_VOID, { CORINFO_TYPE_PTR }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
         { FUNC(CORINFO_HELP_LLVM_EH_UNHANDLED_EXCEPTION) CORINFO_TYPE_VOID, { CORINFO_TYPE_CLASS }, HFIF_SS_ARG },
         { FUNC(CORINFO_HELP_LLVM_RESOLVE_INTERFACE_CALL_TARGET) CORINFO_TYPE_PTR, { CORINFO_TYPE_CLASS, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
-        { FUNC(CORINFO_HELP_LLVM_GET_EXTERNAL_CALL_TARGET) CORINFO_TYPE_PTR, { }, HFIF_NO_RPI_OR_GC | HFIF_NO_VIRTUAL_UNWIND },
         { FUNC(CORINFO_HELP_LLVM_STRESS_GC) CORINFO_TYPE_BYREF, { CORINFO_TYPE_BYREF, CORINFO_TYPE_PTR }, HFIF_SS_ARG },
     };
     // clang-format on

src/coreclr/jit/llvm.h

@@ -308,6 +308,7 @@ class Llvm
 
     // Shared between virtual unwind frame insertion and LSSA.
     unsigned m_initialUnwindIndex = UNWIND_INDEX_NONE;
+    CORINFO_GENERIC_HANDLE m_ehInfoSymbol = nullptr;
 
     // Shared between unwind index insertion and EH codegen.
     ArrayStack<unsigned>* m_unwindIndexMap = nullptr;

src/coreclr/jit/llvmlower.cpp

@@ -243,6 +243,8 @@ void Llvm::lowerBlock(BasicBlock* block)
 
 void Llvm::lowerRange(BasicBlock* block, LIR::Range& range)
 {
+    BasicBlock* savedBlock = m_currentBlock;
+    LIR::Range* savedRange = m_currentRange;
     m_currentBlock = block;
     m_currentRange = ⦥
 
@@ -253,8 +255,8 @@ void Llvm::lowerRange(BasicBlock* block, LIR::Range& range)
 
     INDEBUG(range.CheckLIR(_compiler, /* checkUnusedValues */ true));
 
-    m_currentBlock = nullptr;
-    m_currentRange = nullptr;
+    m_currentBlock = savedBlock;
+    m_currentRange = savedRange;
 }
 
 void Llvm::lowerNode(GenTree* node)
@@ -795,21 +797,19 @@ void Llvm::lowerUnmanagedCall(GenTreeCall* callNode)
     // two or more consecutive PI calls.
     if (!callNode->IsSuppressGCTransition())
     {
+        // TODO-LLVM-Upstream: don't allocate lvaInlinedPInvokeFrameVar (its size is zero).
         assert(_compiler->opts.ShouldUsePInvokeHelpers()); // No inline transition support yet.
         assert(_compiler->lvaInlinedPInvokeFrameVar != BAD_VAR_NUM);
 
         // Insert CORINFO_HELP_JIT_PINVOKE_BEGIN.
-        GenTreeLclFld* frameAddr = _compiler->gtNewLclVarAddrNode(_compiler->lvaInlinedPInvokeFrameVar);
-        GenTreeCall* helperCall = _compiler->gtNewHelperCallNode(CORINFO_HELP_JIT_PINVOKE_BEGIN, TYP_VOID, frameAddr);
-        CurrentRange().InsertBefore(callNode, frameAddr, helperCall);
-        lowerNode(frameAddr);
+        GenTreeCall* helperCall = _compiler->gtNewHelperCallNode(CORINFO_HELP_JIT_PINVOKE_BEGIN, TYP_VOID);
+        CurrentRange().InsertBefore(callNode, helperCall);
         lowerNode(helperCall);
 
         // Insert CORINFO_HELP_JIT_PINVOKE_END. No need to explicitly lower the call/local address as the
         // normal lowering loop will pick them up.
-        frameAddr = _compiler->gtNewLclVarAddrNode(_compiler->lvaInlinedPInvokeFrameVar);
-        helperCall = _compiler->gtNewHelperCallNode(CORINFO_HELP_JIT_PINVOKE_END, TYP_VOID, frameAddr);
-        CurrentRange().InsertAfter(callNode, frameAddr, helperCall);
+        helperCall = _compiler->gtNewHelperCallNode(CORINFO_HELP_JIT_PINVOKE_END, TYP_VOID);
+        CurrentRange().InsertAfter(callNode, helperCall);
     }
 
     if (callNode->gtCallType != CT_INDIRECT)
@@ -1598,20 +1598,28 @@ bool Llvm::addVirtualUnwindFrameForExceptionHandling()
                 CORINFO_GENERIC_HANDLE ehInfoSymbol =
                     m_llvm->GetSparseVirtualUnwindInfo(&clauses.BottomRef(), clauses.Height());
 
-                GenTree* ehInfoNode =
-                    m_compiler->gtNewIconHandleNode(reinterpret_cast<size_t>(ehInfoSymbol), GTF_ICON_CONST_PTR);
-                GenTree* unwindFrameLclAddr = m_compiler->gtNewLclVarAddrNode(unwindFrameLclNum);
-                GenTreeIntCon* initialUnwindIndexNode = m_compiler->gtNewIconNode(m_initialIndexValue, TYP_I_IMPL);
-                GenTreeCall* initializeCall =
-                    m_compiler->gtNewHelperCallNode(CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME, TYP_VOID,
-                        unwindFrameLclAddr, ehInfoNode, initialUnwindIndexNode);
-
-                LIR::Range initRange;
-                initRange.InsertAtEnd(unwindFrameLclAddr);
-                initRange.InsertAtEnd(ehInfoNode);
-                initRange.InsertAtEnd(initialUnwindIndexNode);
-                initRange.InsertAtEnd(initializeCall);
-                m_llvm->lowerAndInsertIntoFirstBlock(std::move(initRange));
+                // For frames with an RPI transition, we will use RPI helpers that combine the transitions with unwind
+                // frame linking. This is not just an optimization - the helpers hardcode where the RPI frame is located
+                // and the presence of an unwind frame, itself hard-set to be at a zero offset, affects that.
+                if (!m_compiler->opts.IsReversePInvoke())
+                {
+                    GenTree* ehInfoNode =
+                        m_compiler->gtNewIconHandleNode(reinterpret_cast<size_t>(ehInfoSymbol), GTF_ICON_CONST_PTR);
+                    GenTree* unwindFrameLclAddr = m_compiler->gtNewLclVarAddrNode(unwindFrameLclNum);
+                    GenTreeIntCon* initialUnwindIndexNode = m_compiler->gtNewIconNode(m_initialIndexValue, TYP_I_IMPL);
+                    GenTreeCall* initializeCall =
+                        m_compiler->gtNewHelperCallNode(CORINFO_HELP_LLVM_EH_PUSH_VIRTUAL_UNWIND_FRAME, TYP_VOID,
+                            unwindFrameLclAddr, ehInfoNode, initialUnwindIndexNode);
+
+                    LIR::Range initRange;
+                    initRange.InsertAtEnd(unwindFrameLclAddr);
+                    initRange.InsertAtEnd(ehInfoNode);
+                    initRange.InsertAtEnd(initialUnwindIndexNode);
+                    initRange.InsertAtEnd(initializeCall);
+                    m_llvm->lowerAndInsertIntoFirstBlock(std::move(initRange));
+                }
+
+                m_llvm->m_ehInfoSymbol = ehInfoSymbol;
                 m_llvm->m_sparseVirtualUnwindFrameLclNum = unwindFrameLclNum;
             }
 
@@ -1630,7 +1638,7 @@ bool Llvm::addVirtualUnwindFrameForExceptionHandling()
             }
 
             // Explicit pops are only needed for explicitly linked (via TLS) sparse frames.
-            if (m_llvm->m_sparseVirtualUnwindFrameLclNum != BAD_VAR_NUM)
+            if ((m_llvm->m_sparseVirtualUnwindFrameLclNum != BAD_VAR_NUM) && !m_compiler->opts.IsReversePInvoke())
             {
                 for (BasicBlock* block : m_compiler->Blocks())
                 {