1 //===- JitRunner.h - MLIR CPU Execution Driver Library ----------*- 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 is a library that provides a shared implementation for command line
10 // utilities that execute an MLIR file on the CPU by translating MLIR to LLVM
11 // IR before JIT-compiling and executing the latter.
12 //
13 // The translation can be customized by providing an MLIR to MLIR
14 // transformation.
15 //
16 //===----------------------------------------------------------------------===//
17 
18 #ifndef MLIR_SUPPORT_JITRUNNER_H_
19 #define MLIR_SUPPORT_JITRUNNER_H_
20 
21 #include "llvm/ADT/STLExtras.h"
22 #include "llvm/ExecutionEngine/Orc/Core.h"
23 
24 namespace llvm {
25 class Module;
26 class LLVMContext;
27 
28 namespace orc {
29 class MangleAndInterner;
30 } // namespace orc
31 } // namespace llvm
32 
33 namespace mlir {
34 
35 class DialectRegistry;
36 class ModuleOp;
37 struct LogicalResult;
38 
39 struct JitRunnerConfig {
40   /// MLIR transformer applied after parsing the input into MLIR IR and before
41   /// passing the MLIR module to the ExecutionEngine.
42   llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer = nullptr;
43 
44   /// A custom function that is passed to ExecutionEngine. It processes MLIR
45   /// module and creates LLVM IR module.
46   llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
47                                                    llvm::LLVMContext &)>
48       llvmModuleBuilder = nullptr;
49 
50   /// A callback to register symbols with ExecutionEngine at runtime.
51   llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
52       runtimesymbolMap = nullptr;
53 };
54 
55 /// Entry point for all CPU runners. Expects the common argc/argv arguments for
56 /// standard C++ main functions. The supplied dialect registry is expected to
57 /// contain any registers that appear in the input IR, they will be loaded
58 /// on-demand by the parser.
59 int JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
60                   JitRunnerConfig config = {});
61 
62 } // namespace mlir
63 
64 #endif // MLIR_SUPPORT_JITRUNNER_H_
65