1 //===- ConvertGPUToSPIRV.cpp - Convert GPU ops to SPIR-V dialect ----------===//
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 the conversion patterns from GPU ops to SPIR-V dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRV.h"
13 #include "mlir/Dialect/GPU/GPUDialect.h"
14 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
15 #include "mlir/Dialect/SPIRV/SPIRVLowering.h"
16 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
17 #include "mlir/IR/Module.h"
18
19 using namespace mlir;
20
21 namespace {
22 /// Pattern lowering GPU block/thread size/id to loading SPIR-V invocation
23 /// builtin variables.
24 template <typename SourceOp, spirv::BuiltIn builtin>
25 class LaunchConfigConversion : public SPIRVOpLowering<SourceOp> {
26 public:
27 using SPIRVOpLowering<SourceOp>::SPIRVOpLowering;
28
29 LogicalResult
30 matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
31 ConversionPatternRewriter &rewriter) const override;
32 };
33
34 /// Pattern lowering subgoup size/id to loading SPIR-V invocation
35 /// builtin variables.
36 template <typename SourceOp, spirv::BuiltIn builtin>
37 class SingleDimLaunchConfigConversion : public SPIRVOpLowering<SourceOp> {
38 public:
39 using SPIRVOpLowering<SourceOp>::SPIRVOpLowering;
40
41 LogicalResult
42 matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
43 ConversionPatternRewriter &rewriter) const override;
44 };
45
46 /// This is separate because in Vulkan workgroup size is exposed to shaders via
47 /// a constant with WorkgroupSize decoration. So here we cannot generate a
48 /// builtin variable; instead the information in the `spv.entry_point_abi`
49 /// attribute on the surrounding FuncOp is used to replace the gpu::BlockDimOp.
50 class WorkGroupSizeConversion : public SPIRVOpLowering<gpu::BlockDimOp> {
51 public:
52 using SPIRVOpLowering<gpu::BlockDimOp>::SPIRVOpLowering;
53
54 LogicalResult
55 matchAndRewrite(gpu::BlockDimOp op, ArrayRef<Value> operands,
56 ConversionPatternRewriter &rewriter) const override;
57 };
58
59 /// Pattern to convert a kernel function in GPU dialect within a spv.module.
60 class GPUFuncOpConversion final : public SPIRVOpLowering<gpu::GPUFuncOp> {
61 public:
62 using SPIRVOpLowering<gpu::GPUFuncOp>::SPIRVOpLowering;
63
64 LogicalResult
65 matchAndRewrite(gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
66 ConversionPatternRewriter &rewriter) const override;
67
68 private:
69 SmallVector<int32_t, 3> workGroupSizeAsInt32;
70 };
71
72 /// Pattern to convert a gpu.module to a spv.module.
73 class GPUModuleConversion final : public SPIRVOpLowering<gpu::GPUModuleOp> {
74 public:
75 using SPIRVOpLowering<gpu::GPUModuleOp>::SPIRVOpLowering;
76
77 LogicalResult
78 matchAndRewrite(gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
79 ConversionPatternRewriter &rewriter) const override;
80 };
81
82 /// Pattern to convert a gpu.return into a SPIR-V return.
83 // TODO: This can go to DRR when GPU return has operands.
84 class GPUReturnOpConversion final : public SPIRVOpLowering<gpu::ReturnOp> {
85 public:
86 using SPIRVOpLowering<gpu::ReturnOp>::SPIRVOpLowering;
87
88 LogicalResult
89 matchAndRewrite(gpu::ReturnOp returnOp, ArrayRef<Value> operands,
90 ConversionPatternRewriter &rewriter) const override;
91 };
92
93 } // namespace
94
95 //===----------------------------------------------------------------------===//
96 // Builtins.
97 //===----------------------------------------------------------------------===//
98
getLaunchConfigIndex(Operation * op)99 static Optional<int32_t> getLaunchConfigIndex(Operation *op) {
100 auto dimAttr = op->getAttrOfType<StringAttr>("dimension");
101 if (!dimAttr) {
102 return {};
103 }
104 if (dimAttr.getValue() == "x") {
105 return 0;
106 } else if (dimAttr.getValue() == "y") {
107 return 1;
108 } else if (dimAttr.getValue() == "z") {
109 return 2;
110 }
111 return {};
112 }
113
114 template <typename SourceOp, spirv::BuiltIn builtin>
matchAndRewrite(SourceOp op,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const115 LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
116 SourceOp op, ArrayRef<Value> operands,
117 ConversionPatternRewriter &rewriter) const {
118 auto index = getLaunchConfigIndex(op);
119 if (!index)
120 return failure();
121
122 // SPIR-V invocation builtin variables are a vector of type <3xi32>
123 auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
124 rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
125 op, rewriter.getIntegerType(32), spirvBuiltin,
126 rewriter.getI32ArrayAttr({index.getValue()}));
127 return success();
128 }
129
130 template <typename SourceOp, spirv::BuiltIn builtin>
131 LogicalResult
matchAndRewrite(SourceOp op,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const132 SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
133 SourceOp op, ArrayRef<Value> operands,
134 ConversionPatternRewriter &rewriter) const {
135 auto spirvBuiltin = spirv::getBuiltinVariableValue(op, builtin, rewriter);
136 rewriter.replaceOp(op, spirvBuiltin);
137 return success();
138 }
139
matchAndRewrite(gpu::BlockDimOp op,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const140 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
141 gpu::BlockDimOp op, ArrayRef<Value> operands,
142 ConversionPatternRewriter &rewriter) const {
143 auto index = getLaunchConfigIndex(op);
144 if (!index)
145 return failure();
146
147 auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
148 auto val = workGroupSizeAttr.getValue<int32_t>(index.getValue());
149 auto convertedType = typeConverter.convertType(op.getResult().getType());
150 if (!convertedType)
151 return failure();
152 rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
153 op, convertedType, IntegerAttr::get(convertedType, val));
154 return success();
155 }
156
157 //===----------------------------------------------------------------------===//
158 // GPUFuncOp
159 //===----------------------------------------------------------------------===//
160
161 // Legalizes a GPU function as an entry SPIR-V function.
162 static spirv::FuncOp
lowerAsEntryFunction(gpu::GPUFuncOp funcOp,SPIRVTypeConverter & typeConverter,ConversionPatternRewriter & rewriter,spirv::EntryPointABIAttr entryPointInfo,ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo)163 lowerAsEntryFunction(gpu::GPUFuncOp funcOp, SPIRVTypeConverter &typeConverter,
164 ConversionPatternRewriter &rewriter,
165 spirv::EntryPointABIAttr entryPointInfo,
166 ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo) {
167 auto fnType = funcOp.getType();
168 if (fnType.getNumResults()) {
169 funcOp.emitError("SPIR-V lowering only supports entry functions"
170 "with no return values right now");
171 return nullptr;
172 }
173 if (fnType.getNumInputs() != argABIInfo.size()) {
174 funcOp.emitError(
175 "lowering as entry functions requires ABI info for all arguments");
176 return nullptr;
177 }
178 // Update the signature to valid SPIR-V types and add the ABI
179 // attributes. These will be "materialized" by using the
180 // LowerABIAttributesPass.
181 TypeConverter::SignatureConversion signatureConverter(fnType.getNumInputs());
182 {
183 for (auto argType : enumerate(funcOp.getType().getInputs())) {
184 auto convertedType = typeConverter.convertType(argType.value());
185 signatureConverter.addInputs(argType.index(), convertedType);
186 }
187 }
188 auto newFuncOp = rewriter.create<spirv::FuncOp>(
189 funcOp.getLoc(), funcOp.getName(),
190 rewriter.getFunctionType(signatureConverter.getConvertedTypes(),
191 llvm::None));
192 for (const auto &namedAttr : funcOp.getAttrs()) {
193 if (namedAttr.first == impl::getTypeAttrName() ||
194 namedAttr.first == SymbolTable::getSymbolAttrName())
195 continue;
196 newFuncOp.setAttr(namedAttr.first, namedAttr.second);
197 }
198
199 rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
200 newFuncOp.end());
201 if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), typeConverter,
202 &signatureConverter)))
203 return nullptr;
204 rewriter.eraseOp(funcOp);
205
206 spirv::setABIAttrs(newFuncOp, entryPointInfo, argABIInfo);
207 return newFuncOp;
208 }
209
210 /// Populates `argABI` with spv.interface_var_abi attributes for lowering
211 /// gpu.func to spv.func if no arguments have the attributes set
212 /// already. Returns failure if any argument has the ABI attribute set already.
213 static LogicalResult
getDefaultABIAttrs(MLIRContext * context,gpu::GPUFuncOp funcOp,SmallVectorImpl<spirv::InterfaceVarABIAttr> & argABI)214 getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
215 SmallVectorImpl<spirv::InterfaceVarABIAttr> &argABI) {
216 for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
217 if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
218 argIndex, spirv::getInterfaceVarABIAttrName()))
219 return failure();
220 // Vulkan's interface variable requirements needs scalars to be wrapped in a
221 // struct. The struct held in storage buffer.
222 Optional<spirv::StorageClass> sc;
223 if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
224 sc = spirv::StorageClass::StorageBuffer;
225 argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
226 }
227 return success();
228 }
229
matchAndRewrite(gpu::GPUFuncOp funcOp,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const230 LogicalResult GPUFuncOpConversion::matchAndRewrite(
231 gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
232 ConversionPatternRewriter &rewriter) const {
233 if (!gpu::GPUDialect::isKernel(funcOp))
234 return failure();
235
236 SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
237 if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
238 argABI.clear();
239 for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
240 // If the ABI is already specified, use it.
241 auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
242 argIndex, spirv::getInterfaceVarABIAttrName());
243 if (!abiAttr) {
244 funcOp.emitRemark(
245 "match failure: missing 'spv.interface_var_abi' attribute at "
246 "argument ")
247 << argIndex;
248 return failure();
249 }
250 argABI.push_back(abiAttr);
251 }
252 }
253
254 auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
255 if (!entryPointAttr) {
256 funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
257 return failure();
258 }
259 spirv::FuncOp newFuncOp = lowerAsEntryFunction(
260 funcOp, typeConverter, rewriter, entryPointAttr, argABI);
261 if (!newFuncOp)
262 return failure();
263 newFuncOp.removeAttr(Identifier::get(gpu::GPUDialect::getKernelFuncAttrName(),
264 rewriter.getContext()));
265 return success();
266 }
267
268 //===----------------------------------------------------------------------===//
269 // ModuleOp with gpu.module.
270 //===----------------------------------------------------------------------===//
271
matchAndRewrite(gpu::GPUModuleOp moduleOp,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const272 LogicalResult GPUModuleConversion::matchAndRewrite(
273 gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
274 ConversionPatternRewriter &rewriter) const {
275 auto spvModule = rewriter.create<spirv::ModuleOp>(
276 moduleOp.getLoc(), spirv::AddressingModel::Logical,
277 spirv::MemoryModel::GLSL450);
278
279 // Move the region from the module op into the SPIR-V module.
280 Region &spvModuleRegion = spvModule.body();
281 rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
282 spvModuleRegion.begin());
283 // The spv.module build method adds a block with a terminator. Remove that
284 // block. The terminator of the module op in the remaining block will be
285 // legalized later.
286 rewriter.eraseBlock(&spvModuleRegion.back());
287 rewriter.eraseOp(moduleOp);
288 return success();
289 }
290
291 //===----------------------------------------------------------------------===//
292 // GPU return inside kernel functions to SPIR-V return.
293 //===----------------------------------------------------------------------===//
294
matchAndRewrite(gpu::ReturnOp returnOp,ArrayRef<Value> operands,ConversionPatternRewriter & rewriter) const295 LogicalResult GPUReturnOpConversion::matchAndRewrite(
296 gpu::ReturnOp returnOp, ArrayRef<Value> operands,
297 ConversionPatternRewriter &rewriter) const {
298 if (!operands.empty())
299 return failure();
300
301 rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
302 return success();
303 }
304
305 //===----------------------------------------------------------------------===//
306 // GPU To SPIRV Patterns.
307 //===----------------------------------------------------------------------===//
308
309 namespace {
310 #include "GPUToSPIRV.cpp.inc"
311 }
312
populateGPUToSPIRVPatterns(MLIRContext * context,SPIRVTypeConverter & typeConverter,OwningRewritePatternList & patterns)313 void mlir::populateGPUToSPIRVPatterns(MLIRContext *context,
314 SPIRVTypeConverter &typeConverter,
315 OwningRewritePatternList &patterns) {
316 populateWithGenerated(context, &patterns);
317 patterns.insert<
318 GPUFuncOpConversion, GPUModuleConversion, GPUReturnOpConversion,
319 LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
320 LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
321 LaunchConfigConversion<gpu::ThreadIdOp,
322 spirv::BuiltIn::LocalInvocationId>,
323 SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
324 spirv::BuiltIn::SubgroupId>,
325 SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
326 spirv::BuiltIn::NumSubgroups>,
327 SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
328 spirv::BuiltIn::SubgroupSize>,
329 WorkGroupSizeConversion>(context, typeConverter);
330 }
331