1 //===- ExecutionEngine.cpp - MLIR Execution engine and utils --------------===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements the execution engine for MLIR modules based on LLVM Orc 10 // JIT engine. 11 // 12 //===----------------------------------------------------------------------===// 13 #include "mlir/ExecutionEngine/ExecutionEngine.h" 14 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 15 #include "mlir/IR/Function.h" 16 #include "mlir/IR/Module.h" 17 #include "mlir/Support/FileUtilities.h" 18 #include "mlir/Target/LLVMIR.h" 19 20 #include "llvm/ExecutionEngine/JITEventListener.h" 21 #include "llvm/ExecutionEngine/ObjectCache.h" 22 #include "llvm/ExecutionEngine/Orc/CompileUtils.h" 23 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" 24 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" 25 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" 26 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" 27 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" 28 #include "llvm/ExecutionEngine/SectionMemoryManager.h" 29 #include "llvm/IR/IRBuilder.h" 30 #include "llvm/Support/Debug.h" 31 #include "llvm/Support/Error.h" 32 #include "llvm/Support/Host.h" 33 #include "llvm/Support/TargetRegistry.h" 34 #include "llvm/Support/ToolOutputFile.h" 35 36 #define DEBUG_TYPE "execution-engine" 37 38 using namespace mlir; 39 using llvm::dbgs; 40 using llvm::Error; 41 using llvm::errs; 42 using llvm::Expected; 43 using llvm::LLVMContext; 44 using llvm::MemoryBuffer; 45 using llvm::MemoryBufferRef; 46 using llvm::Module; 47 using llvm::SectionMemoryManager; 48 using llvm::StringError; 49 using llvm::Triple; 50 using llvm::orc::DynamicLibrarySearchGenerator; 51 using llvm::orc::ExecutionSession; 52 using llvm::orc::IRCompileLayer; 53 using llvm::orc::JITTargetMachineBuilder; 54 using llvm::orc::RTDyldObjectLinkingLayer; 55 using llvm::orc::ThreadSafeModule; 56 using llvm::orc::TMOwningSimpleCompiler; 57 58 /// Wrap a string into an llvm::StringError. 59 static Error make_string_error(const Twine &message) { 60 return llvm::make_error<StringError>(message.str(), 61 llvm::inconvertibleErrorCode()); 62 } 63 64 void SimpleObjectCache::notifyObjectCompiled(const Module *M, 65 MemoryBufferRef ObjBuffer) { 66 cachedObjects[M->getModuleIdentifier()] = MemoryBuffer::getMemBufferCopy( 67 ObjBuffer.getBuffer(), ObjBuffer.getBufferIdentifier()); 68 } 69 70 std::unique_ptr<MemoryBuffer> SimpleObjectCache::getObject(const Module *M) { 71 auto I = cachedObjects.find(M->getModuleIdentifier()); 72 if (I == cachedObjects.end()) { 73 LLVM_DEBUG(dbgs() << "No object for " << M->getModuleIdentifier() 74 << " in cache. Compiling.\n"); 75 return nullptr; 76 } 77 LLVM_DEBUG(dbgs() << "Object for " << M->getModuleIdentifier() 78 << " loaded from cache.\n"); 79 return MemoryBuffer::getMemBuffer(I->second->getMemBufferRef()); 80 } 81 82 void SimpleObjectCache::dumpToObjectFile(StringRef outputFilename) { 83 // Set up the output file. 84 std::string errorMessage; 85 auto file = openOutputFile(outputFilename, &errorMessage); 86 if (!file) { 87 llvm::errs() << errorMessage << "\n"; 88 return; 89 } 90 91 // Dump the object generated for a single module to the output file. 92 assert(cachedObjects.size() == 1 && "Expected only one object entry."); 93 auto &cachedObject = cachedObjects.begin()->second; 94 file->os() << cachedObject->getBuffer(); 95 file->keep(); 96 } 97 98 void ExecutionEngine::dumpToObjectFile(StringRef filename) { 99 cache->dumpToObjectFile(filename); 100 } 101 102 // Setup LLVM target triple from the current machine. 103 bool ExecutionEngine::setupTargetTriple(Module *llvmModule) { 104 // Setup the machine properties from the current architecture. 105 auto targetTriple = llvm::sys::getDefaultTargetTriple(); 106 std::string errorMessage; 107 auto target = llvm::TargetRegistry::lookupTarget(targetTriple, errorMessage); 108 if (!target) { 109 errs() << "NO target: " << errorMessage << "\n"; 110 return true; 111 } 112 std::unique_ptr<llvm::TargetMachine> machine( 113 target->createTargetMachine(targetTriple, "generic", "", {}, {})); 114 llvmModule->setDataLayout(machine->createDataLayout()); 115 llvmModule->setTargetTriple(targetTriple); 116 return false; 117 } 118 119 static std::string makePackedFunctionName(StringRef name) { 120 return "_mlir_" + name.str(); 121 } 122 123 // For each function in the LLVM module, define an interface function that wraps 124 // all the arguments of the original function and all its results into an i8** 125 // pointer to provide a unified invocation interface. 126 static void packFunctionArguments(Module *module) { 127 auto &ctx = module->getContext(); 128 llvm::IRBuilder<> builder(ctx); 129 DenseSet<llvm::Function *> interfaceFunctions; 130 for (auto &func : module->getFunctionList()) { 131 if (func.isDeclaration()) { 132 continue; 133 } 134 if (interfaceFunctions.count(&func)) { 135 continue; 136 } 137 138 // Given a function `foo(<...>)`, define the interface function 139 // `mlir_foo(i8**)`. 140 auto newType = llvm::FunctionType::get( 141 builder.getVoidTy(), builder.getInt8PtrTy()->getPointerTo(), 142 /*isVarArg=*/false); 143 auto newName = makePackedFunctionName(func.getName()); 144 auto funcCst = module->getOrInsertFunction(newName, newType); 145 llvm::Function *interfaceFunc = cast<llvm::Function>(funcCst.getCallee()); 146 interfaceFunctions.insert(interfaceFunc); 147 148 // Extract the arguments from the type-erased argument list and cast them to 149 // the proper types. 150 auto bb = llvm::BasicBlock::Create(ctx); 151 bb->insertInto(interfaceFunc); 152 builder.SetInsertPoint(bb); 153 llvm::Value *argList = interfaceFunc->arg_begin(); 154 SmallVector<llvm::Value *, 8> args; 155 args.reserve(llvm::size(func.args())); 156 for (auto &indexedArg : llvm::enumerate(func.args())) { 157 llvm::Value *argIndex = llvm::Constant::getIntegerValue( 158 builder.getInt64Ty(), APInt(64, indexedArg.index())); 159 llvm::Value *argPtrPtr = builder.CreateGEP(argList, argIndex); 160 llvm::Value *argPtr = builder.CreateLoad(argPtrPtr); 161 argPtr = builder.CreateBitCast( 162 argPtr, indexedArg.value().getType()->getPointerTo()); 163 llvm::Value *arg = builder.CreateLoad(argPtr); 164 args.push_back(arg); 165 } 166 167 // Call the implementation function with the extracted arguments. 168 llvm::Value *result = builder.CreateCall(&func, args); 169 170 // Assuming the result is one value, potentially of type `void`. 171 if (!result->getType()->isVoidTy()) { 172 llvm::Value *retIndex = llvm::Constant::getIntegerValue( 173 builder.getInt64Ty(), APInt(64, llvm::size(func.args()))); 174 llvm::Value *retPtrPtr = builder.CreateGEP(argList, retIndex); 175 llvm::Value *retPtr = builder.CreateLoad(retPtrPtr); 176 retPtr = builder.CreateBitCast(retPtr, result->getType()->getPointerTo()); 177 builder.CreateStore(result, retPtr); 178 } 179 180 // The interface function returns void. 181 builder.CreateRetVoid(); 182 } 183 } 184 185 ExecutionEngine::ExecutionEngine(bool enableObjectCache, 186 bool enableGDBNotificationListener, 187 bool enablePerfNotificationListener) 188 : cache(enableObjectCache ? new SimpleObjectCache() : nullptr), 189 gdbListener(enableGDBNotificationListener 190 ? llvm::JITEventListener::createGDBRegistrationListener() 191 : nullptr), 192 perfListener(enablePerfNotificationListener 193 ? llvm::JITEventListener::createPerfJITEventListener() 194 : nullptr) {} 195 196 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create( 197 ModuleOp m, std::function<Error(llvm::Module *)> transformer, 198 Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel, 199 ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache, 200 bool enableGDBNotificationListener, bool enablePerfNotificationListener) { 201 auto engine = std::make_unique<ExecutionEngine>( 202 enableObjectCache, enableGDBNotificationListener, 203 enablePerfNotificationListener); 204 205 std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext); 206 auto llvmModule = translateModuleToLLVMIR(m); 207 if (!llvmModule) 208 return make_string_error("could not convert to LLVM IR"); 209 // FIXME: the triple should be passed to the translation or dialect conversion 210 // instead of this. Currently, the LLVM module created above has no triple 211 // associated with it. 212 setupTargetTriple(llvmModule.get()); 213 packFunctionArguments(llvmModule.get()); 214 215 // Clone module in a new LLVMContext since translateModuleToLLVMIR buries 216 // ownership too deeply. 217 // TODO(zinenko): Reevaluate model of ownership of LLVMContext in LLVMDialect. 218 std::unique_ptr<Module> deserModule = 219 LLVM::cloneModuleIntoNewContext(ctx.get(), llvmModule.get()); 220 auto dataLayout = deserModule->getDataLayout(); 221 222 // Callback to create the object layer with symbol resolution to current 223 // process and dynamically linked libraries. 224 auto objectLinkingLayerCreator = [&](ExecutionSession &session, 225 const Triple &TT) { 226 auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>( 227 session, []() { return std::make_unique<SectionMemoryManager>(); }); 228 229 // Register JIT event listeners if they are enabled. 230 if (engine->gdbListener) 231 objectLayer->registerJITEventListener(*engine->gdbListener); 232 if (engine->perfListener) 233 objectLayer->registerJITEventListener(*engine->perfListener); 234 235 // Resolve symbols from shared libraries. 236 for (auto libPath : sharedLibPaths) { 237 auto mb = llvm::MemoryBuffer::getFile(libPath); 238 if (!mb) { 239 errs() << "Fail to create MemoryBuffer for: " << libPath << "\n"; 240 continue; 241 } 242 auto &JD = session.createBareJITDylib(std::string(libPath)); 243 auto loaded = DynamicLibrarySearchGenerator::Load( 244 libPath.data(), dataLayout.getGlobalPrefix()); 245 if (!loaded) { 246 errs() << "Could not load " << libPath << ":\n " << loaded.takeError() 247 << "\n"; 248 continue; 249 } 250 JD.addGenerator(std::move(*loaded)); 251 cantFail(objectLayer->add(JD, std::move(mb.get()))); 252 } 253 254 return objectLayer; 255 }; 256 257 // Callback to inspect the cache and recompile on demand. This follows Lang's 258 // LLJITWithObjectCache example. 259 auto compileFunctionCreator = [&](JITTargetMachineBuilder JTMB) 260 -> Expected<std::unique_ptr<IRCompileLayer::IRCompiler>> { 261 if (jitCodeGenOptLevel) 262 JTMB.setCodeGenOptLevel(jitCodeGenOptLevel.getValue()); 263 auto TM = JTMB.createTargetMachine(); 264 if (!TM) 265 return TM.takeError(); 266 return std::make_unique<TMOwningSimpleCompiler>(std::move(*TM), 267 engine->cache.get()); 268 }; 269 270 // Create the LLJIT by calling the LLJITBuilder with 2 callbacks. 271 auto jit = 272 cantFail(llvm::orc::LLJITBuilder() 273 .setCompileFunctionCreator(compileFunctionCreator) 274 .setObjectLinkingLayerCreator(objectLinkingLayerCreator) 275 .create()); 276 277 // Add a ThreadSafemodule to the engine and return. 278 ThreadSafeModule tsm(std::move(deserModule), std::move(ctx)); 279 if (transformer) 280 cantFail(tsm.withModuleDo( 281 [&](llvm::Module &module) { return transformer(&module); })); 282 cantFail(jit->addIRModule(std::move(tsm))); 283 engine->jit = std::move(jit); 284 285 // Resolve symbols that are statically linked in the current process. 286 llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib(); 287 mainJD.addGenerator( 288 cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess( 289 dataLayout.getGlobalPrefix()))); 290 291 return std::move(engine); 292 } 293 294 Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const { 295 auto expectedSymbol = jit->lookup(makePackedFunctionName(name)); 296 297 // JIT lookup may return an Error referring to strings stored internally by 298 // the JIT. If the Error outlives the ExecutionEngine, it would want have a 299 // dangling reference, which is currently caught by an assertion inside JIT 300 // thanks to hand-rolled reference counting. Rewrap the error message into a 301 // string before returning. Alternatively, ORC JIT should consider copying 302 // the string into the error message. 303 if (!expectedSymbol) { 304 std::string errorMessage; 305 llvm::raw_string_ostream os(errorMessage); 306 llvm::handleAllErrors(expectedSymbol.takeError(), 307 [&os](llvm::ErrorInfoBase &ei) { ei.log(os); }); 308 return make_string_error(os.str()); 309 } 310 311 auto rawFPtr = expectedSymbol->getAddress(); 312 auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr); 313 if (!fptr) 314 return make_string_error("looked up function is null"); 315 return fptr; 316 } 317 318 Error ExecutionEngine::invoke(StringRef name, MutableArrayRef<void *> args) { 319 auto expectedFPtr = lookup(name); 320 if (!expectedFPtr) 321 return expectedFPtr.takeError(); 322 auto fptr = *expectedFPtr; 323 324 (*fptr)(args.data()); 325 326 return Error::success(); 327 } 328