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

src/coreclr/jit/codegen.h

@@ -589,6 +589,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,7 +5650,7 @@ void CodeGen::genFnProlog()
 #endif // defined(DEBUG) && defined(TARGET_XARCH)
 
 #else  // defined(TARGET_WASM)
-    // TODO-WASM: prolog zeroing, shadow stack maintenance
+    genWasmLocals();
     GetEmitter()->emitMarkPrologEnd();
 #endif // !defined(TARGET_WASM)
 

src/coreclr/jit/codegenwasm.cpp

@@ -15,6 +15,26 @@ void CodeGen::genMarkLabelsForCodegen()
     // We mark labels as needed in genEmitStartBlock.
 }
 
+//------------------------------------------------------------------------
+// genWasmLocals: generate wasm local declarations
+//
+// TODO-WASM: pre-declare all "register" locals
+void CodeGen::genWasmLocals()
+{
+    // 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;
+    }
+}
+
 void CodeGen::genFnEpilog(BasicBlock* block)
 {
 #ifdef DEBUG
@@ -39,9 +59,16 @@ 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: we need to handle the end-of-function case if we reach the end of a codegen for a function
+    // and do NOT have an epilog. In those cases we currently will not emit an end instruction.
+    if (block->IsLast() || compiler->bbIsFuncletBeg(block->Next()))
+    {
+        instGen(INS_end);
+    }
+    else
+    {
+        instGen(INS_return);
+    }
 }
 
 void CodeGen::genCaptureFuncletPrologEpilogInfo()

src/coreclr/jit/emit.h

@@ -638,6 +638,9 @@ class emitter
         bool              idCatchRet;    // Instruction is for a catch 'return'
         CORINFO_SIG_INFO* idCallSig;     // Used to report native call site signatures to the EE
         BasicBlock*       idTargetBlock; // Target block for branches
+#ifdef TARGET_WASM
+        int lclOffset; // Base index of the WASM locals being declared
+#endif
     };
 
 #ifdef TARGET_ARM
@@ -1296,6 +1299,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 +2345,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 +3322,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,16 @@ 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(RAW_ULEB128, 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_DECL,  IS_NONE, NONE) // <ULEB128 immediate> <byte>
 
 #undef IF_DEF
 #endif // !DEFINE_ID_OPS

src/coreclr/jit/emitwasm.cpp

@@ -116,6 +116,70 @@ bool emitter::emitInsIsStore(instruction ins)
     return false;
 }
 
+//-----------------------------------------------------------------------------
+// emitNewInstrLclVarDecl: Construct an instrDesc corresponding to a wasm local
+// declaration.
+//
+// Arguments:
+//   attr        - emit attributes
+//   localCount  - the count of locals in this declaration
+//   type        - the type of local in the declaration
+//   lclOffset   - used to provide the starting index of this local
+//
+// Notes:
+//   `lclOffset` is stored as debug info attached to the instruction,
+//    so the offset will only be used if m_debugInfoSize > 0
+emitter::instrDesc* emitter::emitNewInstrLclVarDecl(emitAttr       attr,
+                                                    cnsval_ssize_t localCount,
+                                                    WasmValueType  type,
+                                                    int            lclOffset)
+{
+    instrDescLclVarDecl* id = static_cast<instrDescLclVarDecl*>(emitAllocAnyInstr(sizeof(instrDescLclVarDecl), attr));
+    id->idLclCnt(localCount);
+    id->idLclType(type);
+
+    if (m_debugInfoSize > 0)
+    {
+        id->idDebugOnlyInfo()->lclOffset = lclOffset;
+    }
+
+    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.
+//
+// Arguments:
+//   ins      - instruction to emit
+//   imm      - immediate value (local count)
+//   valType  - value type of the local variable
+//   offs     - local variable offset (= count of preceding locals) for debug info
+void emitter::emitIns_I_Ty(instruction ins, cnsval_ssize_t imm, WasmValueType valType, int offs)
+{
+    instrDesc* id  = emitNewInstrLclVarDecl(EA_8BYTE, imm, valType, offs);
+    insFormat  fmt = emitInsFormat(ins);
+
+    id->idIns(ins);
+    id->idInsFmt(fmt);
+
+    dispIns(id);
+    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 +219,11 @@ size_t emitter::emitSizeOfInsDsc(instrDesc* id) const
         return sizeof(instrDescCns);
     }
 
+    if (id->idIsLclVarDecl())
+    {
+        return sizeof(instrDescLclVarDecl);
+    }
+
     return sizeof(instrDesc);
 }
 
@@ -181,6 +250,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
@@ -199,10 +285,18 @@ unsigned emitter::instrDesc::idCodeSize() const
         case IF_BLOCK:
             size += 1;
             break;
-        case IF_LABEL:
+        case IF_RAW_ULEB128:
             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;
@@ -337,7 +431,7 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
             dst += emitRawBytes(dst, &bits, sizeof(cnsval_ssize_t));
             break;
         }
-        case IF_LABEL:
+        case IF_RAW_ULEB128:
         {
             cnsval_ssize_t constant = emitGetInsSC(id);
             dst += emitOutputULEB128(dst, (uint64_t)constant);
@@ -354,6 +448,15 @@ size_t emitter::emitOutputInstr(insGroup* ig, instrDesc* id, BYTE** dp)
             dst += emitOutputULEB128(dst, offset);
             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);
+            dst += emitOutputByte(dst, valType);
+            break;
+        }
         default:
             NYI_WASM("emitOutputInstr");
             break;
@@ -470,7 +573,7 @@ void emitter::emitDispIns(
         case IF_BLOCK:
             break;
 
-        case IF_LABEL:
+        case IF_RAW_ULEB128:
         case IF_ULEB128:
         {
             cnsval_ssize_t imm = emitGetInsSC(id);
@@ -479,6 +582,34 @@ void emitter::emitDispIns(
         }
         break;
 
+        case IF_LOCAL_DECL:
+        {
+            cnsval_ssize_t count   = emitGetLclVarDeclCount(id);
+            WasmValueType  valType = emitGetLclVarDeclType(id);
+            int            offs    = id->idDebugOnlyInfo()->lclOffset;
+            assert(count > 0); // we should not be declaring a local entry with zero count
+
+            if (m_debugInfoSize > 0)
+            {
+                // With debug info: print the local offsets being declared
+                if (count > 1)
+                {
+                    printf("[%llu..%llu] type=%s", offs, offs + (uint64_t)count - 1, WasmValueTypeName(valType));
+                }
+                else // single local case
+                {
+                    printf("[%llu] type=%s", offs, WasmValueTypeName(valType));
+                }
+            }
+            else
+            {
+                // No debug info case: just print the count and type of the locals
+                printf(" count=%llu type=%s", (uint64_t)count, WasmValueTypeName(valType));
+                break;
+            }
+        }
+        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, int offs);
 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, int lclOffset);
+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

@@ -26,7 +26,9 @@
 //
 INST(invalid,     "INVALID",     0, IF_NONE,    BAD_CODE)
 INST(unreachable, "unreachable", 0, IF_OPCODE,  0x00)
-INST(label,       "label",       0, IF_LABEL,   0x00)
+INST(label,       "label",       0, IF_RAW_ULEB128,   0x00)
+INST(local_cnt,   "local.cnt",   0, IF_RAW_ULEB128, 0x00)
+INST(local_decl,  "local",       0, IF_LOCAL_DECL, 0x00)
 INST(nop,         "nop",         0, IF_OPCODE,  0x01)
 INST(block,       "block",       0, IF_BLOCK,   0x02)
 INST(loop,        "loop",        0, IF_BLOCK,   0x03)

src/coreclr/jit/registeropswasm.cpp

@@ -61,6 +61,23 @@ regNumber MakeWasmReg(unsigned index, var_types type)
     return static_cast<regNumber>(regValue);
 }
 
+const char* WasmValueTypeName(WasmValueType type)
+{
+    // clang-format off
+    static const char* const WasmValueTypeNames[] = {
+        "Invalid",
+        "i32",
+        "i64",
+        "f32",
+        "f64",
+    };
+    static_assert(ArrLen(WasmValueTypeNames) == static_cast<unsigned>(WasmValueType::Count));
+    // clang-format on
+
+    assert(WasmValueType::Invalid <= type && type < WasmValueType::Count);
+    return WasmValueTypeNames[(unsigned)type];
+}
+
 //------------------------------------------------------------------------
 // UnpackWasmReg: Extract the WASM local's index and type out of a register
 //