jit-rt: batched symbol lookup

ldc-developers · Nov 17, 2024 · 045bf6a · 045bf6a
1 parent e41d3b4
commit 045bf6a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/runtime/jit-rt/cpp-so/compile.cpp b/runtime/jit-rt/cpp-so/compile.cpp
@@ -252,15 +252,6 @@ void generateBind(const Context &context, DynamicCompilerContext &jitContext,
 void applyBind(const Context &context, DynamicCompilerContext &jitContext,
                const JitModuleInfo &moduleInfo) {
   auto &layout = jitContext.getDataLayout();
-  std::vector<std::string> names{};
-  for (auto &elem : moduleInfo.getBindHandles()) {
-    names.emplace_back(elem.name);
-  }
-  auto results = cantFail(jitContext.lookupMany(names));
-  for (auto &symbol : results) {
-    jitContext.addSymbol((*symbol.first).str(),
-                         symbol.second.getAddress().toPtr<void *>());
-  }
   for (auto &elem : moduleInfo.getBindHandles()) {
     auto decorated = decorate(elem.name, layout);
     auto symbol = jitContext.lookup(elem.name);
@@ -453,12 +444,24 @@ void rtCompileProcessImplSoInternal(const RtCompileModuleList *modlist_head,
   JitFinaliser jitFinalizer(myJit);
   /*if (myJit.isMainContext())*/ {
     interruptPoint(context, "Resolve functions");
+    std::vector<std::string> names{};
+    for (auto &elem : moduleInfo.functions()) {
+      names.emplace_back(elem.name);
+    }
+    for (auto &elem : moduleInfo.getBindHandles()) {
+      names.emplace_back(elem.name);
+    }
+    auto results = cantFail(myJit.lookupMany(names));
+    for (auto &symbol : results) {
+      myJit.addSymbol((*symbol.first).str(),
+                      symbol.second.getAddress().toPtr<void *>());
+    }
     for (auto &&fun : moduleInfo.functions()) {
       if (fun.thunkVar == nullptr) {
         continue;
       }
       auto decorated = decorate(fun.name, layout);
-      auto symbol = myJit.lookup(fun.name.data());
+      auto symbol = myJit.lookup(fun.name.str());
       auto addr = resolveSymbol(symbol);
       if (nullptr == addr) {
         std::string desc = std::string("Symbol not found in jitted code: \"") +

diff --git a/runtime/jit-rt/cpp-so/jit_context.cpp b/runtime/jit-rt/cpp-so/jit_context.cpp
@@ -186,6 +186,10 @@ DynamicCompilerContext::findSymbol(const std::string &name) {
 llvm::Expected<llvm::orc::ExecutorSymbolDef>
 DynamicCompilerContext::lookup(const std::string &name) {
   auto mangled = mangler(name);
+  auto symbol = findSymbol((*mangled).str());
+  if (symbol) {
+    return *symbol;
+  }
   return execSession->lookup({&moduleHandle}, mangled);
 }
 
@@ -204,7 +208,7 @@ DynamicCompilerContext::lookupMany(const std::vector<std::string> &names) {
 void DynamicCompilerContext::clearSymMap() { symMap.clear(); }
 
 void DynamicCompilerContext::addSymbol(std::string &&name, void *value) {
-  symMap.emplace(std::make_pair(*mangler(name), value));
+  symMap[(*mangler(name)).str()] = value;
 }
 
 void DynamicCompilerContext::reset() {