[Mlir-commits] [mlir] 0b3841e - [mlir] Move symbol loading from mlir-cpu-runner to ExecutionEngine.
Ingo Müller
llvmlistbot at llvm.org
Fri Jun 16 07:50:17 PDT 2023
Author: Ingo Müller
Date: 2023-06-16T14:50:14Z
New Revision: 0b3841eb97f5c7951c2375e5b1a38ea89d9d6a37
URL: https://github.com/llvm/llvm-project/commit/0b3841eb97f5c7951c2375e5b1a38ea89d9d6a37
DIFF: https://github.com/llvm/llvm-project/commit/0b3841eb97f5c7951c2375e5b1a38ea89d9d6a37.diff
LOG: [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
Added:
Modified:
mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
mlir/lib/ExecutionEngine/AsyncRuntime.cpp
mlir/lib/ExecutionEngine/ExecutionEngine.cpp
mlir/lib/ExecutionEngine/JitRunner.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
index c62a18d091b73..573a94d34b1c6 100644
--- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -71,7 +71,15 @@ struct ExecutionEngineOptions {
std::optional<llvm::CodeGenOpt::Level> 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<StringRef> 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<void *> &);
+
+ /// 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 @@ class ExecutionEngine {
/// Perf notification listener.
llvm::JITEventListener *perfListener;
+
+ /// Destroy functions in the libraries loaded by the ExecutionEngine that are
+ /// called when this ExecutionEngine is destructed.
+ SmallVector<LibraryDestroyFn> destroyFns;
};
} // namespace mlir
diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp
index 00cd27b25be80..d7c09f9826935 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<void *> &exportSymbols);
+extern "C" API void
+__mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols);
// NOLINTNEXTLINE(*-identifier-naming): externally called.
-void __mlir_runner_init(llvm::StringMap<void *> &exportSymbols) {
+void __mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols) {
auto exportSymbol = [&](llvm::StringRef name, auto ptr) {
assert(exportSymbols.count(name) == 0 && "symbol already exists");
exportSymbols[name] = reinterpret_cast<void *>(ptr);
@@ -526,7 +527,9 @@ void __mlir_runner_init(llvm::StringMap<void *> &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 655d093e17643..da80dbc9f6a03 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<std::unique_ptr<ExecutionEngine>>
ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
std::unique_ptr<llvm::TargetMachine> 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<SmallString<256>, 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<void *> exportSymbols;
+ SmallVector<LibraryDestroyFn> 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<LibraryInitFn>(initSym);
+ initFn(exportSymbols);
+
+ auto destroyFn = reinterpret_cast<LibraryDestroyFn>(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 b4279af7642f1..7a02ecb8376cf 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<llvm::CodeGenOpt::Level>(*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<SmallString<256>, 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<StringRef, 4> executionEngineLibs;
-
- using MlirRunnerInitFn = void (*)(llvm::StringMap<void *> &);
- using MlirRunnerDestroyFn = void (*)();
-
- llvm::StringMap<void *> exportSymbols;
- SmallVector<MlirRunnerDestroyFn> 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<MlirRunnerInitFn>(initSym);
- initFn(exportSymbols);
-
- auto destroyFn = reinterpret_cast<MlirRunnerDestroyFn>(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<StringRef, 4> 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();
}
More information about the Mlir-commits
mailing list