Hack to preserve modules following compilation

gmarkall · Mar 15, 2023 · 0c6b605 · 0c6b605
1 parent 1116373
commit 0c6b605
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/ffi/orcjit.cpp b/ffi/orcjit.cpp
@@ -34,6 +34,8 @@ inline OwningBinary<ObjectFile> *unwrap(LLVMObjectFileRef OF) {
 
 extern "C" {
 
+std::vector<ThreadSafeModule> GlobalTSMs;
+
 API_EXPORT(LLVMOrcLLJITRef)
 LLVMPY_CreateLLJITCompiler(LLVMTargetMachineRef tm, const char **OutError) {
     LLVMOrcLLJITRef jit;
@@ -62,6 +64,11 @@ LLVMPY_CreateLLJITCompiler(LLVMTargetMachineRef tm, const char **OutError) {
 
     auto error = LLVMOrcCreateLLJIT(&jit, builder);
 
+    unwrap(jit)->getIRCompileLayer().setNotifyCompiled([](MaterializationResponsibility &R, ThreadSafeModule TSM) {
+        // Move TSM to take ownership and preserve the module.
+        GlobalTSMs.push_back(std::move(TSM));
+    });
+
     if (error) {
         char *message = LLVMGetErrorMessage(error);
         *OutError = LLVMPY_CreateString(message);