| 94 | } |
| 95 | |
| 96 | CompiledModuleSP ThrustJITv2::CompileModule( |
| 97 | std::function<bool(llvm::Module&)> irGenerator, |
| 98 | const std::string& funcName) { |
| 99 | { |
| 100 | std::unique_lock lock(mutex_); |
| 101 | if (auto cached = compiledModuleCache_.get(funcName); cached != nullptr) { |
| 102 | return cached; |
| 103 | } |
| 104 | |
| 105 | if (compilingFunctions_.count(funcName) > 0) { |
| 106 | compilingCv_.wait( |
| 107 | lock, [&]() { return compilingFunctions_.count(funcName) == 0; }); |
| 108 | return compiledModuleCache_.get(funcName); |
| 109 | } |
| 110 | |
| 111 | compilingFunctions_.insert(funcName); |
| 112 | } |
| 113 | |
| 114 | auto clearCompilingFlag = [this, &funcName]() { |
| 115 | std::lock_guard lock(mutex_); |
| 116 | compilingFunctions_.erase(funcName); |
| 117 | compilingCv_.notify_all(); |
| 118 | }; |
| 119 | |
| 120 | auto llvmContext = std::make_unique<llvm::LLVMContext>(); |
| 121 | auto llvmModule = std::make_unique<llvm::Module>(funcName, *llvmContext); |
| 122 | llvmModule->setDataLayout(jit_->getDataLayout()); |
| 123 | if (irGenerator(*llvmModule)) { |
| 124 | clearCompilingFlag(); |
| 125 | return nullptr; |
| 126 | } |
| 127 | |
| 128 | std::vector<std::string> funcNames; |
| 129 | for (auto& function : *llvmModule) { |
| 130 | if (!function.isDeclaration()) { |
| 131 | funcNames.emplace_back(function.getName().str()); |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | auto resourceTracker = jit_->getMainJITDylib().createResourceTracker(); |
| 136 | llvm::orc::ResourceKey resourceKey = 0; |
| 137 | if (auto keyErr = resourceTracker->withResourceKeyDo( |
| 138 | [&](llvm::orc::ResourceKey key) { resourceKey = key; })) { |
| 139 | llvm::handleAllErrors(std::move(keyErr), [&](llvm::ErrorInfoBase& eib) { |
| 140 | llvm::errs() << "[JITv2] ResourceTracker key Error: " << eib.message() |
| 141 | << '\n'; |
| 142 | }); |
| 143 | clearCompilingFlag(); |
| 144 | return nullptr; |
| 145 | } |
| 146 | |
| 147 | auto err = jit_->addIRModule( |
| 148 | resourceTracker, |
| 149 | llvm::orc::ThreadSafeModule( |
| 150 | std::move(llvmModule), std::move(llvmContext))); |
| 151 | if (err) { |
| 152 | takeTrackedObjectSize(resourceKey); |
| 153 | llvm::handleAllErrors(std::move(err), [&](llvm::ErrorInfoBase& eib) { |