From 8a7c0b5be13e3b191bec86c6897977eb2f3769ab Mon Sep 17 00:00:00 2001 From: cl507523 Date: Thu, 2 Apr 2026 07:54:29 +0000 Subject: [PATCH] feat(evm): add lazy JIT with OSR for multipass mode Background JIT compilation runs in a separate thread while the interpreter handles the first invocation. At JUMP/JUMPI, the interpreter checks whether JIT code has become available and promotes execution mid-flight via On-Stack Replacement (saving stack/memory to EVMInstance and entering the JIT at the target PC through its indirect jump table). Enabled by default; disable with DTVM_EVM_LAZY_JIT=false or set_option("lazy","false"). Co-Authored-By: Claude Opus 4.6 --- src/action/compiler.cpp | 8 +- src/compiler/evm_compiler.cpp | 80 ++++++++++--------- .../evm_frontend/evm_mir_compiler.cpp | 48 +++++++++++ src/compiler/evm_frontend/evm_mir_compiler.h | 5 ++ src/evm/interpreter.cpp | 44 ++++++++++ src/evm/interpreter.h | 6 ++ src/runtime/config.h | 4 +- src/runtime/evm_instance.h | 16 ++++ src/runtime/evm_module.cpp | 19 ++++- src/runtime/evm_module.h | 9 ++- src/runtime/runtime.cpp | 17 +++- src/vm/dt_evmc_vm.cpp | 22 +++++ 12 files changed, 229 insertions(+), 49 deletions(-) diff --git a/src/action/compiler.cpp b/src/action/compiler.cpp index 33317e03e..69126438b 100644 --- a/src/action/compiler.cpp +++ b/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 diff --git a/src/compiler/evm_compiler.cpp b/src/compiler/evm_compiler.cpp index 04d45ad60..171ba629f 100644 --- a/src/compiler/evm_compiler.cpp +++ b/src/compiler/evm_compiler.cpp @@ -56,26 +56,27 @@ 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(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(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(FuncAddr), \ FuncSize) @@ -83,32 +84,39 @@ void EagerEVMJITCompiler::compile() { #define JIT_DUMP_WRITE_FUNC(...) #endif // ZEN_ENABLE_LINUX_PERF - auto &CodeMPool = EVMMod->getJITCodeMemPool(); - uint8_t *JITCode = const_cast(CodeMPool.getMemStart()); + auto &CodeMPool = EVMMod->getJITCodeMemPool(); + uint8_t *JITCode = const_cast(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 diff --git a/src/compiler/evm_frontend/evm_mir_compiler.cpp b/src/compiler/evm_frontend/evm_mir_compiler.cpp index cd98f29cb..74336bd1c 100644 --- a/src/compiler/evm_frontend/evm_mir_compiler.cpp +++ b/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(true, &Ctx.VoidType, OSRPCVal, + OSRPCVar->getVarIdx()); + // Create OSR dispatch BB (populated in finalizeOSR after all JUMPDESTs) + OSRDispatchBB = createBasicBlock(); + MInstruction *OSRCheck = createInstruction( + false, CmpInstruction::Predicate::ICMP_NE, &Ctx.I64Type, OSRPCVal, Zero); + MBasicBlock *NormalEntryBB = createBasicBlock(); + createInstruction(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(true, &Ctx.VoidType, OSRTarget, + JumpTargetVar->getVarIdx()); + createInstruction(true, Ctx, IndirectJumpBB); + addSuccessor(IndirectJumpBB); + } else { + // No JUMPDEST in bytecode — OSR at a non-existent jump target + MBasicBlock *FailBB = + getOrCreateExceptionSetBB(ErrorCode::EVMBadJumpDestination); + createInstruction(true, Ctx, FailBB); + addSuccessor(FailBB); + } + + setInsertBlock(SavedBB); +} + void EVMMirBuilder::finalizeEVMBase() { + finalizeOSR(); const auto &ExceptionSetBBs = CurFunc->getExceptionSetBBs(); VariableIdx ExceptionIDIdx = diff --git a/src/compiler/evm_frontend/evm_mir_compiler.h b/src/compiler/evm_frontend/evm_mir_compiler.h index 13124a41a..541c1b591 100644 --- a/src/compiler/evm_frontend/evm_mir_compiler.h +++ b/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; diff --git a/src/evm/interpreter.cpp b/src/evm/interpreter.cpp index d67e32cb9..0b6b233e1 100644 --- a/src/evm/interpreter.cpp +++ b/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(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(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); +} diff --git a/src/evm/interpreter.h b/src/evm/interpreter.h index 35e6ec24c..f2bee57c3 100644 --- a/src/evm/interpreter.h +++ b/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 { diff --git a/src/runtime/config.h b/src/runtime/config.h index c8a9b301f..3503f45c8 100644 --- a/src/runtime/config.h +++ b/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() { diff --git a/src/runtime/evm_instance.h b/src/runtime/evm_instance.h index 4e1d99cd0..d13b2759f 100644 --- a/src/runtime/evm_instance.h +++ b/src/runtime/evm_instance.h @@ -88,6 +88,11 @@ class EVMInstance final : public RuntimeObject { // ==================== 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 { return static_cast(offsetof(EVMInstance, MemorySize)); } + static constexpr int32_t getOSRPCOffset() { + static_assert(offsetof(EVMInstance, OSRPC) <= + std::numeric_limits::max(), + "EVMInstance offsets should fit in 32-bit signed range"); + return static_cast(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 { 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 HostArgScratch{}; }; diff --git a/src/runtime/evm_module.cpp b/src/runtime/evm_module.cpp index 673054110..7cb5578a3 100644 --- a/src/runtime/evm_module.cpp +++ b/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; diff --git a/src/runtime/evm_module.h b/src/runtime/evm_module.h index 43094717e..12ce5cf40 100644 --- a/src/runtime/evm_module.h +++ b/src/runtime/evm_module.h @@ -7,7 +7,9 @@ #include "evm/evm_cache.h" #include "evmc/evmc.hpp" #include "runtime/module.h" +#include #include +#include #ifdef ZEN_ENABLE_JIT namespace COMPILER { @@ -47,13 +49,13 @@ class EVMModule final : public BaseModule { #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 { #ifdef ZEN_ENABLE_JIT common::CodeMemPool JITCodeMemPool; - void *JITCode = nullptr; + std::atomic JITCode{nullptr}; size_t JITCodeSize = 0; + std::thread JITCompileThread; #endif // ZEN_ENABLE_JIT }; diff --git a/src/runtime/runtime.cpp b/src/runtime/runtime.cpp index 0f2a116cb..9963c3408 100644 --- a/src/runtime/runtime.cpp +++ b/src/runtime/runtime.cpp @@ -691,8 +691,23 @@ void Runtime::callEVMMainOnPhysStack(EVMInstance &Inst, evmc_message &Msg, MsgWithCode.code_size = Inst.getModule()->CodeSize; Inst.setExeResult(evmc::Result{EVMC_SUCCESS, 0, 0}); Inst.pushMessage(&MsgWithCode); - if (getConfig().Mode == RunMode::InterpMode) { + + bool UseInterpreter = (getConfig().Mode == RunMode::InterpMode); +#ifdef ZEN_ENABLE_JIT + if (!UseInterpreter && !Inst.getModule()->getJITCode()) { + UseInterpreter = true; + } +#endif + + if (UseInterpreter) { callEVMInInterpMode(Inst, MsgWithCode, Result); +#ifdef ZEN_ENABLE_JIT + // OSR: interpreter promoted to JIT mid-execution + if (Inst.getOSRPC() != 0) { + callEVMInJITMode(Inst, MsgWithCode, Result); + Inst.setOSRPC(0); + } +#endif } else { #ifdef ZEN_ENABLE_JIT callEVMInJITMode(Inst, MsgWithCode, Result); diff --git a/src/vm/dt_evmc_vm.cpp b/src/vm/dt_evmc_vm.cpp index 57f6e4d72..deeb93b4a 100644 --- a/src/vm/dt_evmc_vm.cpp +++ b/src/vm/dt_evmc_vm.cpp @@ -264,6 +264,17 @@ enum evmc_set_option_result set_option(evmc_vm *VMInstance, const char *Name, } else { return EVMC_SET_OPTION_INVALID_VALUE; } +#ifdef ZEN_ENABLE_MULTIPASS_JIT + } else if (std::strcmp(Name, "lazy") == 0) { + if (std::strcmp(Value, "true") == 0) { + VM->Config.EnableMultipassLazy = true; + return EVMC_SET_OPTION_SUCCESS; + } else if (std::strcmp(Value, "false") == 0) { + VM->Config.EnableMultipassLazy = false; + return EVMC_SET_OPTION_SUCCESS; + } + return EVMC_SET_OPTION_INVALID_VALUE; +#endif } return EVMC_SET_OPTION_INVALID_NAME; } @@ -601,6 +612,17 @@ DTVM::DTVM() ZEN_LOG_WARN("ignore invalid DTVM_EVM_ENABLE_GAS_METERING=%s", EnableGas); } } + +#ifdef ZEN_ENABLE_MULTIPASS_JIT + if (const char *Lazy = std::getenv("DTVM_EVM_LAZY_JIT"); Lazy != nullptr) { + bool ParsedLazy = false; + if (parseBoolEnvValue(Lazy, ParsedLazy)) { + Config.EnableMultipassLazy = ParsedLazy; + } else { + ZEN_LOG_WARN("ignore invalid DTVM_EVM_LAZY_JIT=%s", Lazy); + } + } +#endif } } // namespace