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