Skip to content

Commit

Permalink
[mlir][ExecutionEngine] Only load JITDyLibs without init/destroy funcs.
Browse files Browse the repository at this point in the history
In https://reviews.llvm.org/D153029, I moved the loading/unloading
mechanisms of shared libraries from the JIT runner to the execution
engine in order to make that mechanism available in the latter
(including its Python bindings). However, I realized that I introduced a
small change in semantic: previously, the JIT runner checked for the
presence of init/destroy functions and only loaded the library as
JITDyLib if they were not present. After I moved the code, all libraries
were loaded as JITDyLib, even if they registered their symbols
explicitly in their init function. I am not sure if this is really a
problem but (1) the previous behavior was different and (2) I guess it
could cause a problem if some symbols are exported through the init
function *and*  have public visibility. This patch reestablishes the
original behaviour in the new place of the code.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D153249
  • Loading branch information
ingomueller-net committed Jun 19, 2023
1 parent 53f6229 commit 0eb0fec
Showing 1 changed file with 29 additions and 27 deletions.
56 changes: 29 additions & 27 deletions mlir/lib/ExecutionEngine/ExecutionEngine.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,33 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
return absPath;
});

// If shared library implements custom execution layer library init and
// destroy functions, we'll use them to register the library. Otherwise, load
// the library as JITDyLib below.
llvm::StringMap<void *> exportSymbols;
SmallVector<LibraryDestroyFn> destroyFns;
SmallVector<StringRef> jitDyLibPaths;

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) {
jitDyLibPaths.push_back(libPath);
continue;
}

auto initFn = reinterpret_cast<LibraryInitFn>(initSym);
initFn(exportSymbols);

auto destroyFn = reinterpret_cast<LibraryDestroyFn>(destroySim);
destroyFns.push_back(destroyFn);
}
engine->destroyFns = std::move(destroyFns);

// Callback to create the object layer with symbol resolution to current
// process and dynamically linked libraries.
auto objectLinkingLayerCreator = [&](ExecutionSession &session,
Expand All @@ -308,7 +335,7 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
}

// Resolve symbols from shared libraries.
for (auto &libPath : sharedLibPaths) {
for (auto &libPath : jitDyLibPaths) {
auto mb = llvm::MemoryBuffer::getFile(libPath);
if (!mb) {
errs() << "Failed to create MemoryBuffer for: " << libPath
Expand All @@ -317,7 +344,7 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
}
auto &jd = session.createBareJITDylib(std::string(libPath));
auto loaded = DynamicLibrarySearchGenerator::Load(
libPath.str().str().c_str(), dataLayout.getGlobalPrefix());
libPath.str().c_str(), dataLayout.getGlobalPrefix());
if (!loaded) {
errs() << "Could not load " << libPath << ":\n " << loaded.takeError()
<< "\n";
Expand Down Expand Up @@ -362,31 +389,6 @@ 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();
Expand Down

0 comments on commit 0eb0fec

Please sign in to comment.