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