[Wasm RyuJit] fix funclet prolog and epilog codegen

src/coreclr/jit/codegen.h

@@ -608,7 +608,7 @@ class CodeGen final : public CodeGenInterface
     void genReserveFuncletProlog(BasicBlock* block);
     void genReserveFuncletEpilog(BasicBlock* block);
     void genFuncletProlog(BasicBlock* block);
-    void genFuncletEpilog();
+    void genFuncletEpilog(BasicBlock* block);
     void genCaptureFuncletPrologEpilogInfo();
 
     void genUpdateCurrentFunclet(BasicBlock* block);

src/coreclr/jit/codegenarm.cpp

@@ -2371,7 +2371,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *  Generates code for an EH funclet epilog.
  */
 
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)

src/coreclr/jit/codegenarm64.cpp

@@ -1495,7 +1495,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *  See the description of frame shapes at genFuncletProlog().
  */
 
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)

src/coreclr/jit/codegencommon.cpp

@@ -2014,6 +2014,18 @@ void CodeGen::genGenerateCode(void** codePtr, uint32_t* nativeSizeOfCode)
     }
 #endif // defined(TARGET_WASM)
 #endif // DEBUG
+
+#if defined(TARGET_WASM)
+    // Also fail at this point for any method with funclets, since the Wasm we produce
+    // for such methods requires post-processing by the host before it can be validated.
+    // TODO-WASM: Remove this once the host can do the processing.
+    //
+    if ((JitConfig.JitWasmFunclets() == 0) && (m_compiler->compFuncCount() > 1))
+    {
+        JITDUMP("Failing R2R codegen because method has funclets.\n");
+        implReadyToRunUnsupported();
+    }
+#endif
 }
 
 //----------------------------------------------------------------------

src/coreclr/jit/codegenlinear.cpp

@@ -420,8 +420,6 @@ void CodeGen::genCodeForBlock(BasicBlock* block)
         GetEmitter()->emitSetFirstColdIGCookie(block->bbEmitCookie);
     }
 
-    genEmitStartBlock(block);
-
     // Both stacks are always empty on entry to a basic block.
     assert(genStackLevel == 0);
     genAdjustStackLevel(block);
@@ -444,6 +442,8 @@ void CodeGen::genCodeForBlock(BasicBlock* block)
         genReserveFuncletProlog(block);
     }
 
+    genEmitStartBlock(block);
+
     // Clear compCurStmt and compCurLifeTree.
     m_compiler->compCurStmt     = nullptr;
     m_compiler->compCurLifeTree = nullptr;

src/coreclr/jit/codegenloongarch64.cpp

@@ -475,7 +475,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *  Generates code for an EH funclet epilog.
  */
 
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)

src/coreclr/jit/codegenriscv64.cpp

@@ -463,7 +463,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *
  *  See the description of frame shapes at genFuncletProlog().
  */
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)

src/coreclr/jit/codegenwasm.cpp

@@ -311,10 +311,21 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
 //------------------------------------------------------------------------
 // genFuncletEpilog: codegen for funclet epilogs.
 //
-// For Wasm, funclet epilogs are empty
+// Arguments:
+//   block - funclet epilog block
 //
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* block)
 {
+    ScopedSetVariable<bool> _setGeneratingEpilog(&m_compiler->compGeneratingEpilog, true);
+
+    if (block->IsLast() || m_compiler->bbIsFuncletBeg(block->Next()))
+    {
+        instGen(INS_end);
+    }
+    else
+    {
+        instGen(INS_return);
+    }
 }
 
 //------------------------------------------------------------------------

src/coreclr/jit/codegenxarch.cpp

@@ -10865,7 +10865,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *  Note that we don't do anything with unwind codes, because AMD64 only cares about unwind codes for the prolog.
  */
 
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)
@@ -10992,7 +10992,7 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
  *  Generates code for an EH funclet epilog.
  */
 
-void CodeGen::genFuncletEpilog()
+void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
 {
 #ifdef DEBUG
     if (verbose)

src/coreclr/jit/emit.cpp

@@ -2226,7 +2226,7 @@ void emitter::emitGeneratePrologEpilog()
             case IGPT_FUNCLET_EPILOG:
                 INDEBUG(++funcletEpilogCnt);
                 emitBegFuncletEpilog(igPh);
-                codeGen->genFuncletEpilog();
+                codeGen->genFuncletEpilog(igPhBB);
                 emitEndFuncletEpilog();
                 break;
 

src/coreclr/jit/jitconfigvalues.h

@@ -882,6 +882,8 @@ CONFIG_INTEGER(JitWasmNyiToR2RUnsupported, "JitWasmNyiToR2RUnsupported", 0)
 // Specify methods that will fail with R2R unsupported after codegen.
 // Useful for bypassing methods that compile cleanly but have invalid Wasm codegen.
 CONFIG_STRING(JitR2RUnsupportedRange, "JitR2RUnsupportedRange")
+// Enable processing methods with funclets.
+RELEASE_CONFIG_INTEGER(JitWasmFunclets, "JitWasmFunclets", 0)
 #endif // defined(TARGET_WASM)
 
 // Allow to enregister locals with struct type.

src/coreclr/jit/regallocwasm.cpp

@@ -713,7 +713,7 @@ void WasmRegAlloc::CollectReference(GenTree* node)
     PerFuncletData* const data = m_perFuncletData[m_currentFunclet];
     VirtualRegReferences* refs = data->m_virtualRegRefs;
 
-    // We may make multiple consecutive collection calls for the same node.
+    // We may make multiple collection calls for the same node.
     // We only want to collect it once.
     //
     if (data->m_lastVirtualRegRefsCount > 0)
@@ -777,6 +777,8 @@ void WasmRegAlloc::RequestTemporaryRegisterForMultiplyUsedNode(GenTree* node)
     regNumber reg = AllocateTemporaryRegister(node->TypeGet());
     assert((node->GetRegNum() == REG_NA) && "Trying to double-assign a temporary register");
     node->SetRegNum(reg);
+
+    CollectReference(node);
 }
 
 //------------------------------------------------------------------------
@@ -799,7 +801,6 @@ void WasmRegAlloc::ConsumeTemporaryRegForOperand(GenTree* operand DEBUGARG(const
 
     regNumber reg = ReleaseTemporaryRegister(genActualType(operand));
     assert((reg == operand->GetRegNum()) && "Temporary reg being consumed out of order");
-    CollectReference(operand);
 
     operand->gtLIRFlags &= ~LIR::Flags::MultiplyUsed;
     JITDUMP("Consumed a temporary reg for [%06u]: %s\n", Compiler::dspTreeID(operand), reason);
@@ -819,6 +820,8 @@ void WasmRegAlloc::ConsumeTemporaryRegForOperand(GenTree* operand DEBUGARG(const
 //
 regNumber WasmRegAlloc::RequestInternalRegister(GenTree* node, var_types type)
 {
+    JITDUMP("Requesting internal %s register for [%06u]\n", varTypeName(type), Compiler::dspTreeID(node));
+
     regNumber reg = AllocateTemporaryRegister(type);
     m_codeGen->internalRegisters.Add(node, reg);
     CollectReference(node);