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