1 | //===- GPUToSPIRVPass.cpp - GPU to SPIR-V Passes --------------------------===// |
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 file implements a pass to convert a kernel function in the GPU Dialect |
10 | // into a spirv.module operation. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h" |
15 | |
16 | #include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h" |
17 | #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h" |
18 | #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h" |
19 | #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h" |
20 | #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h" |
21 | #include "mlir/Dialect/Func/IR/FuncOps.h" |
22 | #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
23 | #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" |
24 | #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" |
25 | #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h" |
26 | #include "mlir/IR/PatternMatch.h" |
27 | |
28 | namespace mlir { |
29 | #define GEN_PASS_DEF_CONVERTGPUTOSPIRV |
30 | #include "mlir/Conversion/Passes.h.inc" |
31 | } // namespace mlir |
32 | |
33 | using namespace mlir; |
34 | |
35 | namespace { |
36 | /// Pass to lower GPU Dialect to SPIR-V. The pass only converts the gpu.func ops |
37 | /// inside gpu.module ops. i.e., the function that are referenced in |
38 | /// gpu.launch_func ops. For each such function |
39 | /// |
40 | /// 1) Create a spirv::ModuleOp, and clone the function into spirv::ModuleOp |
41 | /// (the original function is still needed by the gpu::LaunchKernelOp, so cannot |
42 | /// replace it). |
43 | /// |
44 | /// 2) Lower the body of the spirv::ModuleOp. |
45 | struct GPUToSPIRVPass final : impl::ConvertGPUToSPIRVBase<GPUToSPIRVPass> { |
46 | explicit GPUToSPIRVPass(bool mapMemorySpace) |
47 | : mapMemorySpace(mapMemorySpace) {} |
48 | void runOnOperation() override; |
49 | |
50 | private: |
51 | bool mapMemorySpace; |
52 | }; |
53 | |
54 | void GPUToSPIRVPass::runOnOperation() { |
55 | MLIRContext *context = &getContext(); |
56 | ModuleOp module = getOperation(); |
57 | |
58 | SmallVector<Operation *, 1> gpuModules; |
59 | OpBuilder builder(context); |
60 | |
61 | auto targetEnvSupportsKernelCapability = [](gpu::GPUModuleOp moduleOp) { |
62 | Operation *gpuModule = moduleOp.getOperation(); |
63 | auto targetAttr = spirv::lookupTargetEnvOrDefault(op: gpuModule); |
64 | spirv::TargetEnv targetEnv(targetAttr); |
65 | return targetEnv.allows(spirv::Capability::Kernel); |
66 | }; |
67 | |
68 | module.walk([&](gpu::GPUModuleOp moduleOp) { |
69 | // Clone each GPU kernel module for conversion, given that the GPU |
70 | // launch op still needs the original GPU kernel module. |
71 | // For Vulkan Shader capabilities, we insert the newly converted SPIR-V |
72 | // module right after the original GPU module, as that's the expectation of |
73 | // the in-tree Vulkan runner. |
74 | // For OpenCL Kernel capabilities, we insert the newly converted SPIR-V |
75 | // module inside the original GPU module, as that's the expectaion of the |
76 | // normal GPU compilation pipeline. |
77 | if (targetEnvSupportsKernelCapability(moduleOp)) { |
78 | builder.setInsertionPoint(moduleOp.getBody(), |
79 | moduleOp.getBody()->begin()); |
80 | } else { |
81 | builder.setInsertionPoint(moduleOp.getOperation()); |
82 | } |
83 | gpuModules.push_back(builder.clone(*moduleOp.getOperation())); |
84 | }); |
85 | |
86 | // Run conversion for each module independently as they can have different |
87 | // TargetEnv attributes. |
88 | for (Operation *gpuModule : gpuModules) { |
89 | spirv::TargetEnvAttr targetAttr = |
90 | spirv::lookupTargetEnvOrDefault(gpuModule); |
91 | |
92 | // Map MemRef memory space to SPIR-V storage class first if requested. |
93 | if (mapMemorySpace) { |
94 | spirv::MemorySpaceToStorageClassMap memorySpaceMap = |
95 | targetEnvSupportsKernelCapability( |
96 | dyn_cast<gpu::GPUModuleOp>(gpuModule)) |
97 | ? spirv::mapMemorySpaceToOpenCLStorageClass |
98 | : spirv::mapMemorySpaceToVulkanStorageClass; |
99 | spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap); |
100 | spirv::convertMemRefTypesAndAttrs(gpuModule, converter); |
101 | |
102 | // Check if there are any illegal ops remaining. |
103 | std::unique_ptr<ConversionTarget> target = |
104 | spirv::getMemorySpaceToStorageClassTarget(*context); |
105 | gpuModule->walk([&target, this](Operation *childOp) { |
106 | if (target->isIllegal(childOp)) { |
107 | childOp->emitOpError("failed to legalize memory space" ); |
108 | signalPassFailure(); |
109 | return WalkResult::interrupt(); |
110 | } |
111 | return WalkResult::advance(); |
112 | }); |
113 | } |
114 | |
115 | std::unique_ptr<ConversionTarget> target = |
116 | SPIRVConversionTarget::get(targetAttr); |
117 | |
118 | SPIRVConversionOptions options; |
119 | options.use64bitIndex = this->use64bitIndex; |
120 | SPIRVTypeConverter typeConverter(targetAttr, options); |
121 | populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter); |
122 | |
123 | RewritePatternSet patterns(context); |
124 | populateGPUToSPIRVPatterns(typeConverter, patterns); |
125 | populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(typeConverter, |
126 | patterns); |
127 | |
128 | // TODO: Change SPIR-V conversion to be progressive and remove the following |
129 | // patterns. |
130 | ScfToSPIRVContext scfContext; |
131 | populateSCFToSPIRVPatterns(typeConverter, scfContext, patterns); |
132 | mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns); |
133 | populateMemRefToSPIRVPatterns(typeConverter, patterns); |
134 | populateFuncToSPIRVPatterns(typeConverter, patterns); |
135 | |
136 | if (failed(applyFullConversion(gpuModule, *target, std::move(patterns)))) |
137 | return signalPassFailure(); |
138 | } |
139 | |
140 | // For OpenCL, the gpu.func op in the original gpu.module op needs to be |
141 | // replaced with an empty func.func op with the same arguments as the gpu.func |
142 | // op. The func.func op needs gpu.kernel attribute set. |
143 | module.walk([&](gpu::GPUModuleOp moduleOp) { |
144 | if (targetEnvSupportsKernelCapability(moduleOp)) { |
145 | moduleOp.walk([&](gpu::GPUFuncOp funcOp) { |
146 | builder.setInsertionPoint(funcOp); |
147 | auto newFuncOp = builder.create<func::FuncOp>( |
148 | funcOp.getLoc(), funcOp.getName(), funcOp.getFunctionType()); |
149 | auto entryBlock = newFuncOp.addEntryBlock(); |
150 | builder.setInsertionPointToEnd(entryBlock); |
151 | builder.create<func::ReturnOp>(funcOp.getLoc()); |
152 | newFuncOp->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), |
153 | builder.getUnitAttr()); |
154 | funcOp.erase(); |
155 | }); |
156 | } |
157 | }); |
158 | } |
159 | |
160 | } // namespace |
161 | |
162 | std::unique_ptr<OperationPass<ModuleOp>> |
163 | mlir::createConvertGPUToSPIRVPass(bool mapMemorySpace) { |
164 | return std::make_unique<GPUToSPIRVPass>(args&: mapMemorySpace); |
165 | } |
166 | |