1 //===- ExecutionEngine.h - MLIR Execution engine and utils -----*- C++ -*--===//
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 provides a JIT-backed execution engine for MLIR modules.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_
14 #define MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_
15 
16 #include "mlir/Support/LLVM.h"
17 #include "llvm/ExecutionEngine/ObjectCache.h"
18 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
19 #include "llvm/IR/LLVMContext.h"
20 #include "llvm/Support/Error.h"
21 
22 #include <functional>
23 #include <memory>
24 
25 namespace llvm {
26 template <typename T> class Expected;
27 class Module;
28 class ExecutionEngine;
29 class JITEventListener;
30 class MemoryBuffer;
31 } // namespace llvm
32 
33 namespace mlir {
34 
35 class ModuleOp;
36 
37 /// A simple object cache following Lang's LLJITWithObjectCache example.
38 class SimpleObjectCache : public llvm::ObjectCache {
39 public:
40   void notifyObjectCompiled(const llvm::Module *M,
41                             llvm::MemoryBufferRef ObjBuffer) override;
42   std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *M) override;
43 
44   /// Dump cached object to output file `filename`.
45   void dumpToObjectFile(StringRef filename);
46 
47 private:
48   llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
49 };
50 
51 /// JIT-backed execution engine for MLIR modules.  Assumes the module can be
52 /// converted to LLVM IR.  For each function, creates a wrapper function with
53 /// the fixed interface
54 ///
55 ///     void _mlir_funcName(void **)
56 ///
57 /// where the only argument is interpreted as a list of pointers to the actual
58 /// arguments of the function, followed by a pointer to the result.  This allows
59 /// the engine to provide the caller with a generic function pointer that can
60 /// be used to invoke the JIT-compiled function.
61 class ExecutionEngine {
62 public:
63   ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener,
64                   bool enablePerfNotificationListener);
65 
66   /// Creates an execution engine for the given module.
67   ///
68   /// If `llvmModuleBuilder` is provided, it will be used to create LLVM module
69   /// from the given MLIR module. Otherwise, a default `translateModuleToLLVMIR`
70   /// function will be used to translate MLIR module to LLVM IR.
71   ///
72   /// If `transformer` is provided, it will be called on the LLVM module during
73   /// JIT-compilation and can be used, e.g., for reporting or optimization.
74   ///
75   /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for
76   /// target code generation.
77   ///
78   /// If `sharedLibPaths` are provided, the underlying JIT-compilation will
79   /// open and link the shared libraries for symbol resolution.
80   ///
81   /// If `enableObjectCache` is set, the JIT compiler will create one to store
82   /// the object generated for the given module.
83   ///
84   /// If enable `enableGDBNotificationListener` is set, the JIT compiler will
85   /// notify the llvm's global GDB notification listener.
86   ///
87   /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
88   /// the llvm's global Perf notification listener.
89   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
90   create(ModuleOp m,
91          llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
92                                                           llvm::LLVMContext &)>
93              llvmModuleBuilder = nullptr,
94          llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {},
95          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
96          ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
97          bool enableGDBNotificationListener = true,
98          bool enablePerfNotificationListener = true);
99 
100   /// Looks up a packed-argument function with the given name and returns a
101   /// pointer to it.  Propagates errors in case of failure.
102   llvm::Expected<void (*)(void **)> lookup(StringRef name) const;
103 
104   /// Invokes the function with the given name passing it the list of arguments
105   /// as a list of opaque pointers.
106   llvm::Error invoke(StringRef name, MutableArrayRef<void *> args = llvm::None);
107 
108   /// Set the target triple on the module. This is implicitly done when creating
109   /// the engine.
110   static bool setupTargetTriple(llvm::Module *llvmModule);
111 
112   /// Dump object code to output file `filename`.
113   void dumpToObjectFile(StringRef filename);
114 
115   /// Register symbols with this ExecutionEngine.
116   void registerSymbols(
117       llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
118           symbolMap);
119 
120 private:
121   /// Ordering of llvmContext and jit is important for destruction purposes: the
122   /// jit must be destroyed before the context.
123   llvm::LLVMContext llvmContext;
124 
125   /// Underlying LLJIT.
126   std::unique_ptr<llvm::orc::LLJIT> jit;
127 
128   /// Underlying cache.
129   std::unique_ptr<SimpleObjectCache> cache;
130 
131   /// GDB notification listener.
132   llvm::JITEventListener *gdbListener;
133 
134   /// Perf notification listener.
135   llvm::JITEventListener *perfListener;
136 };
137 
138 } // end namespace mlir
139 
140 #endif // MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_
141