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