[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/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

@@ -11974,6 +11974,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 +11996,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(lvaArg0Var)->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";

src/coreclr/jit/targetwasm.cpp

@@ -48,10 +48,82 @@ ABIPassingInformation WasmClassifier::Classify(Compiler*    comp,
 {
     if (type == TYP_STRUCT)
     {
-        NYI_WASM("WasmClassifier::Classify - structs");
+        bool      passByRef = true;
+        var_types abiType   = TYP_I_IMPL;
+
+        CORINFO_CLASS_HANDLE clsHnd = structLayout->GetClassHandle();
+        assert(clsHnd != NO_CLASS_HANDLE);
+        assert(comp->info.compCompHnd->isValueClass(clsHnd));
+
+        unsigned const structSize = comp->info.compCompHnd->getClassSize(clsHnd);
+
+        switch (structSize)
+        {
+            case 1:
+            case 2:
+            case 4:
+            case 8:
+            {
+                var_types primtiveType = GetPrimitiveTypeForTrivialStruct(comp, clsHnd);
+                if (primtiveType != TYP_UNDEF)
+                {
+                    abiType   = genActualType(primtiveType);
+                    passByRef = false;
+                }
+                break;
+            }
+            default:
+                break;
+        }
+
+        regNumber         reg = MakeWasmReg(m_localIndex++, abiType);
+        ABIPassingSegment seg = ABIPassingSegment::InRegister(reg, 0, genTypeSize(abiType));
+        return ABIPassingInformation::FromSegment(comp, passByRef, seg);
     }
 
     regNumber         reg = MakeWasmReg(m_localIndex++, genActualType(type));
     ABIPassingSegment seg = ABIPassingSegment::InRegister(reg, 0, genTypeSize(type));
     return ABIPassingInformation::FromSegmentByValue(comp, seg);
 }
+
+//-----------------------------------------------------------------------------
+// GetPrimitiveTypeForTrivialStruct:
+//   Get the primitive type for a trivial struct for the Wasm ABI.
+//
+// Parameters:
+//   comp   - Compiler instance
+//   clsHnd - Class handle of the struct
+//
+// Returns:
+//   The primitive type for the struct, or TYP_UNDEF if it cannot be represented as a primitive.
+//
+// TODO-Wasm: Union types? 128-bit types? SIMD?
+//
+var_types WasmClassifier::GetPrimitiveTypeForTrivialStruct(Compiler* comp, CORINFO_CLASS_HANDLE clsHnd)
+{
+    for (;;)
+    {
+        // all of class chain must be of value type and must have only one field
+        if (!comp->info.compCompHnd->isValueClass(clsHnd) ||
+            comp->info.compCompHnd->getClassNumInstanceFields(clsHnd) != 1)
+        {
+            return TYP_UNDEF;
+        }
+
+        CORINFO_CLASS_HANDLE* pClsHnd   = &clsHnd;
+        CORINFO_FIELD_HANDLE  fldHnd    = comp->info.compCompHnd->getFieldInClass(clsHnd, 0);
+        CorInfoType           fieldType = comp->info.compCompHnd->getFieldType(fldHnd, pClsHnd);
+
+        var_types vt = JITtype2varType(fieldType);
+
+        if (fieldType == CORINFO_TYPE_VALUECLASS)
+        {
+            clsHnd = *pClsHnd;
+        }
+        else
+        {
+            assert(vt != TYP_STRUCT);
+            return vt;
+        }
+    }
+}