[Mlir-commits] [mlir] 0eb0fec - [mlir][ExecutionEngine] Only load JITDyLibs without init/destroy funcs.

[Ingo Müller]

Author: Ingo Müller
Date: 2023-06-19T07:38:51Z
New Revision: 0eb0fecbc544a0416b70f4450a38dfb27c9b5b9a

URL: https://github.com/llvm/llvm-project/commit/0eb0fecbc544a0416b70f4450a38dfb27c9b5b9a
DIFF: https://github.com/llvm/llvm-project/commit/0eb0fecbc544a0416b70f4450a38dfb27c9b5b9a.diff

LOG: [mlir][ExecutionEngine] Only load JITDyLibs without init/destroy funcs.

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

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index da80dbc9f6a03..ef3f57d3a069a 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -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,
@@ -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
@@ -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";
@@ -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();