1 //===- mlir-spirv-cpu-runner.cpp - MLIR SPIR-V Execution on CPU -----------===//
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 GPU module and host part to LLVM IR before
11 // JIT-compiling and executing.
12 //
13 //===----------------------------------------------------------------------===//
14
15 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
16 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
17 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
18 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
19 #include "mlir/Dialect/Func/IR/FuncOps.h"
20 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
21 #include "mlir/Dialect/GPU/Transforms/Passes.h"
22 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
23 #include "mlir/Dialect/MemRef/IR/MemRef.h"
24 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
25 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
26 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
27 #include "mlir/ExecutionEngine/JitRunner.h"
28 #include "mlir/ExecutionEngine/OptUtils.h"
29 #include "mlir/Pass/Pass.h"
30 #include "mlir/Pass/PassManager.h"
31 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
32 #include "mlir/Target/LLVMIR/Export.h"
33
34 #include "llvm/IR/LLVMContext.h"
35 #include "llvm/IR/Module.h"
36 #include "llvm/Linker/Linker.h"
37 #include "llvm/Support/InitLLVM.h"
38 #include "llvm/Support/TargetSelect.h"
39
40 using namespace mlir;
41
42 /// A utility function that builds llvm::Module from two nested MLIR modules.
43 ///
44 /// module @main {
45 /// module @kernel {
46 /// // Some ops
47 /// }
48 /// // Some other ops
49 /// }
50 ///
51 /// Each of these two modules is translated to LLVM IR module, then they are
52 /// linked together and returned.
53 static std::unique_ptr<llvm::Module>
convertMLIRModule(ModuleOp module,llvm::LLVMContext & context)54 convertMLIRModule(ModuleOp module, llvm::LLVMContext &context) {
55 // Verify that there is only one nested module.
56 auto modules = module.getOps<ModuleOp>();
57 if (!llvm::hasSingleElement(modules)) {
58 module.emitError("The module must contain exactly one nested module");
59 return nullptr;
60 }
61
62 // Translate nested module and erase it.
63 ModuleOp nested = *modules.begin();
64 std::unique_ptr<llvm::Module> kernelModule =
65 translateModuleToLLVMIR(nested, context);
66 nested.erase();
67
68 std::unique_ptr<llvm::Module> mainModule =
69 translateModuleToLLVMIR(module, context);
70 llvm::Linker::linkModules(*mainModule, std::move(kernelModule));
71 return mainModule;
72 }
73
runMLIRPasses(ModuleOp module)74 static LogicalResult runMLIRPasses(ModuleOp module) {
75 PassManager passManager(module.getContext());
76 applyPassManagerCLOptions(passManager);
77 passManager.addPass(createGpuKernelOutliningPass());
78 passManager.addPass(createConvertGPUToSPIRVPass());
79
80 OpPassManager &nestedPM = passManager.nest<spirv::ModuleOp>();
81 nestedPM.addPass(spirv::createLowerABIAttributesPass());
82 nestedPM.addPass(spirv::createUpdateVersionCapabilityExtensionPass());
83 passManager.addPass(createLowerHostCodeToLLVMPass());
84 passManager.addPass(createConvertSPIRVToLLVMPass());
85 return passManager.run(module);
86 }
87
main(int argc,char ** argv)88 int main(int argc, char **argv) {
89 llvm::InitLLVM y(argc, argv);
90
91 llvm::InitializeNativeTarget();
92 llvm::InitializeNativeTargetAsmPrinter();
93
94 mlir::JitRunnerConfig jitRunnerConfig;
95 jitRunnerConfig.mlirTransformer = runMLIRPasses;
96 jitRunnerConfig.llvmModuleBuilder = convertMLIRModule;
97
98 mlir::DialectRegistry registry;
99 registry.insert<mlir::arith::ArithmeticDialect, mlir::LLVM::LLVMDialect,
100 mlir::gpu::GPUDialect, mlir::spirv::SPIRVDialect,
101 mlir::func::FuncDialect, mlir::memref::MemRefDialect>();
102 mlir::registerLLVMDialectTranslation(registry);
103
104 return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
105 }
106