From 0d20dff30d761584ad87ff1727376bcd9c577c08 Mon Sep 17 00:00:00 2001 From: tamnd <1218621+tamnd@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:16:52 +0700 Subject: [PATCH] Give the COFF JIT somewhere to find libc (#322) mojo run on Windows compiles a program, links it, and then cannot resolve __chkstk, memcpy, _setmode, _write and five kernel32 calls. All nine exist on the machine. Nothing was telling the JIT where to look. Two things were missing. The platform stdlib, which is the dylib holding the process wide search generator, was in the search order a lookup starts from but in nobody's link order, and the link order is what resolution walks once materialization is under way. That gap never showed on ELF or Mach-O because dlsym on the compiler runtime handle keeps walking into that library's own dependencies, and libc is one of them. GetProcAddress reads one module's export table and stops, so on Windows the platform stdlib is the only thing that can answer and it has to be reachable. The second is that the process wide generator is not enough on its own here. A Windows build of mojo links the C runtime statically, so ucrtbase.dll is not loaded and nothing in the process exports _setmode or _write. Naming ucrtbase, kernel32 and ntdll outright covers all nine, and it also takes memcpy away from whatever order EnumProcessModules happens to return, since ntdll and ucrtbase both have one. Both changes are COFF only. The other two platforms already work and there is no reason to move their resolution around. --- overlay/MANIFEST | 1 + .../lib/ExecutionEngine/ExecutionEngine.cpp | 431 ++++++++++++++++++ 2 files changed, 432 insertions(+) create mode 100644 overlay/Mojo/lib/ExecutionEngine/ExecutionEngine.cpp diff --git a/overlay/MANIFEST b/overlay/MANIFEST index c72fff7..fffc8bc 100644 --- a/overlay/MANIFEST +++ b/overlay/MANIFEST @@ -32,6 +32,7 @@ edit 59f6f330c21fe83dc6241eb535e4692521d45656 Mojo/lib/Compiler/ObjectCompiler/O edit d916230e014c471d175bdaff41f1d9eed6184b55 Mojo/lib/CompilerRT/Memory.h edit d9f712923d162a14e6d7618bd9838221c4310878 Mojo/lib/CompilerRT/Python.cpp edit ae692501c5389d3afbd979de1534965ddf7db2c8 Mojo/lib/CompilerRT/System.cpp +edit cdfe0469ac49c9b42fdf334de45aaa6af44ee6bd Mojo/lib/ExecutionEngine/ExecutionEngine.cpp edit 3be1f0d71246e8f079a5695fd1c881cea6d3db1e Mojo/lib/KGENToLLVM/CABILowering.cpp edit 2d04b3c80d840066ae526e282562f8aec44c1f41 Mojo/lib/KGENToLLVM/CABILowering.h new - Mojo/lib/KGENToLLVM/CABIWin64.cpp diff --git a/overlay/Mojo/lib/ExecutionEngine/ExecutionEngine.cpp b/overlay/Mojo/lib/ExecutionEngine/ExecutionEngine.cpp new file mode 100644 index 0000000..cff1ed4 --- /dev/null +++ b/overlay/Mojo/lib/ExecutionEngine/ExecutionEngine.cpp @@ -0,0 +1,431 @@ +//===----------------------------------------------------------------------===// +// Copyright (c) 2026, Modular Inc. All rights reserved. +// +// Licensed under the Apache License v2.0 with LLVM Exceptions: +// https://llvm.org/LICENSE.txt +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +//===----------------------------------------------------------------------===// + +#include "Mojo/ExecutionEngine/ExecutionEngine.h" +#include "Mojo/ExecutionEngine/JIT/StaticArchiveLayer.h" +#include "Mojo/Support/Configuration.h" +#include "Support/ErrorOr.h" +#include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h" +#include "llvm/ExecutionEngine/Orc/COFFPlatform.h" +#include "llvm/ExecutionEngine/Orc/Core.h" +#include "llvm/ExecutionEngine/Orc/Debugging/DebugInfoSupport.h" +#include "llvm/ExecutionEngine/Orc/Debugging/DebuggerSupportPlugin.h" +#include "llvm/ExecutionEngine/Orc/Debugging/ELFDebugObjectPlugin.h" +#include "llvm/ExecutionEngine/Orc/Debugging/PerfSupportPlugin.h" +#include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h" +#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h" +#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h" +#include "llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h" +#include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderPerf.h" +#include "llvm/IR/Mangler.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/Process.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/TargetParser/Host.h" + +using namespace M; +using namespace KGEN; + +/// A standard name (that a user is unlikely to create) that we can use for a +/// JITDylib to define platform-specific symbols we want to be in the JIT'ed +/// address space. +static constexpr StringLiteral platformStdlibName = "$platform-stdlib"; +static constexpr StringLiteral compilerRTlibName = "$compilerrt-lib"; +static constexpr StringLiteral mlirclibName = "$mlirc-lib"; + +//===----------------------------------------------------------------------===// +// ExecutionEngine implementation +//===----------------------------------------------------------------------===// + +/// Set up the ORC platform for the various different binary formats/platforms +/// we support. This requires that we have an ExecutionSession *and* an +/// ObjectLinkingLayer. +/// +/// The main reason to use the platform like this is that it automatically sets +/// up the various symbols that complex code will need to execute on a target. +static ErrorOrSuccess setupPlatform(llvm::orc::JITDylib &platformStdlib, + llvm::orc::ExecutionSession &session, + llvm::orc::DylibManager &dylibMgr) { + // Add the current process symbols in. + // NOTE: COFF JIT currently doesn't support in process symbols, as it can + // currently hit conflicts with symbols in the current COFF ORC runtime. + auto generator = toModularErrorOr( + llvm::orc::EPCDynamicLibrarySearchGenerator::GetForTargetProcess( + session, dylibMgr)); + if (generator.isError()) + return generator.takeError(); + platformStdlib.addGenerator(std::move(*generator)); + + // The generator above searches the modules the process already has loaded, + // which is enough on the other two platforms and is not enough here. A + // Windows build of mojo links the C runtime statically, so ucrtbase.dll is + // not among those modules and nothing in the process exports `_setmode` or + // `_write` for JIT'd code to find. Naming the three libraries outright also + // takes the answer away from whatever order EnumProcessModules happens to + // return, which is worth something for `memcpy`, since ntdll and ucrtbase + // both have one. + // + // These are the system's own copies, found the way LoadLibrary finds them, + // and the two that are already loaded cost a reference count and nothing + // else. Between them they cover everything a Mojo program has needed so far: + // the CRT entry points from ucrtbase, the console and handle calls from + // kernel32, and `__chkstk` from either of the other two. + if (session.getTargetTriple().isOSBinFormatCOFF()) { + for (const char *library : {"ucrtbase.dll", "kernel32.dll", "ntdll.dll"}) { + auto systemGenerator = toModularErrorOr( + llvm::orc::EPCDynamicLibrarySearchGenerator::Load(session, dylibMgr, + library)); + if (systemGenerator.isError()) + return Error(Twine("error '") + + Twine(systemGenerator.getError()) + "' while loading " + + library); + platformStdlib.addGenerator(std::move(*systemGenerator)); + } + } + + return success(); +} + +/// Initialize the mlirc and CompilerRT dylib. +static ErrorOrSuccess +initializeCompilerRT(llvm::orc::ExecutionSession &session, + llvm::orc::DylibManager &dylibMgr, MojoConfig &cfg, + const llvm::DataLayout &layout, + const ExecutionEngineOptions &options) { + std::error_code ec; + + // mlirc dylib. Grab the symbols from the current process. + { + auto *libJD = &session.createBareJITDylib(mlirclibName.str()); + libJD->addGenerator(llvm::cantFail( + llvm::orc::EPCDynamicLibrarySearchGenerator::GetForTargetProcess( + session, dylibMgr, + [=](const llvm::orc::SymbolStringPtr &symbolStringPtr) { + StringRef name = *symbolStringPtr; + // On MachO, the symbol names start with `_`. + return name.starts_with("mlir") || name.starts_with("_mlir"); + }))); + } + + // CompilerRT dylib. + std::string compilerRTPath = cfg.getCompilerRTPath().str(); + if (!std::filesystem::exists(compilerRTPath, ec) || ec) + return Error(std::string("unable to locate compiler_rt ") + compilerRTPath); + + auto *libJD = &session.createBareJITDylib(compilerRTlibName.str()); + + SmallVector paths = options.libraryPaths; + paths.push_back(compilerRTPath); + + for (StringRef libPath : paths) { + auto generatorOr = + toModularErrorOr(llvm::orc::EPCDynamicLibrarySearchGenerator::Load( + session, dylibMgr, libPath.str().c_str())); + if (generatorOr.isError()) { + return Error(Twine("error '") + Twine(generatorOr.getError()) + + "' while loading compiler runtime library from '" + + libPath.str().c_str() + "'"); + } + libJD->addGenerator(std::move(*generatorOr)); + } + + // Allow pulling in sanitizer methods from the current process, as we + // currently can't activate any of these runtimes otherwise (they must + // generally be loaded first in the host process). + libJD->addGenerator(llvm::cantFail( + llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess( + layout.getGlobalPrefix(), + [=](const llvm::orc::SymbolStringPtr &symbolStringPtr) { + return llvm::any_of(ArrayRef{"__asan", "__tsan"}, + [&](StringRef prefix) { + return (*symbolStringPtr).starts_with(prefix); + }); + }))); + + return success(); +} + +M::ErrorOr> +ExecutionEngine::create(ExecutionEngineOptions options, + const llvm::TargetMachine &tm) { + // Create the data layout from the target machine. + const llvm::DataLayout &layout = tm.createDataLayout(); + const llvm::Triple &tt = tm.getTargetTriple(); + + // Construct the ExecutionSession. The user may have passed in an + // ExecutorProcessControl that we need to use. + std::unique_ptr epc = + std::move(options.epc); + + // The JIT-link memory manager is owned by the ObjectLinkingLayer so build it + // up front. + // + // JIT mapper reservation granularity. Every fresh reservation by LLVM's + // MapperJITLinkMemoryManager is rounded up to this size and mmapped + // PROT_READ|PROT_WRITE, so it sets a floor on the JIT region's virtual + // footprint — which counts against RLIMIT_AS on Linux. Large compiles + // still work: the mapper reserves additional slabs on demand. + // + // Do not increase this unless you know what you are doing. In the past, + // setting this to 1 GiB caused sporadic OoM crashes on memory-constrained + // runners. + size_t slabSize = size_t{64} * 1024 * 1024; + auto managerOr = + toModularErrorOr(llvm::orc::MapperJITLinkMemoryManager::CreateWithMapper< + llvm::orc::InProcessMemoryMapper>(slabSize)); + if (managerOr.isError()) + return managerOr.takeError(); + + if (!epc) { + auto pageSize = toModularErrorOr(llvm::sys::Process::getPageSize()); + if (pageSize.isError()) + return pageSize.takeError(); + + epc = std::make_unique( + std::make_shared(), + std::make_unique(), tt, *pageSize); + } + auto sessionPtr = + std::make_unique(std::move(epc)); + + // Create a default DylibManager from the ExecutorProcessControl. + auto dylibMgrOr = toModularErrorOr( + sessionPtr->getExecutorProcessControl().createDefaultDylibMgr()); + if (dylibMgrOr.isError()) + return dylibMgrOr.takeError(); + + // Now we can actually create the ExecutionEngine. + auto ee = std::unique_ptr( + new ExecutionEngine(std::move(sessionPtr), layout)); + ee->dylibMgr = std::move(*dylibMgrOr); + + // Open the config object so we can use it. + auto cfgOr = MojoConfig::open(); + if (cfgOr.isError()) + return cfgOr.takeError(); + MojoConfig cfg = std::move(*cfgOr); + + // Construct the object linking layer; it takes ownership of the manager. + ee->objectLayer = std::make_unique( + *ee->executionSession, std::move(*managerOr)); + + // Construct the platform stdlib - this way we don't have to worry about + // whether or not we have it later on. + llvm::orc::JITDylib &platformStdlib = + ee->executionSession->createBareJITDylib(platformStdlibName.str()); + + // If we have the platform support library, use it. This requires the + // compilation target to be a subset of the host process, so disable it for + // cross-compilation. + if (!options.crossCompiling) { + if (auto err = + setupPlatform(platformStdlib, *ee->executionSession, *ee->dylibMgr)) + return err.takeError(); + } + + if (options.registerDebugPlugins) { + llvm::orc::ExecutionSession &session = *ee->executionSession; + + // Get the registrar for the GDB JIT loader interface. + if (tt.isOSBinFormatMachO()) { + // Create and register the JIT DebugInfo plugin. The GDB alloc-action + // symbol is resolved from the session's bootstrap JITDylib, which the + // in-process executor populates automatically. + auto plugin = + toModularErrorOr(llvm::orc::GDBJITDebugInfoRegistrationPlugin::Create( + session, session.getBootstrapJITDylib())); + if (plugin.isError()) + return plugin.takeError(); + + ee->objectLayer->addPlugin(std::move(*plugin)); + } else if (tt.isOSBinFormatELF()) { + // Register the ELFDebugObjectPlugin. + llvm::Error error = llvm::Error::success(); + auto plugin = std::make_unique( + session, /*RequireDebugSections=*/true, error); + if (auto errOr = toModularErrorOr(std::move(error)); failed(errOr)) + return errOr.takeError(); + ee->objectLayer->addPlugin(std::move(plugin)); + } + } + + if (options.registerPerfPlugins) { + auto debugInfo = llvm::orc::DebugInfoPreservationPlugin::Create(); + if (!debugInfo) + return toModularError(debugInfo.takeError()); + ee->objectLayer->addPlugin(std::move(debugInfo.get())); + auto perf = std::make_unique( + ee->objectLayer->getExecutionSession().getExecutorProcessControl(), + llvm::orc::ExecutorAddr::fromPtr(&llvm_orc_registerJITLoaderPerfStart), + llvm::orc::ExecutorAddr::fromPtr(&llvm_orc_registerJITLoaderPerfEnd), + llvm::orc::ExecutorAddr::fromPtr(&llvm_orc_registerJITLoaderPerfImpl), + true, true); + ee->objectLayer->addPlugin(std::move(perf)); + } + + // Add the platform dylib to the search order. + if (auto err = ee->addToSearchOrder(platformStdlibName, &platformStdlib)) + return err.takeError(); + + // Prepare the CompilerRT dylib. + if (auto err = initializeCompilerRT(*ee->executionSession, *ee->dylibMgr, cfg, + layout, options)) + return err.takeError(); + + return std::move(ee); +} + +ErrorOr> +ExecutionEngine::createWithStandardLayers(ExecutionEngineOptions options, + const llvm::TargetMachine &tm) { + auto engineOr = ExecutionEngine::create(std::move(options), tm); + if (engineOr.isError()) + return engineOr.takeError(); + + // Add the standard layers. + (*engineOr)->addLayer((*engineOr)->getLinkingLayer()); + + return std::move(*engineOr); +} + +ExecutionEngine::ExecutionEngine( + std::unique_ptr session, + const llvm::DataLayout &dl) + : executionSession(std::move(session)), + // Parse the layout so that we own the underlying memory. DataLayout is a + // bit weird, it seems like it has some internal data structures that + // every instance shares. + dataLayout(dl.getStringRepresentation()) {} + +ExecutionEngine::~ExecutionEngine() { + if (!executionSession) + return; + + // If the execution engine has initialized the ORC runtime, the ELFNix and + // COFF platform implementations need manual shutdown. The MachOPlatform + // implementation is more sophisticated and performs shutdown automatically + // through the JITLink LinkGraph allocation actions. + const llvm::Triple &triple = executionSession->getTargetTriple(); + if (executionSession->getPlatform() && + // FIXME: On Windows, this complains about symbol not found. The Windows + // build seems happy even without the shutdown, so disable it for now. + (triple.isOSBinFormatELF() /*|| triple.isOSBinFormatCOFF()*/)) { + ErrorOr shutdown = + lookup(triple.isOSBinFormatELF() ? "__orc_rt_elfnix_platform_shutdown" + : "__orc_rt_coff_platform_shutdown"); + if (shutdown.isError()) { + llvm::report_fatal_error( + Twine("failed to find ELF/COFF platform shutdown function: ") + + shutdown.takeError().get()); + } + struct OrcRTCWrapperFunctionResult { + char *data; + size_t size; + }; + shutdown->invoke(nullptr, 0); + } + + if (auto err = executionSession->endSession()) + executionSession->reportError(std::move(err)); +} + +ErrorOr ExecutionEngine::lookup(StringRef symbol) { + return lookupWithSearchOrder(searchOrder, symbol); +} + +ErrorOr ExecutionEngine::lookup(StringRef libName, + StringRef symbol) { + llvm::orc::JITDylib *dylib = executionSession->getJITDylibByName(libName); + if (!dylib) + return Error("could not find JITDylib with name: " + libName); + + return lookupWithSearchOrder(llvm::orc::makeJITDylibSearchOrder({dylib}), + symbol); +} + +ErrorOrSuccess +ExecutionEngine::runProgram(StringRef libName, StringRef entryPoint, + function_ref runFn) { + using namespace llvm::orc; + // There is not global ctor/dtor in mojo. + + // Lookup the entry point symbol and directly invoke it rather than going + // through the runtime. + ErrorOr mainFn = lookup(entryPoint); + if (mainFn.isError()) + return mainFn.takeError(); + if (ErrorOrSuccess err = runFn(mainFn->getFunctionPointer())) + return err.takeError(); + return success(); +} + +llvm::orc::SymbolStringPtr +KGEN::ExecutionEngine::mangleAndIntern(StringRef name) { + std::string mangledName; + llvm::raw_string_ostream mangledNameStream(mangledName); + llvm::Mangler::getNameWithPrefix(mangledNameStream, name, dataLayout); + return executionSession->intern(mangledName); +} + +ErrorOrSuccess ExecutionEngine::addToSearchOrder(StringRef name, + llvm::orc::JITDylib *dylib) { + [[maybe_unused]] auto [_, didInsert] = knownDylibs.insert(name); + assert(didInsert && "must have uniquely-named dylibs"); + + // If this isn't the platform stdlib, setup CompilerRT and mlirc. + if (name != platformStdlibName) { + dylib->addToLinkOrder( + *executionSession->getJITDylibByName(compilerRTlibName)); + dylib->addToLinkOrder(*executionSession->getJITDylibByName(mlirclibName)); + + // The platform stdlib is in the search order a lookup starts from, but + // until now it was in nobody's link order, which is what a symbol gets + // resolved through once materialization is under way. On ELF and Mach-O + // that gap never showed, because dlsym on the compiler runtime handle keeps + // walking into that library's own dependencies and the C runtime is one of + // them. GetProcAddress does no such walking: it reads one module's export + // table and stops. So on Windows the only thing that can answer for a + // libc symbol is the platform stdlib, and it has to be reachable from here. + // + // Last, so that anything the Mojo runtime defines still wins over the + // system copy of the same name. + if (executionSession->getTargetTriple().isOSBinFormatCOFF()) + dylib->addToLinkOrder( + *executionSession->getJITDylibByName(platformStdlibName)); + } + + // Use higher preference for newer dylibs. + searchOrder.insert(searchOrder.begin(), + {dylib, llvm::orc::JITDylibLookupFlags::MatchAllSymbols}); + return success(); +} + +ErrorOr ExecutionEngine::lookupWithSearchOrder( + const llvm::orc::JITDylibSearchOrder &order, llvm::StringRef symbol) { + // Look up this symbol with the search order provided. + llvm::Expected sym = + executionSession->lookup(order, mangleAndIntern(symbol)); + if (sym) + return CompiledFunc(sym->getAddress().toPtr()); + + // Check to see if any of the layers have errors. + auto found = llvm::find_if( + layers, [](const auto &layer) { return layer->hasError(); }); + // If not, return the error returned by the ORC. + if (found == layers.end()) + return toModularError(sym.takeError()); + + // Add the additional context from the layer's error. + return Error(llvm::toString(sym.takeError()) + + " (from the layer: " + (*found)->takeError().get() + ")"); +}