xref: /llvm-project/mlir/include/mlir/ExecutionEngine/JitRunner.h (revision db791b278a414fb6df1acc1799adcf11d8fb9169)
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_EXECUTIONENGINE_JITRUNNER_H
19 #define MLIR_EXECUTIONENGINE_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 struct LogicalResult;
28 
29 namespace orc {
30 class MangleAndInterner;
31 } // namespace orc
32 } // namespace llvm
33 
34 namespace mlir {
35 
36 class DialectRegistry;
37 class Operation;
38 
39 /// JitRunner command line options used by JitRunnerConfig methods
40 struct JitRunnerOptions {
41   /// The name of the main function
42   llvm::StringRef mainFuncName;
43   /// The type of the main function (as string, from cmd-line)
44   llvm::StringRef mainFuncType;
45 };
46 
47 /// Configuration to override functionality of the JitRunner
48 struct JitRunnerConfig {
49   /// MLIR transformer applied after parsing the input into MLIR IR and before
50   /// passing the MLIR IR to the ExecutionEngine.
51   llvm::function_ref<llvm::LogicalResult(mlir::Operation *,
52                                          JitRunnerOptions &options)>
53       mlirTransformer = nullptr;
54 
55   /// A custom function that is passed to ExecutionEngine. It processes MLIR and
56   /// creates an LLVM IR module.
57   llvm::function_ref<std::unique_ptr<llvm::Module>(Operation *,
58                                                    llvm::LLVMContext &)>
59       llvmModuleBuilder = nullptr;
60 
61   /// A callback to register symbols with ExecutionEngine at runtime.
62   llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
63       runtimesymbolMap = nullptr;
64 };
65 
66 /// Entry point for all CPU runners. Expects the common argc/argv arguments for
67 /// standard C++ main functions. The supplied dialect registry is expected to
68 /// contain any registers that appear in the input IR, they will be loaded
69 /// on-demand by the parser.
70 int JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
71                   JitRunnerConfig config = {});
72 
73 } // namespace mlir
74 
75 #endif // MLIR_EXECUTIONENGINE_JITRUNNER_H
76