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