[RyuJit/WASM] Register allocator skeleton (#122414)

Force all locals to stack and add an interface for the SP/FP.

Implement `GT_LCL_ADDR` / `GT_LCL_FLD` as samples.

Next step is to add the stackifier (it'll run before RA and liveness,
since it may need to introduce locals and we want to be able to
enregister them). After that - parameter homing, and that should bring
us to syntactically (if not yet semantically) valid WASM.

---------

Co-authored-by: Andy Ayers <[email protected]>

Author: SingleAccretion

src/coreclr/jit/CMakeLists.txt

@@ -162,6 +162,7 @@ set( JIT_SOURCES
   rangecheckcloning.cpp
   rationalize.cpp
   redundantbranchopts.cpp
+  regalloc.cpp
   regset.cpp
   scev.cpp
   scopeinfo.cpp
@@ -182,7 +183,6 @@ set( JIT_SOURCES
 set ( JIT_NATIVE_TARGET_SOURCES
   lsra.cpp
   lsrabuild.cpp
-  regalloc.cpp
   regMaskTPOps.cpp
   gcdecode.cpp
   gcencode.cpp
@@ -385,6 +385,7 @@ set( JIT_HEADERS
   rangecheckcloning.h
   rationalize.h
   regalloc.h
+  regallocimpl.h
   register.h
   regset.h
   scev.h

src/coreclr/jit/codegeninterface.h

@@ -254,6 +254,42 @@ class CodeGenInterface
         m_cgFrameRequired = value;
     }
 
+#if !HAS_FIXED_REGISTER_SET
+private:
+    // For targets without a fixed SP/FP, these are the registers with which they are associated.
+    PhasedVar<regNumber> m_cgStackPointerReg = REG_NA;
+    PhasedVar<regNumber> m_cgFramePointerReg = REG_NA;
+
+public:
+    void SetStackPointerReg(regNumber reg)
+    {
+        assert(reg != REG_NA);
+        m_cgStackPointerReg = reg;
+    }
+    void SetFramePointerReg(regNumber reg)
+    {
+        assert(reg != REG_NA);
+        m_cgFramePointerReg = reg;
+    }
+    regNumber GetStackPointerReg() const
+    {
+        return m_cgStackPointerReg;
+    }
+    regNumber GetFramePointerReg() const
+    {
+        return m_cgFramePointerReg;
+    }
+#else  // HAS_FIXED_REGISTER_SET
+    regNumber GetStackPointerReg() const
+    {
+        return REG_SPBASE;
+    }
+    regNumber GetFramePointerReg() const
+    {
+        return REG_FPBASE;
+    }
+#endif // HAS_FIXED_REGISTER_SET
+
 public:
     int genCallerSPtoFPdelta() const;
     int genCallerSPtoInitialSPdelta() const;

src/coreclr/jit/codegenwasm.cpp

@@ -9,6 +9,14 @@
 #include "codegen.h"
 #include "fgwasm.h"
 
+#ifdef TARGET_64BIT
+static const instruction INS_I_const = INS_i64_const;
+static const instruction INS_I_add   = INS_i64_add;
+#else  // !TARGET_64BIT
+static const instruction INS_I_const = INS_i32_const;
+static const instruction INS_I_add   = INS_i32_add;
+#endif // !TARGET_64BIT
+
 void CodeGen::genMarkLabelsForCodegen()
 {
     // No work needed here for now.
@@ -268,6 +276,14 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genCodeForCompare(treeNode->AsOp());
             break;
 
+        case GT_LCL_ADDR:
+            genCodeForLclAddr(treeNode->AsLclFld());
+            break;
+
+        case GT_LCL_FLD:
+            genCodeForLclFld(treeNode->AsLclFld());
+            break;
+
         case GT_LCL_VAR:
             genCodeForLclVar(treeNode->AsLclVar());
             break;
@@ -741,6 +757,38 @@ void CodeGen::genCodeForNegNot(GenTreeOp* tree)
     genProduceReg(tree);
 }
 
+//------------------------------------------------------------------------
+// genCodeForLclAddr: Generates the code for GT_LCL_ADDR.
+//
+// Arguments:
+//    lclAddrNode - the node.
+//
+void CodeGen::genCodeForLclAddr(GenTreeLclFld* lclAddrNode)
+{
+    assert(lclAddrNode->OperIs(GT_LCL_ADDR));
+
+    GetEmitter()->emitIns_I(INS_local_get, EA_PTRSIZE, WasmRegToIndex(GetFramePointerReg()));
+    GetEmitter()->emitIns_S(INS_I_const, EA_PTRSIZE, lclAddrNode->GetLclNum(), lclAddrNode->GetLclOffs());
+    GetEmitter()->emitIns(INS_I_add);
+    genProduceReg(lclAddrNode);
+}
+
+//------------------------------------------------------------------------
+// genCodeForLclFld: Produce code for a GT_LCL_FLD node.
+//
+// Arguments:
+//    tree - the GT_LCL_FLD node
+//
+void CodeGen::genCodeForLclFld(GenTreeLclFld* tree)
+{
+    assert(tree->OperIs(GT_LCL_FLD));
+    LclVarDsc* varDsc = compiler->lvaGetDesc(tree);
+
+    GetEmitter()->emitIns_I(INS_local_get, EA_PTRSIZE, WasmRegToIndex(GetFramePointerReg()));
+    GetEmitter()->emitIns_S(ins_Load(tree->TypeGet()), emitTypeSize(tree), tree->GetLclNum(), tree->GetLclOffs());
+    genProduceReg(tree);
+}
+
 //------------------------------------------------------------------------
 // genCodeForLclVar: Produce code for a GT_LCL_VAR node.
 //
@@ -759,14 +807,14 @@ void CodeGen::genCodeForLclVar(GenTreeLclVar* tree)
     if (!varDsc->lvIsRegCandidate())
     {
         var_types type = varDsc->GetRegisterType(tree);
-        // TODO-WASM: actually local.get the frame base local here.
+        GetEmitter()->emitIns_I(INS_local_get, EA_PTRSIZE, WasmRegToIndex(GetFramePointerReg()));
         GetEmitter()->emitIns_S(ins_Load(type), emitTypeSize(tree), tree->GetLclNum(), 0);
         genProduceReg(tree);
     }
     else
     {
         assert(genIsValidReg(varDsc->GetRegNum()));
-        unsigned wasmLclIndex = UnpackWasmReg(varDsc->GetRegNum());
+        unsigned wasmLclIndex = WasmRegToIndex(varDsc->GetRegNum());
         GetEmitter()->emitIns_I(INS_local_get, emitTypeSize(tree), wasmLclIndex);
     }
 }
@@ -785,23 +833,15 @@ void CodeGen::genCodeForStoreLclVar(GenTreeLclVar* tree)
     assert(!op1->IsMultiRegNode());
     genConsumeRegs(op1);
 
+    // We rewrite all stack stores to STOREIND because the address must be first on the operand stack, so here only
+    // enregistered locals need to be handled.
     LclVarDsc* varDsc    = compiler->lvaGetDesc(tree);
     regNumber  targetReg = varDsc->GetRegNum();
+    assert(genIsValidReg(targetReg) && varDsc->lvIsRegCandidate());
 
-    if (!varDsc->lvIsRegCandidate())
-    {
-        // TODO-WASM: handle these cases in lower/ra.
-        // Emit drop for now to simulate the store effect on the wasm stack.
-        GetEmitter()->emitIns(INS_drop);
-        genUpdateLife(tree);
-    }
-    else
-    {
-        assert(genIsValidReg(targetReg));
-        unsigned wasmLclIndex = UnpackWasmReg(targetReg);
-        GetEmitter()->emitIns_I(INS_local_set, emitTypeSize(tree), wasmLclIndex);
-        genUpdateLifeStore(tree, targetReg, varDsc);
-    }
+    unsigned wasmLclIndex = WasmRegToIndex(targetReg);
+    GetEmitter()->emitIns_I(INS_local_set, emitTypeSize(tree), wasmLclIndex);
+    genUpdateLifeStore(tree, targetReg, varDsc);
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/emitwasm.cpp

@@ -73,7 +73,7 @@ void emitter::emitIns_J(instruction ins, emitAttr attr, cnsval_ssize_t imm, Basi
 }
 
 //------------------------------------------------------------------------
-// emitIns_S: Emit a memory instruction with a stack-based address mode operand.
+// emitIns_S: Emit an instruction with a stack offset immediate.
 //
 void emitter::emitIns_S(instruction ins, emitAttr attr, int varx, int offs)
 {
@@ -498,8 +498,7 @@ void emitter::emitDispIns(
 
         case IF_MEMARG:
         {
-            // TODO-WASM: decide what our strategy for alignment hints is and display these accordingly.
-            unsigned       log2align = emitGetAlignHintLog2(id) + 1;
+            unsigned       log2align = emitGetAlignHintLog2(id);
             cnsval_ssize_t offset    = emitGetInsSC(id);
             printf(" %u %llu", log2align, (uint64_t)offset);
         }

src/coreclr/jit/lclvars.cpp

@@ -4609,7 +4609,14 @@ void Compiler::lvaFixVirtualFrameOffsets()
 
         JITDUMP("--- delta bump %d for FP frame\n", delta);
     }
-#endif // !TARGET_LOONGARCH64 || !TARGET_RISCV64
+#elif defined(TARGET_WASM)
+    else
+    {
+        // The FP always points at the bottom of the fixed portion of the frame.
+        JITDUMP("--- delta bump %d for FP frame\n", codeGen->genTotalFrameSize());
+        delta += codeGen->genTotalFrameSize();
+    }
+#endif
 
     if (opts.IsOSR())
     {
@@ -6260,7 +6267,7 @@ void Compiler::lvaDumpFrameLocation(unsigned lclNum, int minLength)
 #else
     bool EBPbased;
     offset  = lvaFrameAddress(lclNum, &EBPbased);
-    baseReg = EBPbased ? REG_FPBASE : REG_SPBASE;
+    baseReg = EBPbased ? codeGen->GetFramePointerReg() : codeGen->GetStackPointerReg();
 #endif
 
     int printed =

src/coreclr/jit/lower.h

@@ -17,14 +17,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 #include "compiler.h"
 #include "phase.h"
 #include "sideeffects.h"
-
-#if HAS_FIXED_REGISTER_SET
-#include "lsra.h"
-#endif
-
-#ifdef TARGET_WASM
-#include "regallocwasm.h"
-#endif
+#include "regallocimpl.h"
 
 class Lowering final : public Phase
 {