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/ExecutionEngine/SectionMemoryManager.h"
20#include "llvm/IR/LLVMContext.h"
21#include "llvm/Support/Error.h"
22
23#include <functional>
24#include <memory>
25#include <optional>
26
27namespace llvm {
28template <typename T>
29class Expected;
30class Module;
31class ExecutionEngine;
32class JITEventListener;
33class MemoryBuffer;
34} // namespace llvm
35
36namespace mlir {
37
38class Operation;
39
40/// A simple object cache following Lang's LLJITWithObjectCache example.
41class SimpleObjectCache : public llvm::ObjectCache {
42public:
43 void notifyObjectCompiled(const llvm::Module *m,
44 llvm::MemoryBufferRef objBuffer) override;
45 std::unique_ptr<llvm::MemoryBuffer> getObject(const llvm::Module *m) override;
46
47 /// Dump cached object to output file `filename`.
48 void dumpToObjectFile(StringRef filename);
49
50 /// Returns `true` if cache hasn't been populated yet.
51 bool isEmpty();
52
53private:
54 llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> cachedObjects;
55};
56
57struct ExecutionEngineOptions {
58 /// If `llvmModuleBuilder` is provided, it will be used to create an LLVM
59 /// module from the given MLIR IR. Otherwise, a default
60 /// `translateModuleToLLVMIR` function will be used to translate to LLVM IR.
61 llvm::function_ref<std::unique_ptr<llvm::Module>(Operation *,
62 llvm::LLVMContext &)>
63 llvmModuleBuilder = nullptr;
64
65 /// If `transformer` is provided, it will be called on the LLVM module during
66 /// JIT-compilation and can be used, e.g., for reporting or optimization.
67 llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {};
68
69 /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for
70 /// target code generation.
71 std::optional<llvm::CodeGenOptLevel> jitCodeGenOptLevel;
72
73 /// If `sharedLibPaths` are provided, the underlying JIT-compilation will
74 /// open and link the shared libraries for symbol resolution. Libraries that
75 /// are designed to be used with the `ExecutionEngine` may implement a
76 /// loading and unloading protocol: if they implement the two functions with
77 /// the names defined in `kLibraryInitFnName` and `kLibraryDestroyFnName`,
78 /// these functions will be called upon loading the library and upon
79 /// destruction of the `ExecutionEngine`. In the init function, the library
80 /// may provide a list of symbols that it wants to make available to code
81 /// run by the `ExecutionEngine`. If the two functions are not defined, only
82 /// symbols with public visibility are available to the executed code.
83 ArrayRef<StringRef> sharedLibPaths = {};
84
85 /// Specifies an existing `sectionMemoryMapper` to be associated with the
86 /// compiled code. If none is provided, a default memory mapper that directly
87 /// calls into the operating system is used.
88 llvm::SectionMemoryManager::MemoryMapper *sectionMemoryMapper = nullptr;
89
90 /// If `enableObjectCache` is set, the JIT compiler will create one to store
91 /// the object generated for the given module. The contents of the cache can
92 /// be dumped to a file via the `dumpToObjectFile` method.
93 bool enableObjectDump = false;
94
95 /// If enable `enableGDBNotificationListener` is set, the JIT compiler will
96 /// notify the llvm's global GDB notification listener.
97 bool enableGDBNotificationListener = true;
98
99 /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
100 /// the llvm's global Perf notification listener.
101 bool enablePerfNotificationListener = true;
102};
103
104/// JIT-backed execution engine for MLIR. Assumes the IR can be converted to
105/// LLVM IR. For each function, creates a wrapper function with the fixed
106/// interface
107///
108/// void _mlir_funcName(void **)
109///
110/// where the only argument is interpreted as a list of pointers to the actual
111/// arguments of the function, followed by a pointer to the result. This allows
112/// the engine to provide the caller with a generic function pointer that can
113/// be used to invoke the JIT-compiled function.
114class ExecutionEngine {
115public:
116 /// Name of init functions of shared libraries. If a library provides a
117 /// function with this name and the one of the destroy function, this function
118 /// is called upon loading the library.
119 static constexpr const char *const kLibraryInitFnName =
120 "__mlir_execution_engine_init";
121
122 /// Name of destroy functions of shared libraries. If a library provides a
123 /// function with this name and the one of the init function, this function is
124 /// called upon destructing the `ExecutionEngine`.
125 static constexpr const char *const kLibraryDestroyFnName =
126 "__mlir_execution_engine_destroy";
127
128 /// Function type for init functions of shared libraries. The library may
129 /// provide a list of symbols that it wants to make available to code run by
130 /// the `ExecutionEngine`. If the two functions are not defined, only symbols
131 /// with public visibility are available to the executed code.
132 using LibraryInitFn = void (*)(llvm::StringMap<void *> &);
133
134 /// Function type for destroy functions of shared libraries.
135 using LibraryDestroyFn = void (*)();
136
137 ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener,
138 bool enablePerfNotificationListener);
139
140 ~ExecutionEngine();
141
142 /// Creates an execution engine for the given MLIR IR. If TargetMachine is
143 /// not provided, default TM is created (i.e. ignoring any command line flags
144 /// that could affect the set-up).
145 static llvm::Expected<std::unique_ptr<ExecutionEngine>>
146 create(Operation *op, const ExecutionEngineOptions &options = {},
147 std::unique_ptr<llvm::TargetMachine> tm = nullptr);
148
149 /// Looks up a packed-argument function wrapping the function with the given
150 /// name and returns a pointer to it. Propagates errors in case of failure.
151 llvm::Expected<void (*)(void **)> lookupPacked(StringRef name) const;
152
153 /// Looks up the original function with the given name and returns a
154 /// pointer to it. This is not necesarily a packed function. Propagates
155 /// errors in case of failure.
156 llvm::Expected<void *> lookup(StringRef name) const;
157
158 /// Invokes the function with the given name passing it the list of opaque
159 /// pointers to the actual arguments.
160 llvm::Error invokePacked(StringRef name,
161 MutableArrayRef<void *> args = std::nullopt);
162
163 /// Trait that defines how a given type is passed to the JIT code. This
164 /// defaults to passing the address but can be specialized.
165 template <typename T>
166 struct Argument {
167 static void pack(SmallVectorImpl<void *> &args, T &val) {
168 args.push_back(Elt: &val);
169 }
170 };
171
172 /// Tag to wrap an output parameter when invoking a jitted function.
173 template <typename T>
174 struct Result {
175 Result(T &result) : value(result) {}
176 T &value;
177 };
178
179 /// Helper function to wrap an output operand when using
180 /// ExecutionEngine::invoke.
181 template <typename T>
182 static Result<T> result(T &t) {
183 return Result<T>(t);
184 }
185
186 // Specialization for output parameter: their address is forwarded directly to
187 // the native code.
188 template <typename T>
189 struct Argument<Result<T>> {
190 static void pack(SmallVectorImpl<void *> &args, Result<T> &result) {
191 args.push_back(Elt: &result.value);
192 }
193 };
194
195 /// Invokes the function with the given name passing it the list of arguments
196 /// by value. Function result can be obtain through output parameter using the
197 /// `Result` wrapper defined above. For example:
198 ///
199 /// func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface }
200 ///
201 /// can be invoked:
202 ///
203 /// int32_t result = 0;
204 /// llvm::Error error = jit->invoke("foo", 42,
205 /// result(result));
206 template <typename... Args>
207 llvm::Error invoke(StringRef funcName, Args... args) {
208 const std::string adapterName =
209 std::string("_mlir_ciface_") + funcName.str();
210 llvm::SmallVector<void *> argsArray;
211 // Pack every arguments in an array of pointers. Delegate the packing to a
212 // trait so that it can be overridden per argument type.
213 (Argument<Args>::pack(argsArray, args), ...);
214 return invokePacked(name: adapterName, args: argsArray);
215 }
216
217 /// Set the target triple and the data layout for the input module based on
218 /// the input TargetMachine. This is implicitly done when creating the
219 /// engine.
220 static void setupTargetTripleAndDataLayout(llvm::Module *llvmModule,
221 llvm::TargetMachine *tm);
222
223 /// Dump object code to output file `filename`.
224 void dumpToObjectFile(StringRef filename);
225
226 /// Register symbols with this ExecutionEngine.
227 void registerSymbols(
228 llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
229 symbolMap);
230
231private:
232 /// Ordering of llvmContext and jit is important for destruction purposes: the
233 /// jit must be destroyed before the context.
234 llvm::LLVMContext llvmContext;
235
236 /// Underlying LLJIT.
237 std::unique_ptr<llvm::orc::LLJIT> jit;
238
239 /// Underlying cache.
240 std::unique_ptr<SimpleObjectCache> cache;
241
242 /// Names of functions that may be looked up.
243 std::vector<std::string> functionNames;
244
245 /// GDB notification listener.
246 llvm::JITEventListener *gdbListener;
247
248 /// Perf notification listener.
249 llvm::JITEventListener *perfListener;
250
251 /// Destroy functions in the libraries loaded by the ExecutionEngine that are
252 /// called when this ExecutionEngine is destructed.
253 SmallVector<LibraryDestroyFn> destroyFns;
254};
255
256} // namespace mlir
257
258#endif // MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_
259

source code of mlir/include/mlir/ExecutionEngine/ExecutionEngine.h