From 0b3841eb97f5c7951c2375e5b1a38ea89d9d6a37 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Ingo=20M=C3=BCller?= Date: Thu, 15 Jun 2023 14:33:22 +0000 Subject: [PATCH] [mlir] Move symbol loading from mlir-cpu-runner to ExecutionEngine. Both the mlir-cpu-runner and the execution engine allow to provide a list of shared libraries that should be loaded into the process such that the jitted code can use the symbols from those libraries. The runner had implemented a protocol that allowed libraries to control which symbols it wants to provide in that context (with a function called __mlir_runner_init). In absence of that, the runner would rely on the loading mechanism of the execution engine, which didn't do anything particular with the symbols, i.e., only symbols with public visibility were visible to jitted code. Libraries used a mix of the two mechanisms: while the runner utils and C runner utils libs (and potentially others) used public visibility, the async runtime lib (as the only one in the monorepo) used the loading protocol. As a consequence, the async runtime library could not be used through the Python bindings of the execution engine. This patch moves the loading protocol from the runner to the execution engine. For the runner, this should not change anything: it lets the execution engine handle the loading which now implements the same protocol that the runner had implemented before. However, the Python binding now get to benefit from the loading protocol as well, so the async runtime library (and potentially other out-of-tree libraries) can now be used in that context. Reviewed By: mehdi_amini Differential Revision: https://reviews.llvm.org/D153029 --- .../include/mlir/ExecutionEngine/ExecutionEngine.h | 37 ++++++++++++- mlir/lib/ExecutionEngine/AsyncRuntime.cpp | 9 ++-- mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 56 +++++++++++++++++++- mlir/lib/ExecutionEngine/JitRunner.cpp | 61 ++-------------------- 4 files changed, 99 insertions(+), 64 deletions(-) diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h index c62a18d..573a94d 100644 --- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h +++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h @@ -71,7 +71,15 @@ struct ExecutionEngineOptions { std::optional jitCodeGenOptLevel; /// If `sharedLibPaths` are provided, the underlying JIT-compilation will - /// open and link the shared libraries for symbol resolution. + /// open and link the shared libraries for symbol resolution. Libraries that + /// are designed to be used with the `ExecutionEngine` may implement a + /// loading and unloading protocol: if they implement the two functions with + /// the names defined in `kLibraryInitFnName` and `kLibraryDestroyFnName`, + /// these functions will be called upon loading the library and upon + /// destruction of the `ExecutionEngine`. In the init function, the library + /// may provide a list of symbols that it wants to make available to code + /// run by the `ExecutionEngine`. If the two functions are not defined, only + /// symbols with public visibility are available to the executed code. ArrayRef sharedLibPaths = {}; /// Specifies an existing `sectionMemoryMapper` to be associated with the @@ -105,9 +113,32 @@ struct ExecutionEngineOptions { /// be used to invoke the JIT-compiled function. class ExecutionEngine { public: + /// Name of init functions of shared libraries. If a library provides a + /// function with this name and the one of the destroy function, this function + /// is called upon loading the library. + static constexpr const char *const kLibraryInitFnName = + "__mlir_execution_engine_init"; + + /// Name of destroy functions of shared libraries. If a library provides a + /// function with this name and the one of the init function, this function is + /// called upon destructing the `ExecutionEngine`. + static constexpr const char *const kLibraryDestroyFnName = + "__mlir_execution_engine_destroy"; + + /// Function type for init functions of shared libraries. The library may + /// provide a list of symbols that it wants to make available to code run by + /// the `ExecutionEngine`. If the two functions are not defined, only symbols + /// with public visibility are available to the executed code. + using LibraryInitFn = void (*)(llvm::StringMap &); + + /// Function type for destroy functions of shared libraries. + using LibraryDestroyFn = void (*)(); + ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener, bool enablePerfNotificationListener); + ~ExecutionEngine(); + /// Creates an execution engine for the given MLIR IR. If TargetMachine is /// not provided, default TM is created (i.e. ignoring any command line flags /// that could affect the set-up). @@ -216,6 +247,10 @@ private: /// Perf notification listener. llvm::JITEventListener *perfListener; + + /// Destroy functions in the libraries loaded by the ExecutionEngine that are + /// called when this ExecutionEngine is destructed. + SmallVector destroyFns; }; } // namespace mlir diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp index 00cd27b..d7c09f9 100644 --- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp +++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp @@ -468,10 +468,11 @@ extern "C" void mlirAsyncRuntimePrintCurrentThreadId() { // The bug is fixed in VS2019 16.1. Separating the declaration and definition is // a work around for older versions of Visual Studio. // NOLINTNEXTLINE(*-identifier-naming): externally called. -extern "C" API void __mlir_runner_init(llvm::StringMap &exportSymbols); +extern "C" API void +__mlir_execution_engine_init(llvm::StringMap &exportSymbols); // NOLINTNEXTLINE(*-identifier-naming): externally called. -void __mlir_runner_init(llvm::StringMap &exportSymbols) { +void __mlir_execution_engine_init(llvm::StringMap &exportSymbols) { auto exportSymbol = [&](llvm::StringRef name, auto ptr) { assert(exportSymbols.count(name) == 0 && "symbol already exists"); exportSymbols[name] = reinterpret_cast(ptr); @@ -526,7 +527,9 @@ void __mlir_runner_init(llvm::StringMap &exportSymbols) { } // NOLINTNEXTLINE(*-identifier-naming): externally called. -extern "C" API void __mlir_runner_destroy() { resetDefaultAsyncRuntime(); } +extern "C" API void __mlir_execution_engine_destroy() { + resetDefaultAsyncRuntime(); +} } // namespace runtime } // namespace mlir diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp index 655d093..da80dbc 100644 --- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp +++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp @@ -222,6 +222,12 @@ ExecutionEngine::ExecutionEngine(bool enableObjectDump, } } +ExecutionEngine::~ExecutionEngine() { + // Run all dynamic library destroy callbacks to prepare for the shutdown. + for (LibraryDestroyFn destroy : destroyFns) + destroy(); +} + Expected> ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, std::unique_ptr tm) { @@ -267,6 +273,16 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, auto dataLayout = llvmModule->getDataLayout(); + // Use absolute library path so that gdb can find the symbol table. + SmallVector, 4> sharedLibPaths; + transform( + options.sharedLibPaths, std::back_inserter(sharedLibPaths), + [](StringRef libPath) { + SmallString<256> absPath(libPath.begin(), libPath.end()); + cantFail(llvm::errorCodeToError(llvm::sys::fs::make_absolute(absPath))); + return absPath; + }); + // Callback to create the object layer with symbol resolution to current // process and dynamically linked libraries. auto objectLinkingLayerCreator = [&](ExecutionSession &session, @@ -292,7 +308,7 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, } // Resolve symbols from shared libraries. - for (auto libPath : options.sharedLibPaths) { + for (auto &libPath : sharedLibPaths) { auto mb = llvm::MemoryBuffer::getFile(libPath); if (!mb) { errs() << "Failed to create MemoryBuffer for: " << libPath @@ -301,7 +317,7 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, } auto &jd = session.createBareJITDylib(std::string(libPath)); auto loaded = DynamicLibrarySearchGenerator::Load( - libPath.data(), dataLayout.getGlobalPrefix()); + libPath.str().str().c_str(), dataLayout.getGlobalPrefix()); if (!loaded) { errs() << "Could not load " << libPath << ":\n " << loaded.takeError() << "\n"; @@ -346,6 +362,42 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options, cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess( dataLayout.getGlobalPrefix()))); + // If shared library implements custom execution layer library init and + // destroy functions, we'll use them to register the library. + + llvm::StringMap exportSymbols; + SmallVector destroyFns; + + for (auto &libPath : sharedLibPaths) { + auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary( + libPath.str().str().c_str()); + void *initSym = lib.getAddressOfSymbol(kLibraryInitFnName); + void *destroySim = lib.getAddressOfSymbol(kLibraryDestroyFnName); + + // Library does not provide call backs, rely on symbol visiblity. + if (!initSym || !destroySim) { + continue; + } + + auto initFn = reinterpret_cast(initSym); + initFn(exportSymbols); + + auto destroyFn = reinterpret_cast(destroySim); + destroyFns.push_back(destroyFn); + } + engine->destroyFns = std::move(destroyFns); + + // Build a runtime symbol map from the exported symbols and register them. + auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) { + auto symbolMap = llvm::orc::SymbolMap(); + for (auto &exportSymbol : exportSymbols) + symbolMap[interner(exportSymbol.getKey())] = { + llvm::orc::ExecutorAddr::fromPtr(exportSymbol.getValue()), + llvm::JITSymbolFlags::Exported}; + return symbolMap; + }; + engine->registerSymbols(runtimeSymbolMap); + return std::move(engine); } diff --git a/mlir/lib/ExecutionEngine/JitRunner.cpp b/mlir/lib/ExecutionEngine/JitRunner.cpp index b4279af..7a02ecb 100644 --- a/mlir/lib/ExecutionEngine/JitRunner.cpp +++ b/mlir/lib/ExecutionEngine/JitRunner.cpp @@ -185,65 +185,15 @@ compileAndExecute(Options &options, Operation *module, StringRef entryPoint, if (auto clOptLevel = getCommandLineOptLevel(options)) jitCodeGenOptLevel = static_cast(*clOptLevel); - // If shared library implements custom mlir-runner library init and destroy - // functions, we'll use them to register the library with the execution - // engine. Otherwise we'll pass library directly to the execution engine. - SmallVector, 4> libPaths; - - // Use absolute library path so that gdb can find the symbol table. - transform( - options.clSharedLibs, std::back_inserter(libPaths), - [](std::string libPath) { - SmallString<256> absPath(libPath.begin(), libPath.end()); - cantFail(llvm::errorCodeToError(llvm::sys::fs::make_absolute(absPath))); - return absPath; - }); - - // Libraries that we'll pass to the ExecutionEngine for loading. - SmallVector executionEngineLibs; - - using MlirRunnerInitFn = void (*)(llvm::StringMap &); - using MlirRunnerDestroyFn = void (*)(); - - llvm::StringMap exportSymbols; - SmallVector destroyFns; - - // Handle libraries that do support mlir-runner init/destroy callbacks. - for (auto &libPath : libPaths) { - auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(libPath.c_str()); - void *initSym = lib.getAddressOfSymbol("__mlir_runner_init"); - void *destroySim = lib.getAddressOfSymbol("__mlir_runner_destroy"); - - // Library does not support mlir runner, load it with ExecutionEngine. - if (!initSym || !destroySim) { - executionEngineLibs.push_back(libPath); - continue; - } - - auto initFn = reinterpret_cast(initSym); - initFn(exportSymbols); - - auto destroyFn = reinterpret_cast(destroySim); - destroyFns.push_back(destroyFn); - } - - // Build a runtime symbol map from the config and exported symbols. - auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) { - auto symbolMap = config.runtimeSymbolMap ? config.runtimeSymbolMap(interner) - : llvm::orc::SymbolMap(); - for (auto &exportSymbol : exportSymbols) - symbolMap[interner(exportSymbol.getKey())] = - { llvm::orc::ExecutorAddr::fromPtr(exportSymbol.getValue()), - llvm::JITSymbolFlags::Exported }; - return symbolMap; - }; + SmallVector sharedLibs(options.clSharedLibs.begin(), + options.clSharedLibs.end()); mlir::ExecutionEngineOptions engineOptions; engineOptions.llvmModuleBuilder = config.llvmModuleBuilder; if (config.transformer) engineOptions.transformer = config.transformer; engineOptions.jitCodeGenOptLevel = jitCodeGenOptLevel; - engineOptions.sharedLibPaths = executionEngineLibs; + engineOptions.sharedLibPaths = sharedLibs; engineOptions.enableObjectDump = true; auto expectedEngine = mlir::ExecutionEngine::create(module, engineOptions, std::move(tm)); @@ -251,7 +201,6 @@ compileAndExecute(Options &options, Operation *module, StringRef entryPoint, return expectedEngine.takeError(); auto engine = std::move(*expectedEngine); - engine->registerSymbols(runtimeSymbolMap); auto expectedFPtr = engine->lookupPacked(entryPoint); if (!expectedFPtr) @@ -265,10 +214,6 @@ compileAndExecute(Options &options, Operation *module, StringRef entryPoint, void (*fptr)(void **) = *expectedFPtr; (*fptr)(args); - // Run all dynamic library destroy callbacks to prepare for the shutdown. - for (MlirRunnerDestroyFn destroy : destroyFns) - destroy(); - return Error::success(); } -- 2.7.4