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