[Wasm RyuJit] Call Codegen Fixes for R2R

src/coreclr/jit/codegenwasm.cpp

@@ -2515,6 +2515,9 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
 
     params.wasmSignature = m_compiler->info.compCompHnd->getWasmTypeSymbol(typeStack.Data(), typeStack.Height());
 
+    // A non-null target expression always indicates an indirect call on Wasm,
+    // as currently the only possible result of the target expression would be a
+    // table index which must be used via call_indirect
     if (target != nullptr)
     {
         // Codegen should have already evaluated our target node (last) and pushed it onto the stack,
@@ -2526,52 +2529,28 @@ void CodeGen::genCallInstruction(GenTreeCall* call)
     }
     else
     {
-        // If we have no target and this is a call with indirection cell then
-        // we do an optimization where we load the call address directly from
-        // the indirection cell instead of duplicating the tree. In BuildCall
-        // we ensure that get an extra register for the purpose. Note that for
-        // CFG the call might have changed to
-        // CORINFO_HELP_DISPATCH_INDIRECT_CALL in which case we still have the
-        // indirection cell but we should not try to optimize.
-        WellKnownArg indirectionCellArgKind = WellKnownArg::None;
-        if (!call->IsHelperCall(CORINFO_HELP_DISPATCH_INDIRECT_CALL))
-        {
-            indirectionCellArgKind = call->GetIndirectionCellArgKind();
-        }
+        // Generate a direct call to a non-virtual user defined or helper method
+        assert(call->IsHelperCall() || (call->gtCallType == CT_USER_FUNC));
 
-        if (indirectionCellArgKind != WellKnownArg::None)
-        {
-            assert(call->IsR2ROrVirtualStubRelativeIndir());
+        assert(call->gtEntryPoint.addr == NULL);
 
-            params.callType = EC_INDIR_R;
-            // params.ireg     = targetAddrReg;
-            genEmitCallWithCurrentGC(params);
+        if (call->IsHelperCall())
+        {
+            assert(!call->IsFastTailCall());
+            CorInfoHelpFunc helperNum = m_compiler->eeGetHelperNum(params.methHnd);
+            noway_assert(helperNum != CORINFO_HELP_UNDEF);
+            CORINFO_CONST_LOOKUP helperLookup = m_compiler->compGetHelperFtn(helperNum);
+            assert(helperLookup.accessType == IAT_VALUE);
+            params.addr = helperLookup.addr;
         }
         else
         {
-            // Generate a direct call to a non-virtual user defined or helper method
-            assert(call->IsHelperCall() || (call->gtCallType == CT_USER_FUNC));
-
-            assert(call->gtEntryPoint.addr == NULL);
-
-            if (call->IsHelperCall())
-            {
-                assert(!call->IsFastTailCall());
-                CorInfoHelpFunc helperNum = m_compiler->eeGetHelperNum(params.methHnd);
-                noway_assert(helperNum != CORINFO_HELP_UNDEF);
-                CORINFO_CONST_LOOKUP helperLookup = m_compiler->compGetHelperFtn(helperNum);
-                assert(helperLookup.accessType == IAT_VALUE);
-                params.addr = helperLookup.addr;
-            }
-            else
-            {
-                // Direct call to a non-virtual user function.
-                params.addr = call->gtDirectCallAddress;
-            }
-
-            params.callType = EC_FUNC_TOKEN;
-            genEmitCallWithCurrentGC(params);
+            // Direct call to a non-virtual user function.
+            params.addr = call->gtDirectCallAddress;
         }
+
+        params.callType = EC_FUNC_TOKEN;
+        genEmitCallWithCurrentGC(params);
     }
 }
 
@@ -2676,16 +2655,21 @@ void CodeGen::genEmitHelperCall(unsigned helper, int argSize, emitAttr retSize,
     if (helperIsManaged)
     {
         // Push PEP onto the stack because we are calling a managed helper that expects it as the last parameter.
+        // The helper function address is the address of an indirection cell, so we load from the cell to get the PEP
+        // address to push.
         assert(helperFunction.accessType == IAT_PVALUE);
         GetEmitter()->emitAddressConstant(helperFunction.addr);
+        GetEmitter()->emitIns_I(INS_i32_load, EA_PTRSIZE, 0);
     }
 
     if (params.callType == EC_INDIR_R)
     {
-        // Push the call target onto the wasm evaluation stack by dereferencing the PEP.
+        // Push the call target onto the wasm evaluation stack by dereferencing the indirection cell
+        // and then the PEP pointed to by the indirection cell.
         assert(helperFunction.accessType == IAT_PVALUE);
         GetEmitter()->emitAddressConstant(helperFunction.addr);
         GetEmitter()->emitIns_I(INS_i32_load, EA_PTRSIZE, 0);
+        GetEmitter()->emitIns_I(INS_i32_load, EA_PTRSIZE, 0);
     }
 
     genEmitCallWithCurrentGC(params);

src/coreclr/jit/gentree.h

@@ -5761,7 +5761,7 @@ struct GenTreeCall final : public GenTree
             return WellKnownArg::VirtualStubCell;
         }
 
-#if defined(TARGET_ARMARCH) || defined(TARGET_RISCV64) || defined(TARGET_LOONGARCH64) || defined(TARGET_WASM)
+#if defined(TARGET_ARMARCH) || defined(TARGET_RISCV64) || defined(TARGET_LOONGARCH64)
         // For ARM architectures, we always use an indirection cell for R2R calls.
         if (IsR2RRelativeIndir() && !IsDelegateInvoke())
         {

src/coreclr/jit/lower.cpp

@@ -3022,6 +3022,15 @@ GenTree* Lowering::LowerCall(GenTree* node)
         }
     }
 
+#ifdef TARGET_WASM
+    // For any type of managed call, if we have portable entry points enabled, we need to lower
+    // the call according to the portable entrypoint abi
+    if (!call->IsUnmanaged() && m_compiler->opts.jitFlags->IsSet(JitFlags::JIT_FLAG_PORTABLE_ENTRY_POINTS))
+    {
+        LowerPEPCall(call);
+    }
+#endif // TARGET_WASM
+
     if (varTypeIsStruct(call))
     {
         LowerCallStruct(call);
@@ -3638,6 +3647,73 @@ GenTree* Lowering::LowerTailCallViaJitHelper(GenTreeCall* call, GenTree* callTar
     return result;
 }
 
+#ifdef TARGET_WASM
+//---------------------------------------------------------------------------------------------
+// LowerPEPCall: Lower a call node dispatched through a PortableEntryPoint (PEP)
+//
+// Given a call node with gtControlExpr representing a call target which is the address of a portable entrypoint,
+// this function lowers the call to appropriately dispatch through the portable entrypoint using the Portable
+// entrypoint calling convention.
+// To do this, it:
+//      1. Introduces a new local variable to hold the PEP address
+//      2. Adds a new well-known argument to the call passing this local
+//      3. Rewrites the control expression to indirect through the new local, since for PEP's, the actual call target
+//         must be loaded from the portable entry point address.
+//
+// Arguments:
+//    call         -  The call node to lower. It is expected that the call node has gtControlExpr set to the original
+//                      call target and that the call does not have a PEP arg already.
+//
+// Return Value:
+//    None. The call node is modified in place.
+//
+void Lowering::LowerPEPCall(GenTreeCall* call)
+{
+    JITDUMP("Begin lowering PEP call\n");
+    DISPTREERANGE(BlockRange(), call);
+
+    // PEP call must always have a control expression
+    assert(call->gtControlExpr != nullptr);
+    LIR::Use callTargetUse(BlockRange(), &call->gtControlExpr, call);
+
+    JITDUMP("Creating new local variable for PEP");
+    unsigned int   callTargetLclNum    = callTargetUse.ReplaceWithLclVar(m_compiler);
+    GenTreeLclVar* callTargetLclForArg = m_compiler->gtNewLclvNode(callTargetLclNum, TYP_I_IMPL);
+    DISPTREE(call)
+
+    JITDUMP("Add new arg to call arg list corresponding to PEP target");
+    NewCallArg pepTargetArg =
+        NewCallArg::Primitive(callTargetLclForArg).WellKnown(WellKnownArg::WasmPortableEntryPoint);
+    CallArg* pepArg = call->gtArgs.PushBack(m_compiler, pepTargetArg);
+
+    pepArg->SetEarlyNode(nullptr);
+    pepArg->SetLateNode(callTargetLclForArg);
+    call->gtArgs.PushLateBack(pepArg);
+
+    // Set up ABI information for this arg; PEP's should be passed as the last param to a wasm function
+    unsigned  pepIndex = call->gtArgs.CountArgs() - 1;
+    regNumber pepReg   = MakeWasmReg(pepIndex, WasmValueType::I);
+    pepArg->AbiInfo =
+        ABIPassingInformation::FromSegmentByValue(m_compiler,
+                                                  ABIPassingSegment::InRegister(pepReg, 0, TARGET_POINTER_SIZE));
+    BlockRange().InsertBefore(call, callTargetLclForArg);
+
+    // Lower the new PEP arg now that the call abi info is updated and lcl var is inserted
+    LowerArg(call, pepArg);
+    DISPTREE(call);
+
+    JITDUMP("Rewrite PEP call's control expression to indirect through the new local variable\n");
+    // Rewrite the call's control expression to have an additional load from the PEP local
+    GenTree* controlExpr = call->gtControlExpr;
+    GenTree* target      = Ind(controlExpr);
+    BlockRange().InsertAfter(controlExpr, target);
+    call->gtControlExpr = target;
+
+    JITDUMP("Finished lowering PEP call\n");
+    DISPTREERANGE(BlockRange(), call);
+}
+#endif // TARGET_WASM
+
 //------------------------------------------------------------------------
 // LowerCFGCall: Potentially lower a call to use control-flow guard. This
 // expands indirect calls into either a validate+call sequence or to a dispatch

src/coreclr/jit/lower.h

@@ -153,6 +153,7 @@ class Lowering final : public Phase
     bool     LowerCallMemcmp(GenTreeCall* call, GenTree** next);
     bool     LowerCallMemset(GenTreeCall* call, GenTree** next);
     void     LowerCFGCall(GenTreeCall* call);
+    void     LowerPEPCall(GenTreeCall* call);
     void     MovePutArgNodesUpToCall(GenTreeCall* call);
     void     MovePutArgUpToCall(GenTreeCall* call, GenTree* node);
 #ifndef TARGET_64BIT

src/coreclr/jit/morph.cpp

@@ -1756,19 +1756,23 @@ void CallArgs::AddFinalArgsAndDetermineABIInfo(Compiler* comp, GenTreeCall* call
     // That's ok; after making something a tailcall, we will invalidate this information
     // and reconstruct it if necessary. The tailcalling decision does not change since
     // this is a non-standard arg in a register.
-    bool needsIndirectionCell = call->IsR2RRelativeIndir() && !call->IsDelegateInvoke();
+#ifndef TARGET_WASM
+    bool needsIndirectionCellArg = call->IsR2RRelativeIndir() && !call->IsDelegateInvoke();
+#else
+    bool needsIndirectionCellArg = false;
+#endif
+
 #if defined(TARGET_XARCH)
-    needsIndirectionCell &= call->IsFastTailCall();
+    needsIndirectionCellArg &= call->IsFastTailCall();
 #endif
 
-    if (needsIndirectionCell)
+    if (needsIndirectionCellArg)
     {
         assert(call->gtEntryPoint.addr != nullptr);
 
         size_t   addrValue           = (size_t)call->gtEntryPoint.addr;
         GenTree* indirectCellAddress = comp->gtNewIconHandleNode(addrValue, GTF_ICON_FTN_ADDR);
         INDEBUG(indirectCellAddress->AsIntCon()->gtTargetHandle = (size_t)call->gtCallMethHnd);
-
 #ifdef TARGET_ARM
         // TODO-ARM: We currently do not properly kill this register in LSRA
         // (see getKillSetForCall which does so only for VSD calls).
@@ -1781,12 +1785,7 @@ void CallArgs::AddFinalArgsAndDetermineABIInfo(Compiler* comp, GenTreeCall* call
         // Push the stub address onto the list of arguments.
         NewCallArg indirCellAddrArg =
             NewCallArg::Primitive(indirectCellAddress).WellKnown(WellKnownArg::R2RIndirectionCell);
-#ifdef TARGET_WASM
-        // On wasm we need to ensure we put the indirection cell address last in LIR, after the SP and formal args.
-        PushBack(comp, indirCellAddrArg);
-#else
         InsertAfterThisOrFirst(comp, indirCellAddrArg);
-#endif // TARGET_WASM
     }
 #endif // defined(FEATURE_READYTORUN)