feat(evm): add lazy JIT with OSR for multipass mode

src/action/compiler.cpp

@@ -46,12 +46,8 @@ void performEVMJITCompile(runtime::EVMModule &Mod) {
   switch (Mod.getRuntime()->getConfig().Mode) {
 #ifdef ZEN_ENABLE_MULTIPASS_JIT
   case common::RunMode::MultipassMode: {
-    if (Mod.getRuntime()->getConfig().EnableMultipassLazy) {
-      ZEN_LOG_WARN("EVMJIT does not support lazy compilation now");
-    } else {
-      COMPILER::EagerEVMJITCompiler ECompiler(&Mod);
-      ECompiler.compile();
-    }
+    COMPILER::EagerEVMJITCompiler ECompiler(&Mod);
+    ECompiler.compile();
     break;
   }
 #endif

src/compiler/evm_compiler.cpp

@@ -56,59 +56,67 @@ void EVMJITCompiler::compileEVMToMC(EVMFrontendContext &Ctx, MModule &Mod,
 }
 
 void EagerEVMJITCompiler::compile() {
-  auto Timer = Stats.startRecord(zen::utils::StatisticPhase::JITCompilation);
+  try {
+    auto Timer = Stats.startRecord(zen::utils::StatisticPhase::JITCompilation);
 
-  EVMFrontendContext Ctx;
-  Ctx.setGasMeteringEnabled(Config.EnableEvmGasMetering);
+    EVMFrontendContext Ctx;
+    Ctx.setGasMeteringEnabled(Config.EnableEvmGasMetering);
 #ifdef ZEN_ENABLE_EVM_GAS_REGISTER
-  Ctx.setGasRegisterEnabled(true);
+    Ctx.setGasRegisterEnabled(true);
 #endif
-  Ctx.setRevision(EVMMod->getRevision());
-  Ctx.setBytecode(reinterpret_cast<const Byte *>(EVMMod->Code),
-                  EVMMod->CodeSize);
-  const auto &Cache = EVMMod->getBytecodeCache();
-  Ctx.setGasChunkInfo(Cache.GasChunkEnd.data(), Cache.GasChunkCost.data(),
-                      EVMMod->CodeSize);
+    Ctx.setRevision(EVMMod->getRevision());
+    Ctx.setBytecode(reinterpret_cast<const Byte *>(EVMMod->Code),
+                    EVMMod->CodeSize);
+    const auto &Cache = EVMMod->getBytecodeCache();
+    Ctx.setGasChunkInfo(Cache.GasChunkEnd.data(), Cache.GasChunkCost.data(),
+                        EVMMod->CodeSize);
 
-  MModule Mod(Ctx);
-  buildEVMFunction(Ctx, Mod, *EVMMod);
-  Ctx.CodeMPool = &EVMMod->getJITCodeMemPool();
+    MModule Mod(Ctx);
+    buildEVMFunction(Ctx, Mod, *EVMMod);
+    Ctx.CodeMPool = &EVMMod->getJITCodeMemPool();
 
 #ifdef ZEN_ENABLE_LINUX_PERF
-  utils::JitDumpWriter JitDumpWriter;
+    utils::JitDumpWriter JitDumpWriter;
 #define JIT_DUMP_WRITE_FUNC(FuncName, FuncAddr, FuncSize)                      \
   JitDumpWriter.writeFunc(FuncName, reinterpret_cast<uint64_t>(FuncAddr),      \
                           FuncSize)
 #else
 #define JIT_DUMP_WRITE_FUNC(...)
 #endif // ZEN_ENABLE_LINUX_PERF
 
-  auto &CodeMPool = EVMMod->getJITCodeMemPool();
-  uint8_t *JITCode = const_cast<uint8_t *>(CodeMPool.getMemStart());
+    auto &CodeMPool = EVMMod->getJITCodeMemPool();
+    uint8_t *JITCode = const_cast<uint8_t *>(CodeMPool.getMemStart());
 
-  // EVM has only 1 function, use direct single-threaded compilation
-  compileEVMToMC(Ctx, Mod, 0, Config.DisableMultipassGreedyRA);
-  emitObjectBuffer(&Ctx);
-  ZEN_ASSERT(Ctx.ExternRelocs.empty());
+    // EVM has only 1 function, use direct single-threaded compilation
+    compileEVMToMC(Ctx, Mod, 0, Config.DisableMultipassGreedyRA);
+    emitObjectBuffer(&Ctx);
+    ZEN_ASSERT(Ctx.ExternRelocs.empty());
 
-  uint8_t *JITFuncPtr = Ctx.CodePtr + Ctx.FuncOffsetMap[0];
-  EVMMod->setJITCodeAndSize(JITFuncPtr, Ctx.CodeSize);
+    uint8_t *JITFuncPtr = Ctx.CodePtr + Ctx.FuncOffsetMap[0];
 #ifdef ZEN_ENABLE_LINUX_PERF
-  // Write block symbols instead of EVM_Main
-  // JIT_DUMP_WRITE_FUNC("EVM_Main", JITFuncPtr, Ctx.FuncSizeMap[0]);
-  for (const auto &[BBIdx, BBSymOffset] : Ctx.FuncOffsetMap) {
-    if (BBIdx == 0) {
-      continue;
+    // Write block symbols instead of EVM_Main
+    // JIT_DUMP_WRITE_FUNC("EVM_Main", JITFuncPtr, Ctx.FuncSizeMap[0]);
+    for (const auto &[BBIdx, BBSymOffset] : Ctx.FuncOffsetMap) {
+      if (BBIdx == 0) {
+        continue;
+      }
+      uint8_t *BBCode = Ctx.CodePtr + BBSymOffset;
+      JIT_DUMP_WRITE_FUNC(Ctx.FuncNameMap[BBIdx], BBCode,
+                          Ctx.FuncSizeMap[BBIdx]);
     }
-    uint8_t *BBCode = Ctx.CodePtr + BBSymOffset;
-    JIT_DUMP_WRITE_FUNC(Ctx.FuncNameMap[BBIdx], BBCode, Ctx.FuncSizeMap[BBIdx]);
-  }
 #endif
-  size_t CodeSize = CodeMPool.getMemEnd() - JITCode;
-  platform::mprotect(JITCode, TO_MPROTECT_CODE_SIZE(CodeSize),
-                     PROT_READ | PROT_EXEC);
-  EVMMod->setJITCodeAndSize(JITCode, CodeSize);
-
-  Stats.stopRecord(Timer);
+    size_t CodeSize = CodeMPool.getMemEnd() - JITCode;
+    platform::mprotect(JITCode, TO_MPROTECT_CODE_SIZE(CodeSize),
+                       PROT_READ | PROT_EXEC);
+    // Publish JITCode only after mprotect — atomic release ensures the
+    // interpreter thread sees fully executable code.
+    EVMMod->setJITCodeAndSize(JITFuncPtr, CodeSize);
+
+    Stats.stopRecord(Timer);
+  } catch (const std::exception &E) {
+    ZEN_LOG_ERROR("EVM JIT compilation failed: %s", E.what());
+  } catch (...) {
+    ZEN_LOG_ERROR("EVM JIT compilation failed");
+  }
 }
 } // namespace COMPILER

src/compiler/evm_frontend/evm_mir_compiler.cpp

@@ -279,6 +279,28 @@ void EVMMirBuilder::initEVM(CompilerContext *Context) {
   ReturnBB = createBasicBlock();
   loadEVMInstanceAttr();
 
+  // OSR entry: check if we should start at a specific PC instead of PC=0
+  const int32_t OSRPCOffset = zen::runtime::EVMInstance::getOSRPCOffset();
+  MInstruction *OSRPCVal = getInstanceElement(&Ctx.I64Type, OSRPCOffset);
+  // Clear OSRPC in instance
+  MInstruction *Zero = createIntConstInstruction(&Ctx.I64Type, 0);
+  setInstanceElement(&Ctx.I64Type, Zero, OSRPCOffset);
+  // Save for later use in OSR dispatch
+  OSRPCVar = CurFunc->createVariable(&Ctx.I64Type);
+  createInstruction<DassignInstruction>(true, &Ctx.VoidType, OSRPCVal,
+                                        OSRPCVar->getVarIdx());
+  // Create OSR dispatch BB (populated in finalizeOSR after all JUMPDESTs)
+  OSRDispatchBB = createBasicBlock();
+  MInstruction *OSRCheck = createInstruction<CmpInstruction>(
+      false, CmpInstruction::Predicate::ICMP_NE, &Ctx.I64Type, OSRPCVal, Zero);
+  MBasicBlock *NormalEntryBB = createBasicBlock();
+  createInstruction<BrIfInstruction>(true, Ctx, OSRCheck, OSRDispatchBB,
+                                     NormalEntryBB);
+  addSuccessor(OSRDispatchBB);
+  addSuccessor(NormalEntryBB);
+  setInsertBlock(NormalEntryBB);
+  // Normal execution continues from here (bytecode at PC=0)
+
   GasChunkEnd = EvmCtx->getGasChunkEnd();
   GasChunkCost = EvmCtx->getGasChunkCost();
   GasChunkSize = EvmCtx->getGasChunkSize();
@@ -296,7 +318,33 @@ void EVMMirBuilder::initEVM(CompilerContext *Context) {
 #endif // ZEN_ENABLE_LINUX_PERF
 }
 
+void EVMMirBuilder::finalizeOSR() {
+  if (!OSRDispatchBB)
+    return;
+
+  MBasicBlock *SavedBB = CurBB;
+  setInsertBlock(OSRDispatchBB);
+
+  if (IndirectJumpBB) {
+    // Set JumpTargetVar to OSRPC and dispatch via existing jump table
+    MInstruction *OSRTarget = loadVariable(OSRPCVar);
+    createInstruction<DassignInstruction>(true, &Ctx.VoidType, OSRTarget,
+                                          JumpTargetVar->getVarIdx());
+    createInstruction<BrInstruction>(true, Ctx, IndirectJumpBB);
+    addSuccessor(IndirectJumpBB);
+  } else {
+    // No JUMPDEST in bytecode — OSR at a non-existent jump target
+    MBasicBlock *FailBB =
+        getOrCreateExceptionSetBB(ErrorCode::EVMBadJumpDestination);
+    createInstruction<BrInstruction>(true, Ctx, FailBB);
+    addSuccessor(FailBB);
+  }
+
+  setInsertBlock(SavedBB);
+}
+
 void EVMMirBuilder::finalizeEVMBase() {
+  finalizeOSR();
   const auto &ExceptionSetBBs = CurFunc->getExceptionSetBBs();
 
   VariableIdx ExceptionIDIdx =

src/compiler/evm_frontend/evm_mir_compiler.h

@@ -188,6 +188,7 @@ class EVMMirBuilder final {
   void loadEVMInstanceAttr();
   void initEVM(CompilerContext *Context);
   void finalizeEVMBase();
+  void finalizeOSR();
 
   void meterOpcode(evmc_opcode Opcode, uint64_t PC);
   void meterOpcodeRange(uint64_t StartPC, uint64_t EndPCExclusive);
@@ -891,6 +892,10 @@ class EVMMirBuilder final {
   Variable *JumpTargetVar = nullptr;
   MBasicBlock *IndirectJumpBB = nullptr;
 
+  // OSR (On-Stack Replacement) support
+  Variable *OSRPCVar = nullptr;
+  MBasicBlock *OSRDispatchBB = nullptr;
+
   // Stack check block for stack overflow/underflow checking
   MBasicBlock *StackCheckBB = nullptr;
   Variable *StackTopVar = nullptr;

src/evm/interpreter.cpp

@@ -1957,6 +1957,14 @@ void BaseInterpreter::interpret() {
         break;
       }
 
+#ifdef ZEN_ENABLE_JIT
+      // OSR: if JIT code is ready, promote to JIT at this jump target
+      if (Context.getInstance()->getModule()->getJITCode()) {
+        Context.saveStateToInstance(Dest);
+        Context.OSRPromoted = true;
+        return;
+      }
+#endif
       Frame->Pc = Dest;
       continue;
     }
@@ -1987,6 +1995,14 @@ void BaseInterpreter::interpret() {
         break;
       }
 
+#ifdef ZEN_ENABLE_JIT
+      // OSR: if JIT code is ready, promote to JIT at this jump target
+      if (Context.getInstance()->getModule()->getJITCode()) {
+        Context.saveStateToInstance(Dest);
+        Context.OSRPromoted = true;
+        return;
+      }
+#endif
       Frame->Pc = Dest;
       continue;
     }
@@ -2230,3 +2246,31 @@ void InterpreterExecContext::restoreStateFromInstance(uint64_t StartPC) {
   // Reset execution status
   setStatus(EVMC_SUCCESS);
 }
+
+void InterpreterExecContext::saveStateToInstance(uint64_t TargetPC) {
+  runtime::EVMInstance *Instance = getInstance();
+  EVMFrame *Frame = getCurFrame();
+
+  // Save stack: copy EVMFrame.Stack → EVMInstance.EVMStack
+  constexpr size_t ELEMENT_SIZE = 32;
+  uint8_t *EvmStackData = const_cast<uint8_t *>(Instance->getEVMStack());
+  for (size_t I = 0; I < Frame->Sp; ++I) {
+    uint8_t *ElementData = EvmStackData + (I * ELEMENT_SIZE);
+    const intx::uint256 &Value = Frame->Stack[I];
+    for (size_t J = 0; J < ELEMENT_SIZE / 8; J++) {
+      *reinterpret_cast<uint64_t *>(ElementData) = Value[J];
+      ElementData += 8;
+    }
+  }
+  Instance->setEVMStackSize(Frame->Sp * ELEMENT_SIZE);
+
+  // Save memory: copy EVMFrame.Memory → EVMInstance memory
+  if (!Frame->Memory.empty()) {
+    Instance->expandMemoryNoGas(Frame->Memory.size());
+    std::memcpy(Instance->getMemoryBase(), Frame->Memory.data(),
+                Frame->Memory.size());
+  }
+
+  // Set OSR target PC
+  Instance->setOSRPC(TargetPC);
+}

src/evm/interpreter.h

@@ -123,6 +123,12 @@ class InterpreterExecContext {
 
   // Fallback support: restore execution state from EVMInstance
   void restoreStateFromInstance(uint64_t startPC);
+
+  // OSR support: save execution state to EVMInstance for JIT promotion
+  void saveStateToInstance(uint64_t TargetPC);
+
+  // Set when OSR promotion happened — caller should check and use JIT result
+  bool OSRPromoted = false;
 };
 
 class BaseInterpreter {

src/runtime/config.h

@@ -37,8 +37,8 @@ struct RuntimeConfig {
   bool DisableMultipassMultithread = false;
   // Number of threads for multipass JIT if DisableMultipassMultithread is false
   uint32_t NumMultipassThreads = 8;
-  // Enable multipass lazy mode(on request compile)
-  bool EnableMultipassLazy = false;
+  // Enable multipass lazy mode(background JIT compilation + OSR)
+  bool EnableMultipassLazy = true;
 #endif // ZEN_ENABLE_MULTIPASS_JIT
 
   bool validate() {

src/runtime/evm_instance.h

@@ -88,6 +88,11 @@ class EVMInstance final : public RuntimeObject<EVMInstance> {
   // ==================== Stack Methods ====================
   const uint8_t *getEVMStack() const { return EVMStack; }
   uint64_t getEVMStackSize() const { return EVMStackSize; }
+  void setEVMStackSize(uint64_t Size) { EVMStackSize = Size; }
+
+  // OSR (On-Stack Replacement) support
+  uint64_t getOSRPC() const { return OSRPC; }
+  void setOSRPC(uint64_t PC) { OSRPC = PC; }
 
   // ==================== Evmc Message Stack Methods ====================
   // Note: These methods manage the call stack for JIT host interface functions
@@ -294,6 +299,13 @@ class EVMInstance final : public RuntimeObject<EVMInstance> {
     return static_cast<int32_t>(offsetof(EVMInstance, MemorySize));
   }
 
+  static constexpr int32_t getOSRPCOffset() {
+    static_assert(offsetof(EVMInstance, OSRPC) <=
+                      std::numeric_limits<int32_t>::max(),
+                  "EVMInstance offsets should fit in 32-bit signed range");
+    return static_cast<int32_t>(offsetof(EVMInstance, OSRPC));
+  }
+
   // Capacity for EVMStack: 1024 * 256 / 8 = 32768
   static const size_t EVMStackCapacity = 32768;
 
@@ -395,6 +407,10 @@ class EVMInstance final : public RuntimeObject<EVMInstance> {
   uint8_t EVMStack[EVMStackCapacity];
   uint64_t EVMStackSize = 0;
 
+  // OSR (On-Stack Replacement): when non-zero, JIT function enters at this PC
+  // instead of PC=0. Set by interpreter before promoting to JIT.
+  uint64_t OSRPC = 0;
+
   static constexpr size_t ALIGNMENT = 8;
   alignas(16) std::array<uint8_t, HostArgScratchSize> HostArgScratch{};
 };

src/runtime/evm_module.cpp

@@ -27,6 +27,12 @@ EVMModule::EVMModule(Runtime *RT)
 }
 
 EVMModule::~EVMModule() {
+#ifdef ZEN_ENABLE_JIT
+  if (JITCompileThread.joinable()) {
+    JITCompileThread.join();
+  }
+#endif
+
   if (Name) {
     this->freeSymbol(Name);
     Name = common::WASM_SYMBOL_NULL;
@@ -66,7 +72,18 @@ EVMModuleUniquePtr EVMModule::newEVMModule(Runtime &RT,
   Mod->Host = RT.getEVMHost();
 
   if (RT.getConfig().Mode != common::RunMode::InterpMode) {
-    action::performEVMJITCompile(*Mod);
+#ifdef ZEN_ENABLE_MULTIPASS_JIT
+    if (RT.getConfig().EnableMultipassLazy) {
+      // Eagerly init bytecode cache (needed by both interpreter and JIT)
+      (void)Mod->getBytecodeCache();
+      EVMModule *RawMod = Mod.get();
+      RawMod->JITCompileThread =
+          std::thread([RawMod]() { action::performEVMJITCompile(*RawMod); });
+    } else
+#endif
+    {
+      action::performEVMJITCompile(*Mod);
+    }
   }
 
   return Mod;

src/runtime/evm_module.h

@@ -7,7 +7,9 @@
 #include "evm/evm_cache.h"
 #include "evmc/evmc.hpp"
 #include "runtime/module.h"
+#include <atomic>
 #include <limits>
+#include <thread>
 
 #ifdef ZEN_ENABLE_JIT
 namespace COMPILER {
@@ -47,13 +49,13 @@ class EVMModule final : public BaseModule<EVMModule> {
 #ifdef ZEN_ENABLE_JIT
   common::CodeMemPool &getJITCodeMemPool() { return JITCodeMemPool; }
 
-  void *getJITCode() const { return JITCode; }
+  void *getJITCode() const { return JITCode.load(std::memory_order_acquire); }
 
   size_t getJITCodeSize() const { return JITCodeSize; }
 
   void setJITCodeAndSize(void *Code, size_t Size) {
-    JITCode = Code;
     JITCodeSize = Size;
+    JITCode.store(Code, std::memory_order_release);
   }
 #endif // ZEN_ENABLE_JIT
 
@@ -72,8 +74,9 @@ class EVMModule final : public BaseModule<EVMModule> {
 
 #ifdef ZEN_ENABLE_JIT
   common::CodeMemPool JITCodeMemPool;
-  void *JITCode = nullptr;
+  std::atomic<void *> JITCode{nullptr};
   size_t JITCodeSize = 0;
+  std::thread JITCompileThread;
 #endif // ZEN_ENABLE_JIT
 };