[Wasm RyuJIT] Add Instruction Encoding for Local Var Declarations and Emit Local Declarations

src/coreclr/jit/codegen.h

@@ -13,6 +13,7 @@
 #include "compiler.h" // temporary??
 #include "regset.h"
 #include "jitgcinfo.h"
+#include "wasmtypesdef.h"
 
 class CodeGen final : public CodeGenInterface
 {
@@ -589,6 +590,9 @@ class CodeGen final : public CodeGenInterface
     void genReserveProlog(BasicBlock* block); // currently unused
     void genReserveEpilog(BasicBlock* block);
     void genFnProlog();
+#if defined(TARGET_WASM)
+    void genWasmLocals();
+#endif
     void genFnEpilog(BasicBlock* block);
 
     void genReserveFuncletProlog(BasicBlock* block);

src/coreclr/jit/codegencommon.cpp

@@ -5650,13 +5650,35 @@ void CodeGen::genFnProlog()
 #endif // defined(DEBUG) && defined(TARGET_XARCH)
 
 #else  // defined(TARGET_WASM)
-    // TODO-WASM: prolog zeroing, shadow stack maintenance
+    // TODO-WASM: proper local count, local declarations, and shadow stack maintenance
+    assert(compiler->info.compLocalsCount >= compiler->info.compArgsCount);
+    GetEmitter()->emitIns_I(INS_local_cnt, EA_8BYTE, compiler->info.compLocalsCount - compiler->info.compArgsCount);
+    genWasmLocals();
     GetEmitter()->emitMarkPrologEnd();
 #endif // !defined(TARGET_WASM)
 
     GetEmitter()->emitEndProlog();
 }
 
+#if defined(TARGET_WASM)
+
+//------------------------------------------------------------------------
+// genWasmLocals: generate wasm locals for all locals
+//
+// TODO-WASM: re-evaluate this, we may not want a 1:1 mapping of locals to wasm locals
+// TODO-WASM: pre-declare all "register" locals
+void CodeGen::genWasmLocals()
+{
+    for (unsigned i = 0; i < compiler->info.compLocalsCount - compiler->info.compArgsCount; i++)
+    {
+        LclVarDsc* varDsc = compiler->lvaGetDesc(i);
+        assert(!varDsc->lvIsParam);
+        emitter::instWasmValueType type = GetEmitter()->genWasmTypeFromVarType(varDsc->TypeGet());
+        GetEmitter()->emitIns_I_Ty(INS_local_decl, 1, type);
+    }
+}
+#endif
+
 #if !defined(TARGET_WASM)
 
 //----------------------------------------------------------------------------------

src/coreclr/jit/codegenwasm.cpp

@@ -42,6 +42,7 @@ void CodeGen::genFnEpilog(BasicBlock* block)
     // TODO-WASM-CQ: do not emit "return" in case this is the last block
 
     instGen(INS_return);
+    instGen(INS_end);
 }
 
 void CodeGen::genCaptureFuncletPrologEpilogInfo()

src/coreclr/jit/emit.h

@@ -640,6 +640,10 @@ class emitter
         BasicBlock*       idTargetBlock; // Target block for branches
     };
 
+#if defined(TARGET_WASM)
+#include "wasmtypesdef.h"
+#endif
+
 #ifdef TARGET_ARM
     unsigned insEncodeSetFlags(insFlags sf);
 
@@ -892,6 +896,9 @@ class emitter
         unsigned _idLclFPBase : 1; // access a local on stack - SP based offset
         insOpts  _idInsOpt    : 3; // options for Load/Store instructions
 #endif
+#ifdef TARGET_WASM
+        unsigned _idLclDecl : 1; // is this a local declaration?
+#endif
 
         ////////////////////////////////////////////////////////////////////////
         // Space taken up to here:
@@ -919,7 +926,7 @@ class emitter
 #elif defined(TARGET_AMD64)
 #define ID_EXTRA_BITFIELD_BITS (20)
 #elif defined(TARGET_WASM)
-#define ID_EXTRA_BITFIELD_BITS (-4)
+#define ID_EXTRA_BITFIELD_BITS (-3)
 #else
 #error Unsupported or unset target architecture
 #endif
@@ -1296,6 +1303,18 @@ class emitter
             assert(reg == _idReg1);
         }
 
+#ifdef TARGET_WASM
+        bool idIsLclVarDecl() const
+        {
+            return (_idLclDecl != 0);
+        }
+
+        void idSetIsLclVarDecl(bool isDecl)
+        {
+            _idLclDecl = isDecl ? 1 : 0;
+        }
+#endif
+
 #ifdef TARGET_ARM64
         GCtype idGCrefReg2() const
         {
@@ -2335,6 +2354,25 @@ class emitter
     };
 #endif
 
+#if defined(TARGET_WASM)
+    struct instrDescLclVarDecl : instrDesc
+    {
+        instrDescLclVarDecl() = delete;
+        cnsval_ssize_t    lclCnt;
+        instWasmValueType lclType;
+
+        void idLclType(instWasmValueType type)
+        {
+            lclType = type;
+        }
+
+        void idLclCnt(cnsval_ssize_t cnt)
+        {
+            lclCnt = cnt;
+        }
+    };
+#endif // TARGET_WASM
+
 #ifdef TARGET_RISCV64
     struct instrDescLoadImm : instrDescCns
     {
@@ -3293,6 +3331,11 @@ class emitter
     instrDesc* emitNewInstrCns(emitAttr attr, cnsval_ssize_t cns);
     instrDesc* emitNewInstrDsp(emitAttr attr, target_ssize_t dsp);
     instrDesc* emitNewInstrCnsDsp(emitAttr attr, target_ssize_t cns, int dsp);
+
+#ifdef TARGET_WASM
+
+#endif
+
 #ifdef TARGET_ARM
     instrDesc* emitNewInstrReloc(emitAttr attr, BYTE* addr);
 #endif // TARGET_ARM

src/coreclr/jit/emitfmtswasm.h

@@ -26,15 +26,17 @@ enum ID_OPS
 //                  (unused)
 //////////////////////////////////////////////////////////////////////////////
 
-IF_DEF(NONE,    IS_NONE, NONE)
-IF_DEF(OPCODE,  IS_NONE, NONE) // <opcode>
-IF_DEF(BLOCK,   IS_NONE, NONE) // <opcode> <0x40>
-IF_DEF(LABEL,   IS_NONE, NONE) // <ULEB128 immediate>
-IF_DEF(ULEB128, IS_NONE, NONE) // <opcode> <ULEB128 immediate>
-IF_DEF(SLEB128, IS_NONE, NONE) // <opcode> <LEB128 immediate (signed)>
-IF_DEF(F32,     IS_NONE, NONE) // <opcode> <f32 immediate (stored as 64-bit integer constant)>
-IF_DEF(F64,     IS_NONE, NONE) // <opcode> <f64 immediate (stored as 64-bit integer constant)>
-IF_DEF(MEMARG,  IS_NONE, NONE) // <opcode> <memarg> (<align> <offset>)
+IF_DEF(NONE,       IS_NONE, NONE)
+IF_DEF(OPCODE,     IS_NONE, NONE) // <opcode>
+IF_DEF(BLOCK,      IS_NONE, NONE) // <opcode> <0x40>
+IF_DEF(LABEL,      IS_NONE, NONE) // <ULEB128 immediate>
+IF_DEF(ULEB128,    IS_NONE, NONE) // <opcode> <ULEB128 immediate>
+IF_DEF(SLEB128,    IS_NONE, NONE) // <opcode> <LEB128 immediate (signed)>
+IF_DEF(F32,        IS_NONE, NONE) // <opcode> <f32 immediate (stored as 64-bit integer constant)>
+IF_DEF(F64,        IS_NONE, NONE) // <opcode> <f64 immediate (stored as 64-bit integer constant)>
+IF_DEF(MEMARG,     IS_NONE, NONE) // <opcode> <memarg> (<align> <offset>)
+IF_DEF(LOCAL_CNT,  IS_NONE, NONE) // <ULEB128 immediate>
+IF_DEF(LOCAL_DECL, IS_NONE, NONE) // <ULEB128 immediate> <byte>
 
 #undef IF_DEF
 #endif // !DEFINE_ID_OPS

src/coreclr/jit/emitwasm.cpp

@@ -116,6 +116,45 @@ bool emitter::emitInsIsStore(instruction ins)
     return false;
 }
 
+emitter::instrDesc* emitter::emitNewInstrLclVarDecl(emitAttr attr, cnsval_ssize_t localCount, instWasmValueType type)
+{
+    instrDescLclVarDecl* id = static_cast<instrDescLclVarDecl*>(emitAllocAnyInstr(sizeof(instrDescLclVarDecl), attr));
+    id->idSetIsLclVarDecl(true);
+    id->idLclCnt(localCount);
+    id->idLclType(type);
+
+    return id;
+}
+
+//-----------------------------------------------------------------------------------
+// emitIns_I_Ty: Emit an instruction for a local variable declaration, encoding both
+// a count (immediate) and a value type. This is specifically used for local variable
+// declarations that require both the number of locals and their type to be encoded.
+//
+void emitter::emitIns_I_Ty(instruction ins, cnsval_ssize_t imm, emitter::instWasmValueType valType)
+{
+    instrDesc* id  = this->emitNewInstrLclVarDecl(EA_8BYTE, imm, valType);
+    insFormat  fmt = this->emitInsFormat(ins);
+
+    id->idIns(ins);
+    id->idInsFmt(fmt);
+
+    this->dispIns(id);
+    this->appendToCurIG(id);
+}
+
+emitter::instWasmValueType emitter::emitGetLclVarDeclType(instrDesc* id)
+{
+    assert(id->idIsLclVarDecl());
+    return static_cast<instrDescLclVarDecl*>(id)->lclType;
+}
+
+cnsval_ssize_t emitter::emitGetLclVarDeclCount(instrDesc* id)
+{
+    assert(id->idIsLclVarDecl());
+    return static_cast<instrDescLclVarDecl*>(id)->lclCnt;
+}
+
 emitter::insFormat emitter::emitInsFormat(instruction ins)
 {
     static_assert(IF_COUNT < 255);
@@ -155,6 +194,11 @@ size_t emitter::emitSizeOfInsDsc(instrDesc* id) const
         return sizeof(instrDescCns);
     }
 
+    if (id->idIsLclVarDecl())
+    {
+        return sizeof(instrDescLclVarDecl);
+    }
+
     return sizeof(instrDesc);
 }
 
@@ -200,9 +244,17 @@ unsigned emitter::instrDesc::idCodeSize() const
             size += 1;
             break;
         case IF_LABEL:
+        case IF_LOCAL_CNT:
             assert(!idIsCnsReloc());
             size = SizeOfULEB128(emitGetInsSC(this));
             break;
+        case IF_LOCAL_DECL:
+        {
+            assert(idIsLclVarDecl());
+            instrDescLclVarDecl* idl = static_cast<instrDescLclVarDecl*>(const_cast<instrDesc*>(this));
+            size                     = SizeOfULEB128(idl->lclCnt) + sizeof(emitter::instWasmValueType);
+            break;
+        }
         case IF_ULEB128:
             size += idIsCnsReloc() ? PADDED_RELOC_SIZE : SizeOfULEB128(emitGetInsSC(this));
             break;
@@ -354,6 +406,24 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
             dst += emitOutputULEB128(dst, offset);
             break;
         }
+        case IF_LOCAL_CNT:
+        {
+            cnsval_ssize_t constant = emitGetInsSC(id);
+            dst += emitOutputULEB128(dst, (uint64_t)constant);
+            break;
+        }
+        case IF_LOCAL_DECL:
+        {
+            assert(id->idIsLclVarDecl());
+            cnsval_ssize_t    count   = emitGetLclVarDeclCount(id);
+            instWasmValueType valType = emitGetLclVarDeclType(id);
+            dst += emitOutputULEB128(dst, (uint64_t)count);
+            // TODO-WASM: currently assuming all locals are numtypes which are single byte encoded.
+            // vec types are also single byte encoded. If we end up using reftypes, we'll need to handle the more
+            // complex encoding.
+            dst += emitOutputByte(dst, static_cast<uint8_t>(valType));
+            break;
+        }
         default:
             NYI_WASM("emitOutputInstr");
             break;
@@ -472,13 +542,22 @@ void emitter::emitDispIns(
 
         case IF_LABEL:
         case IF_ULEB128:
+        case IF_LOCAL_CNT:
         {
             cnsval_ssize_t imm = emitGetInsSC(id);
             printf(" %llu", (uint64_t)imm);
             dispJumpTargetIfAny();
         }
         break;
 
+        case IF_LOCAL_DECL:
+        {
+            cnsval_ssize_t    imm     = emitGetLclVarDeclCount(id);
+            instWasmValueType valType = emitGetLclVarDeclType(id);
+            printf(" %llu %s", (uint64_t)imm, instWasmValueTypeToStr(valType));
+        }
+        break;
+
         case IF_SLEB128:
         {
             cnsval_ssize_t imm = emitGetInsSC(id);

src/coreclr/jit/emitwasm.h

@@ -5,6 +5,8 @@
 /*             Debug-only routines to display instructions              */
 /************************************************************************/
 
+#include "wasmtypesdef.h"
+
 #if defined(DEBUG) || defined(LATE_DISASM)
 void getInsSveExecutionCharacteristics(instrDesc* id, insExecutionCharacteristics& result);
 #endif // defined(DEBUG) || defined(LATE_DISASM)
@@ -18,6 +20,7 @@ void emitDispInst(instruction ins);
 public:
 void emitIns(instruction ins);
 void emitIns_I(instruction ins, emitAttr attr, cnsval_ssize_t imm);
+void emitIns_I_Ty(instruction ins, cnsval_ssize_t imm, instWasmValueType valType);
 void emitIns_J(instruction ins, emitAttr attr, cnsval_ssize_t imm, BasicBlock* tgtBlock);
 void emitIns_S(instruction ins, emitAttr attr, int varx, int offs);
 void emitIns_R(instruction ins, emitAttr attr, regNumber reg);
@@ -33,6 +36,10 @@ static unsigned SizeOfSLEB128(int64_t value);
 
 static unsigned emitGetAlignHintLog2(const instrDesc* id);
 
+instrDesc*               emitNewInstrLclVarDecl(emitAttr attr, cnsval_ssize_t localCount, instWasmValueType type);
+static instWasmValueType emitGetLclVarDeclType(instrDesc* id);
+static cnsval_ssize_t    emitGetLclVarDeclCount(instrDesc* id);
+
 /************************************************************************/
 /*  Private members that deal with target-dependent instr. descriptors  */
 /************************************************************************/

src/coreclr/jit/instrswasm.h

@@ -37,6 +37,8 @@ INST(br,          "br",          0, IF_ULEB128, 0x0C)
 INST(br_if,       "br_if",       0, IF_ULEB128, 0x0D)
 INST(br_table,    "br_table",    0, IF_ULEB128, 0x0E)
 INST(return,      "return",      0, IF_OPCODE,  0x0F)
+INST(local_cnt,   "local.cnt",   0, IF_LOCAL_CNT, 0x00)
+INST(local_decl,  "local",       0, IF_LOCAL_DECL, 0x00)
 
 INST(local_get,   "local.get",   0, IF_ULEB128, 0x20)
 INST(i32_load,    "i32.load",    0, IF_MEMARG,  0x28)

src/coreclr/jit/registeropswasm.h

@@ -9,7 +9,6 @@ enum class WasmValueType : unsigned
     F32,
     F64,
     Count,
-
 #ifdef TARGET_64BIT
     I = I64,
 #else