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

src/coreclr/jit/codegen.h

@@ -14,6 +14,10 @@
 #include "regset.h"
 #include "jitgcinfo.h"
 
+#if defined(TARGET_WASM)
+#include "wasmtypesdef.h"
+#endif
+
 class CodeGen final : public CodeGenInterface
 {
     friend class emitter;
@@ -589,6 +593,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,34 @@ 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, (unsigned)WasmValueType::Count - 1);
+    genWasmLocals();
     GetEmitter()->emitMarkPrologEnd();
 #endif // !defined(TARGET_WASM)
 
     GetEmitter()->emitEndProlog();
 }
 
+#if defined(TARGET_WASM)
+
+//------------------------------------------------------------------------
+// genWasmLocals: generate wasm locals for all locals
+//
+// TODO-WASM: pre-declare all "register" locals
+void CodeGen::genWasmLocals()
+{
+    // 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.
+    for (unsigned i = (uint8_t)WasmValueType::Invalid + 1; i < (unsigned)WasmValueType::Count; i++)
+    {
+        GetEmitter()->emitIns_I_Ty(INS_local_decl, 1, static_cast<WasmValueType>(i));
+    }
+}
+#endif
+
 #if !defined(TARGET_WASM)
 
 //----------------------------------------------------------------------------------

src/coreclr/jit/codegenlinear.cpp

@@ -574,7 +574,7 @@ void CodeGen::genCodeForBBlist()
         // For example, if the following IR ends the current block, no code will have been generated for
         // offset 21:
         //
-        //          (  0,  0) [000040] ------------                il_offset void   IL offset: 21
+        //          (  0,  0) [000040] ------------   /             il_offset void   IL offset: 21
         //
         //     N001 (  0,  0) [000039] ------------                nop       void
         //

src/coreclr/jit/codegenwasm.cpp

@@ -40,8 +40,13 @@ void CodeGen::genFnEpilog(BasicBlock* block)
 
     // TODO-WASM: shadow stack maintenance
     // TODO-WASM-CQ: do not emit "return" in case this is the last block
-
     instGen(INS_return);
+    // TODO-WASM: this condition will not be sufficient if we have funclet to determine if we're at the end of
+    // codegen for a method. Revisit later.
+    if (block->IsLast()) 
+    {
+        instGen(INS_end);
+    }
 }
 
 void CodeGen::genCaptureFuncletPrologEpilogInfo()

src/coreclr/jit/emit.h

@@ -1296,6 +1296,13 @@ class emitter
             assert(reg == _idReg1);
         }
 
+#ifdef TARGET_WASM
+        bool idIsLclVarDecl() const
+        {
+            return _idInsFmt == IF_LOCAL_DECL;
+        }
+#endif
+
 #ifdef TARGET_ARM64
         GCtype idGCrefReg2() const
         {
@@ -2335,6 +2342,25 @@ class emitter
     };
 #endif
 
+#if defined(TARGET_WASM)
+    struct instrDescLclVarDecl : instrDesc
+    {
+        instrDescLclVarDecl() = delete;
+        cnsval_ssize_t    lclCnt;
+        WasmValueType lclType;
+
+        void idLclType(WasmValueType type)
+        {
+            lclType = type;
+        }
+
+        void idLclCnt(cnsval_ssize_t cnt)
+        {
+            lclCnt = cnt;
+        }
+    };
+#endif // TARGET_WASM
+
 #ifdef TARGET_RISCV64
     struct instrDescLoadImm : instrDescCns
     {
@@ -3293,6 +3319,7 @@ 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_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,44 @@ bool emitter::emitInsIsStore(instruction ins)
     return false;
 }
 
+emitter::instrDesc* emitter::emitNewInstrLclVarDecl(emitAttr attr, cnsval_ssize_t localCount, WasmValueType type)
+{
+    instrDescLclVarDecl* id = static_cast<instrDescLclVarDecl*>(emitAllocAnyInstr(sizeof(instrDescLclVarDecl), attr));
+    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, WasmValueType 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);
+}
+
+WasmValueType 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 +193,11 @@ size_t emitter::emitSizeOfInsDsc(instrDesc* id) const
         return sizeof(instrDescCns);
     }
 
+    if (id->idIsLclVarDecl())
+    {
+        return sizeof(instrDescLclVarDecl);
+    }
+
     return sizeof(instrDesc);
 }
 
@@ -181,6 +224,23 @@ unsigned emitter::SizeOfSLEB128(int64_t value)
     return (x * 37) >> 8;
 }
 
+static uint8_t getWasmValueTypeCode(WasmValueType type)
+{
+    // clang-format off
+    static const uint8_t typecode_mapping[] = {
+        0x00, // WasmValueType::Invalid = 0,
+        0x7C, // WasmValueType::F64 = 1,
+        0x7D, // WasmValueType::F32 = 2,
+        0x7E, // WasmValueType::I64 = 3,
+        0x7F, // WasmValueType::I32 = 4,
+    };
+    static const int WASM_TYP_COUNT = ArrLen(typecode_mapping);
+    static_assert(ArrLen(typecode_mapping) == (int)WasmValueType::Count);
+    // clang-format on
+
+    return typecode_mapping[static_cast<unsigned>(type)];
+}
+
 unsigned emitter::instrDesc::idCodeSize() const
 {
 #ifdef TARGET_WASM32
@@ -200,9 +260,18 @@ 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));
+            uint8_t typeCode = getWasmValueTypeCode(idl->lclType);
+            size                     = SizeOfULEB128(idl->lclCnt) + sizeof(typeCode);
+            break;
+        }
         case IF_ULEB128:
             size += idIsCnsReloc() ? PADDED_RELOC_SIZE : SizeOfULEB128(emitGetInsSC(this));
             break;
@@ -354,6 +423,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);
+            uint8_t valType = getWasmValueTypeCode(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, valType);
+            break;
+        }
         default:
             NYI_WASM("emitOutputInstr");
             break;
@@ -472,13 +559,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);
+            WasmValueType  valType = emitGetLclVarDeclType(id);
+            printf(" %llu %s", (uint64_t)imm, WasmValueTypeName(valType));
+        }
+        break;
+
         case IF_SLEB128:
         {
             cnsval_ssize_t imm = emitGetInsSC(id);

src/coreclr/jit/emitwasm.h

@@ -18,6 +18,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, WasmValueType 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 +34,10 @@ static unsigned SizeOfSLEB128(int64_t value);
 
 static unsigned emitGetAlignHintLog2(const instrDesc* id);
 
+instrDesc*               emitNewInstrLclVarDecl(emitAttr attr, cnsval_ssize_t localCount, WasmValueType type);
+static WasmValueType 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,14 +9,28 @@ enum class WasmValueType : unsigned
     F32,
     F64,
     Count,
-
 #ifdef TARGET_64BIT
     I = I64,
 #else
     I = I32,
 #endif
 };
 
+static const char *const WasmValueTypeNames[] = {
+    "Invalid",
+    "i32",
+    "i64",
+    "f32",
+    "f64",
+    "Count",
+};
+
+inline const char* WasmValueTypeName(WasmValueType type)
+{
+    static_assert(ArrLen(WasmValueTypeNames) == static_cast<unsigned>(WasmValueType::Count) + 1);
+    return WasmValueTypeNames[static_cast<unsigned>(type)];
+}
+
 regNumber MakeWasmReg(unsigned index, var_types type);
 unsigned  UnpackWasmReg(regNumber reg, WasmValueType* pType = nullptr);
 unsigned  WasmRegToIndex(regNumber reg);