[Wasm RyuJit] Wasm stack arg preliminaries

src/coreclr/jit/abi.h

@@ -338,6 +338,8 @@ class WasmClassifier
 public:
     WasmClassifier(const ClassifierInfo& info);
 
+    static var_types GetPrimitiveTypeForTrivialStruct(Compiler* comp, CORINFO_CLASS_HANDLE clsHnd);
+
     unsigned StackSize()
     {
         return 0;

src/coreclr/jit/codegenwasm.cpp

@@ -32,17 +32,7 @@ void CodeGen::genMarkLabelsForCodegen()
 void CodeGen::genBeginFnProlog()
 {
     // TODO-WASM: proper local count, local declarations, and shadow stack maintenance
-    GetEmitter()->emitIns_I(INS_local_cnt, EA_8BYTE, (unsigned)WasmValueType::Count - 1);
-    // Emit 1 local of each supported value type to ensure
-    // the declarations can be encoded.
-    // TODO-WASM: remove and declare locals based on RA assignments once this is supported.
-    int localOffset  = 0;
-    int countPerType = 1;
-    for (unsigned i = (uint8_t)WasmValueType::Invalid + 1; i < (unsigned)WasmValueType::Count; i++)
-    {
-        GetEmitter()->emitIns_I_Ty(INS_local_decl, countPerType, static_cast<WasmValueType>(i), localOffset);
-        localOffset += countPerType;
-    }
+    GetEmitter()->emitIns_I(INS_local_cnt, EA_8BYTE, 0);
 }
 
 //------------------------------------------------------------------------
@@ -71,7 +61,7 @@ void CodeGen::genAllocLclFrame(unsigned frameSize, regNumber initReg, bool* pIni
         return;
     }
 
-    unsigned initialSPLclIndex = 0; // TODO-WASM: remove this hardcoding once we have the SP arg local.
+    unsigned initialSPLclIndex = compiler->lvaWasmSpArg;
     unsigned spLclIndex        = WasmRegToIndex(spReg);
     assert(initialSPLclIndex == spLclIndex);
     if (frameSize != 0)
@@ -150,6 +140,13 @@ void CodeGen::genHomeRegisterParams(regNumber initReg, bool* initRegStillZeroed)
 
     for (unsigned lclNum = 0; lclNum < compiler->info.compArgsCount; lclNum++)
     {
+        // We don't need to home the stack pointer
+        //
+        if (lclNum == compiler->lvaWasmSpArg)
+        {
+            continue;
+        }
+
         LclVarDsc*                   lclDsc  = compiler->lvaGetDesc(lclNum);
         const ABIPassingInformation& abiInfo = compiler->lvaGetParameterABIInfo(lclNum);
 
@@ -501,6 +498,10 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genCodeForStoreInd(treeNode->AsStoreInd());
             break;
 
+        case GT_PUTARG_REG:
+            genPutArgReg(treeNode->AsOp());
+            break;
+
         default:
 #ifdef DEBUG
             NYIRAW(GenTree::OpName(treeNode->OperGet()));
@@ -1081,6 +1082,20 @@ void CodeGen::genCodeForStoreInd(GenTreeStoreInd* tree)
     genUpdateLife(tree);
 }
 
+//---------------------------------------------------------------------
+// genPutArgReg - generate code for a GT_PUTARG_REG node
+//
+// Arguments
+//    tree - the GT_PUTARG_REG node
+//
+void CodeGen::genPutArgReg(GenTreeOp* tree)
+{
+    assert(tree->OperIs(GT_PUTARG_REG));
+    GenTree* const op1 = tree->gtOp1;
+    genConsumeReg(op1);
+    genProduceReg(tree);
+}
+
 //------------------------------------------------------------------------
 // genCodeForCompare: Produce code for a GT_EQ/GT_NE/GT_LT/GT_LE/GT_GE/GT_GT node.
 //

src/coreclr/jit/compiler.h

@@ -3960,6 +3960,10 @@ class Compiler
                                       // arguments
 #endif                                // TARGET_X86
 
+#if defined(TARGET_WASM)
+    unsigned lvaWasmSpArg = BAD_VAR_NUM; // lcl var index of Wasm stack pointer arg
+#endif // defined(TARGET_WASM)
+
     unsigned lvaInlinedPInvokeFrameVar = BAD_VAR_NUM; // variable representing the InlinedCallFrame
     unsigned lvaReversePInvokeFrameVar = BAD_VAR_NUM; // variable representing the reverse PInvoke frame
     unsigned lvaMonAcquired = BAD_VAR_NUM; // boolean variable introduced into in synchronized methods
@@ -4118,6 +4122,11 @@ class Compiler
     void lvaInitVarArgsHandle(unsigned* curVarNum);
     void lvaInitAsyncContinuation(unsigned* curVarNum);
 
+#if defined(TARGET_WASM)
+    void lvaInitWasmStackPtr(unsigned* curVarNum);
+    void lvaInitWasmPortableEntryPtr(unsigned* curVarNum);
+#endif // defined(TARGET_WASM)
+
     void lvaInitVarDsc(LclVarDsc*              varDsc,
                        unsigned                varNum,
                        CorInfoType             corInfoType,
@@ -10588,9 +10597,10 @@ class Compiler
         unsigned  compILargsCount;   // Number of arguments (incl. implicit but not hidden)
         unsigned  compArgsCount;     // Number of arguments (incl. implicit and     hidden)
 
-        unsigned compRetBuffArg;    // position of hidden return param var (0, 1) (BAD_VAR_NUM means not present);
-        unsigned compTypeCtxtArg;   // position of hidden param for type context for generic code
-                                    // (CORINFO_CALLCONV_PARAMTYPE)
+        unsigned compRetBuffArg;  // position of hidden return param var (0, 1) (BAD_VAR_NUM means not present);
+        unsigned compTypeCtxtArg; // position of hidden param for type context for generic code
+                                  // (CORINFO_CALLCONV_PARAMTYPE)
+
         unsigned       compThisArg; // position of implicit this pointer param (not to be confused with lvaArg0Var)
         unsigned       compILlocalsCount; // Number of vars : args + locals (incl. implicit but not hidden)
         unsigned       compLocalsCount;   // Number of vars : args + locals (incl. implicit and     hidden)

src/coreclr/jit/compiler.hpp

@@ -2996,6 +2996,14 @@ inline unsigned Compiler::compMapILargNum(unsigned ILargNum)
 {
     assert(ILargNum < info.compILargsCount);
 
+#if defined(TARGET_WASM)
+    if (ILargNum >= lvaWasmSpArg)
+    {
+        ILargNum++;
+        assert(ILargNum < info.compLocalsCount); // compLocals count already adjusted.
+    }
+#endif
+
     // Note that this works because if compRetBuffArg/compTypeCtxtArg/lvVarargsHandleArg are not present
     // they will be BAD_VAR_NUM (MAX_UINT), which is larger than any variable number.
     if (ILargNum >= info.compRetBuffArg)

src/coreclr/jit/gentree.cpp

@@ -2344,6 +2344,13 @@ bool GenTreeCall::HasNonStandardAddedArgs(Compiler* compiler) const
 //
 int GenTreeCall::GetNonStandardAddedArgCount(Compiler* compiler) const
 {
+#if defined(TARGET_WASM)
+    // TODO-WASM: may need adjustments for other hidden args
+    // For now: managed calls get extra SP + PortableEntryPoint args, but
+    // we're not adding the PE arg yet. So just note one extra arg.
+    return IsUnmanaged() ? 0 : 1;
+#endif // defined(TARGET_WASM)
+
     if (IsUnmanaged() && !compiler->opts.ShouldUsePInvokeHelpers())
     {
         // R11 = PInvoke cookie param
@@ -11974,6 +11981,12 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
             {
                 ilName = "AsyncCont";
             }
+#if defined(TARGET_WASM)
+            else if (lclNum == lvaWasmSpArg)
+            {
+                ilName = "WasmSP";
+            }
+#endif // defined(TARGET_WASM)
             else
             {
                 ilKind = "tmp";
@@ -11990,7 +12003,7 @@ void Compiler::gtGetLclVarNameInfo(unsigned lclNum, const char** ilKindOut, cons
     }
     else if (lclNum < (compIsForInlining() ? impInlineInfo->InlinerCompiler->info.compArgsCount : info.compArgsCount))
     {
-        if (ilNum == 0 && !info.compIsStatic)
+        if ((ilNum == 0) && !info.compIsStatic)
         {
             ilName = "this";
         }

src/coreclr/jit/gentree.h

@@ -4633,6 +4633,8 @@ enum class WellKnownArg : unsigned
     RuntimeMethodHandle,
     AsyncExecutionContext,
     AsyncSynchronizationContext,
+    WasmShadowStackPointer,
+    WasmPortableEntryPoint
 };
 
 #ifdef DEBUG

src/coreclr/jit/lclvars.cpp

@@ -166,6 +166,11 @@ void Compiler::lvaInitTypeRef()
         info.compArgsCount++;
     }
 
+#if defined(TARGET_WASM)
+    // Wasm passes stack pointer as first arg, portable entry point as last arg
+    info.compArgsCount += 2;
+#endif
+
     lvaCount = info.compLocalsCount = info.compArgsCount + info.compMethodInfo->locals.numArgs;
 
     info.compILlocalsCount = info.compILargsCount + info.compMethodInfo->locals.numArgs;
@@ -338,6 +343,12 @@ void Compiler::lvaInitArgs(bool hasRetBuffArg)
     //----------------------------------------------------------------------
 
     unsigned varNum = 0;
+
+#if defined(TARGET_WASM)
+    // Wasm stack pointer is first arg
+    lvaInitWasmStackPtr(&varNum);
+#endif
+
     // Is there a "this" pointer ?
     lvaInitThisPtr(&varNum);
 
@@ -397,6 +408,11 @@ void Compiler::lvaInitArgs(bool hasRetBuffArg)
     lvaInitVarArgsHandle(&varNum);
 #endif
 
+#if defined(TARGET_WASM)
+    // Wasm portable entry point is the very last arg
+    lvaInitWasmPortableEntryPtr(&varNum);
+#endif
+
     //----------------------------------------------------------------------
 
     // We have set info.compArgsCount in compCompile()
@@ -431,7 +447,12 @@ void Compiler::lvaInitThisPtr(unsigned* curVarNum)
     varDsc->lvIsPtr   = 1;
 
     lvaArg0Var = info.compThisArg = *curVarNum;
+
+#if defined(TARGET_WASM)
+    noway_assert(info.compThisArg == 1);
+#else
     noway_assert(info.compThisArg == 0);
+#endif
 
     if (eeIsValueClass(info.compClassHnd))
     {
@@ -461,6 +482,29 @@ void Compiler::lvaInitRetBuffArg(unsigned* curVarNum, bool useFixedRetBufReg)
     (*curVarNum)++;
 }
 
+#if defined(TARGET_WASM)
+
+/*****************************************************************************/
+void Compiler::lvaInitWasmStackPtr(unsigned* curVarNum)
+{
+    LclVarDsc* varDsc = lvaGetDesc(*curVarNum);
+    varDsc->lvType    = TYP_I_IMPL;
+    varDsc->lvIsParam = 1;
+    lvaWasmSpArg      = *curVarNum;
+    (*curVarNum)++;
+}
+
+void Compiler::lvaInitWasmPortableEntryPtr(unsigned* curVarNum)
+{
+    // This ends up never being used in the JIT, so perhaps it's not needed.
+    LclVarDsc* varDsc = lvaGetDesc(*curVarNum);
+    varDsc->lvType    = TYP_I_IMPL;
+    varDsc->lvIsParam = 1;
+    (*curVarNum)++;
+}
+
+#endif // defined(TARGET_WASM)
+
 //-----------------------------------------------------------------------------
 // lvaInitUserArgs:
 //     Initialize local var descriptions for incoming user arguments
@@ -837,6 +881,13 @@ void Compiler::lvaClassifyParameterABI(Classifier& classifier)
         }
 #endif
 
+#if defined(TARGET_WASM)
+        if (i == lvaWasmSpArg)
+        {
+            wellKnownArg = WellKnownArg::WasmShadowStackPointer;
+        }
+#endif
+
         ABIPassingInformation abiInfo = classifier.Classify(this, dsc->TypeGet(), structLayout, wellKnownArg);
         lvaParameterPassingInfo[i]    = abiInfo;
 
@@ -1126,6 +1177,13 @@ unsigned Compiler::compMap2ILvarNum(unsigned varNum) const
         return (unsigned)ICorDebugInfo::UNKNOWN_ILNUM;
     }
 
+#if defined(TARGET_WASM)
+    if (varNum == lvaWasmSpArg)
+    {
+        return (unsigned)ICorDebugInfo::UNKNOWN_ILNUM;
+    }
+#endif // defined(TARGET_WASM)
+
     unsigned originalVarNum = varNum;
 
     // Now mutate varNum to remove extra parameters from the count.
@@ -1153,6 +1211,13 @@ unsigned Compiler::compMap2ILvarNum(unsigned varNum) const
         varNum--;
     }
 
+#if defined(TARGET_WASM)
+    if (lvaWasmSpArg != BAD_VAR_NUM && originalVarNum > lvaWasmSpArg)
+    {
+        varNum--;
+    }
+#endif
+
     if (varNum >= info.compLocalsCount)
     {
         return (unsigned)ICorDebugInfo::UNKNOWN_ILNUM; // Cannot be mapped
@@ -3716,7 +3781,7 @@ PhaseStatus Compiler::lvaMarkLocalVars()
     // Update bookkeeping on the generic context.
     if (lvaKeepAliveAndReportThis())
     {
-        lvaGetDesc(0u)->lvImplicitlyReferenced = reportParamTypeArg;
+        lvaGetDesc(info.compThisArg)->lvImplicitlyReferenced = reportParamTypeArg;
     }
     else if (lvaReportParamTypeArg())
     {

src/coreclr/jit/morph.cpp

@@ -682,6 +682,10 @@ const char* getWellKnownArgName(WellKnownArg arg)
             return "AsyncExecutionContext";
         case WellKnownArg::AsyncSynchronizationContext:
             return "AsyncSynchronizationContext";
+        case WellKnownArg::WasmShadowStackPointer:
+            return "WasmShadowStackPointer";
+        case WellKnownArg::WasmPortableEntryPoint:
+            return "WasmPortableEntryPoint";
     }
 
     return "N/A";
@@ -1701,6 +1705,15 @@ void CallArgs::AddFinalArgsAndDetermineABIInfo(Compiler* comp, GenTreeCall* call
     // in the implementation of fast tail call.
     // *********** END NOTE *********
 
+#if defined(TARGET_WASM)
+    // On WASM, we need to add an initial hidden argument for the stack pointer for managed calls.
+    if (!call->IsUnmanaged())
+    {
+        GenTree* const stackPointer = comp->gtNewLclVarNode(comp->lvaWasmSpArg, TYP_I_IMPL);
+        PushFront(comp, NewCallArg::Primitive(stackPointer).WellKnown(WellKnownArg::WasmShadowStackPointer));
+    }
+#endif // defined(TARGET_WASM)
+
 #if defined(TARGET_ARM)
     // A non-standard calling convention using wrapper delegate invoke is used on ARM, only, for wrapper
     // delegates. It is used for VSD delegate calls where the VSD custom calling convention ABI requires passing
@@ -1826,6 +1839,10 @@ void CallArgs::AddFinalArgsAndDetermineABIInfo(Compiler* comp, GenTreeCall* call
     }
 #endif
 
+#if defined(TARGET_WASM)
+    // TODO-WASM: pass the portable entry point as the last argument for managed calls
+#endif
+
     ClassifierInfo info;
     info.CallConv = call->GetUnmanagedCallConv();
     // X86 tailcall helper is considered varargs, but not for ABI classification purposes.

src/coreclr/jit/regallocwasm.cpp

@@ -87,6 +87,15 @@ void WasmRegAlloc::IdentifyCandidates()
     for (unsigned lclNum = 0; lclNum < m_compiler->lvaCount; lclNum++)
     {
         LclVarDsc* varDsc = m_compiler->lvaGetDesc(lclNum);
+
+        if (lclNum == m_compiler->lvaWasmSpArg)
+        {
+            varDsc->lvLRACandidate = true;
+            m_spReg                = MakeWasmReg(m_compiler->lvaWasmSpArg, TYP_I_IMPL);
+            varDsc->SetRegNum(m_spReg);
+            continue;
+        }
+
         varDsc->SetRegNum(REG_STK);
         varDsc->lvLRACandidate = isRegCandidate(varDsc);