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   if (failed(spirv::setABIAttrs(newFuncOp, entryPointInfo, argABIInfo)))
212     return nullptr;
213   return newFuncOp;
214 }
215 
216 /// Populates `argABI` with spv.interface_var_abi attributes for lowering
217 /// gpu.func to spv.func if no arguments have the attributes set
218 /// already. Returns failure if any argument has the ABI attribute set already.
219 static LogicalResult
220 getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
221                    SmallVectorImpl<spirv::InterfaceVarABIAttr> &argABI) {
222   spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(funcOp);
223   if (!spirv::needsInterfaceVarABIAttrs(targetEnv))
224     return success();
225 
226   for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
227     if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
228             argIndex, spirv::getInterfaceVarABIAttrName()))
229       return failure();
230     // Vulkan's interface variable requirements needs scalars to be wrapped in a
231     // struct. The struct held in storage buffer.
232     Optional<spirv::StorageClass> sc;
233     if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
234       sc = spirv::StorageClass::StorageBuffer;
235     argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
236   }
237   return success();
238 }
239 
240 LogicalResult GPUFuncOpConversion::matchAndRewrite(
241     gpu::GPUFuncOp funcOp, ArrayRef<Value> operands,
242     ConversionPatternRewriter &rewriter) const {
243   if (!gpu::GPUDialect::isKernel(funcOp))
244     return failure();
245 
246   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
247   if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
248     argABI.clear();
249     for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
250       // If the ABI is already specified, use it.
251       auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
252           argIndex, spirv::getInterfaceVarABIAttrName());
253       if (!abiAttr) {
254         funcOp.emitRemark(
255             "match failure: missing 'spv.interface_var_abi' attribute at "
256             "argument ")
257             << argIndex;
258         return failure();
259       }
260       argABI.push_back(abiAttr);
261     }
262   }
263 
264   auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
265   if (!entryPointAttr) {
266     funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
267     return failure();
268   }
269   spirv::FuncOp newFuncOp = lowerAsEntryFunction(
270       funcOp, *getTypeConverter(), rewriter, entryPointAttr, argABI);
271   if (!newFuncOp)
272     return failure();
273   newFuncOp->removeAttr(Identifier::get(
274       gpu::GPUDialect::getKernelFuncAttrName(), rewriter.getContext()));
275   return success();
276 }
277 
278 //===----------------------------------------------------------------------===//
279 // ModuleOp with gpu.module.
280 //===----------------------------------------------------------------------===//
281 
282 LogicalResult GPUModuleConversion::matchAndRewrite(
283     gpu::GPUModuleOp moduleOp, ArrayRef<Value> operands,
284     ConversionPatternRewriter &rewriter) const {
285   spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(moduleOp);
286   spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv);
287   FailureOr<spirv::MemoryModel> memoryModel = spirv::getMemoryModel(targetEnv);
288   if (failed(memoryModel))
289     return moduleOp.emitRemark("match failure: could not selected memory model "
290                                "based on 'spv.target_env'");
291 
292   // Add a keyword to the module name to avoid symbolic conflict.
293   std::string spvModuleName = (kSPIRVModule + moduleOp.getName()).str();
294   auto spvModule = rewriter.create<spirv::ModuleOp>(
295       moduleOp.getLoc(), addressingModel, memoryModel.getValue(),
296       StringRef(spvModuleName));
297 
298   // Move the region from the module op into the SPIR-V module.
299   Region &spvModuleRegion = spvModule.body();
300   rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
301                               spvModuleRegion.begin());
302   // The spv.module build method adds a block with a terminator. Remove that
303   // block. The terminator of the module op in the remaining block will be
304   // legalized later.
305   rewriter.eraseBlock(&spvModuleRegion.back());
306   rewriter.eraseOp(moduleOp);
307   return success();
308 }
309 
310 //===----------------------------------------------------------------------===//
311 // GPU return inside kernel functions to SPIR-V return.
312 //===----------------------------------------------------------------------===//
313 
314 LogicalResult GPUReturnOpConversion::matchAndRewrite(
315     gpu::ReturnOp returnOp, ArrayRef<Value> operands,
316     ConversionPatternRewriter &rewriter) const {
317   if (!operands.empty())
318     return failure();
319 
320   rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
321   return success();
322 }
323 
324 //===----------------------------------------------------------------------===//
325 // GPU To SPIRV Patterns.
326 //===----------------------------------------------------------------------===//
327 
328 namespace {
329 #include "GPUToSPIRV.cpp.inc"
330 }
331 
332 void mlir::populateGPUToSPIRVPatterns(MLIRContext *context,
333                                       SPIRVTypeConverter &typeConverter,
334                                       OwningRewritePatternList &patterns) {
335   populateWithGenerated(context, patterns);
336   patterns.insert<
337       GPUFuncOpConversion, GPUModuleConversion, GPUReturnOpConversion,
338       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
339       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
340       LaunchConfigConversion<gpu::ThreadIdOp,
341                              spirv::BuiltIn::LocalInvocationId>,
342       SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
343                                       spirv::BuiltIn::SubgroupId>,
344       SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
345                                       spirv::BuiltIn::NumSubgroups>,
346       SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
347                                       spirv::BuiltIn::SubgroupSize>,
348       WorkGroupSizeConversion>(typeConverter, context);
349 }
350