1 | //===- mlir-runner.cpp - MLIR CPU Execution Driver ------------------------===// |
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 | // Main entry point to a command line utility that executes an MLIR file on the |
10 | // CPU by translating MLIR to LLVM IR before JIT-compiling and executing the |
11 | // latter. |
12 | // |
13 | //===----------------------------------------------------------------------===// |
14 | |
15 | #include "mlir/Dialect/LLVMIR/LLVMDialect.h" |
16 | #include "mlir/ExecutionEngine/JitRunner.h" |
17 | #include "mlir/ExecutionEngine/OptUtils.h" |
18 | #include "mlir/IR/Dialect.h" |
19 | #include "mlir/Target/LLVMIR/Dialect/All.h" |
20 | #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" |
21 | #include "mlir/Target/LLVMIR/Export.h" |
22 | |
23 | #include "llvm/IR/LLVMContext.h" |
24 | #include "llvm/IR/Module.h" |
25 | #include "llvm/Linker/Linker.h" |
26 | #include "llvm/Support/CommandLine.h" |
27 | #include "llvm/Support/InitLLVM.h" |
28 | #include "llvm/Support/TargetSelect.h" |
29 | |
30 | using namespace mlir; |
31 | |
32 | // TODO: Consider removing this linking functionality from the SPIR-V CPU Runner |
33 | // flow in favour of a more proper host/device split like other runners. |
34 | // https://github.com/llvm/llvm-project/issues/115348 |
35 | static llvm::cl::opt<bool> LinkNestedModules( |
36 | "link-nested-modules" , |
37 | llvm::cl::desc("Link two nested MLIR modules into a single LLVM IR module. " |
38 | "Useful if both the host and device code can be run on the " |
39 | "same CPU, as in SPIR-V CPU Runner tests." )); |
40 | |
41 | /// A utility function that builds llvm::Module from two nested MLIR modules. |
42 | /// |
43 | /// module @main { |
44 | /// module @kernel { |
45 | /// // Some ops |
46 | /// } |
47 | /// // Some other ops |
48 | /// } |
49 | /// |
50 | /// Each of these two modules is translated to LLVM IR module, then they are |
51 | /// linked together and returned. |
52 | static std::unique_ptr<llvm::Module> |
53 | convertMLIRModule(Operation *op, llvm::LLVMContext &context) { |
54 | auto module = dyn_cast<ModuleOp>(op); |
55 | if (!module) |
56 | return op->emitError(message: "op must be a 'builtin.module" ), nullptr; |
57 | |
58 | std::unique_ptr<llvm::Module> kernelModule; |
59 | if (LinkNestedModules) { |
60 | // Verify that there is only one nested module. |
61 | auto modules = module.getOps<ModuleOp>(); |
62 | if (!llvm::hasSingleElement(modules)) { |
63 | module.emitError("The module must contain exactly one nested module" ); |
64 | return nullptr; |
65 | } |
66 | |
67 | // Translate nested module and erase it. |
68 | ModuleOp nested = *modules.begin(); |
69 | kernelModule = translateModuleToLLVMIR(nested, context); |
70 | nested.erase(); |
71 | } |
72 | |
73 | std::unique_ptr<llvm::Module> mainModule = |
74 | translateModuleToLLVMIR(module, context); |
75 | |
76 | if (LinkNestedModules) |
77 | llvm::Linker::linkModules(Dest&: *mainModule, Src: std::move(kernelModule)); |
78 | |
79 | return mainModule; |
80 | } |
81 | |
82 | int main(int argc, char **argv) { |
83 | llvm::InitLLVM y(argc, argv); |
84 | llvm::InitializeNativeTarget(); |
85 | llvm::InitializeNativeTargetAsmPrinter(); |
86 | llvm::InitializeNativeTargetAsmParser(); |
87 | |
88 | mlir::DialectRegistry registry; |
89 | mlir::registerAllToLLVMIRTranslations(registry); |
90 | |
91 | mlir::JitRunnerConfig jitRunnerConfig; |
92 | jitRunnerConfig.llvmModuleBuilder = convertMLIRModule; |
93 | return mlir::JitRunnerMain(argc, argv, registry, config: jitRunnerConfig); |
94 | } |
95 | |