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
30using 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
35static 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.
52static std::unique_ptr<llvm::Module>
53convertMLIRModule(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
82int 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

source code of mlir/tools/mlir-runner/mlir-runner.cpp