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/SPIRVEnums.h"
17 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
18 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
19 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
20 #include "mlir/IR/BuiltinOps.h"
21 #include "mlir/Transforms/DialectConversion.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, typename SourceOp::Adaptor adaptor,
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, typename SourceOp::Adaptor adaptor,
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, OpAdaptor adaptor,
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, OpAdaptor adaptor,
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, OpAdaptor adaptor,
85                   ConversionPatternRewriter &rewriter) const override;
86 };
87 
88 class GPUModuleEndConversion final
89     : public OpConversionPattern<gpu::ModuleEndOp> {
90 public:
91   using OpConversionPattern::OpConversionPattern;
92 
93   LogicalResult
94   matchAndRewrite(gpu::ModuleEndOp endOp, OpAdaptor adaptor,
95                   ConversionPatternRewriter &rewriter) const override {
96     rewriter.eraseOp(endOp);
97     return success();
98   }
99 };
100 
101 /// Pattern to convert a gpu.return into a SPIR-V return.
102 // TODO: This can go to DRR when GPU return has operands.
103 class GPUReturnOpConversion final : public OpConversionPattern<gpu::ReturnOp> {
104 public:
105   using OpConversionPattern<gpu::ReturnOp>::OpConversionPattern;
106 
107   LogicalResult
108   matchAndRewrite(gpu::ReturnOp returnOp, OpAdaptor adaptor,
109                   ConversionPatternRewriter &rewriter) const override;
110 };
111 
112 /// Pattern to convert a gpu.barrier op into a spv.ControlBarrier op.
113 class GPUBarrierConversion final : public OpConversionPattern<gpu::BarrierOp> {
114 public:
115   using OpConversionPattern::OpConversionPattern;
116 
117   LogicalResult
118   matchAndRewrite(gpu::BarrierOp barrierOp, OpAdaptor adaptor,
119                   ConversionPatternRewriter &rewriter) const override;
120 };
121 
122 } // namespace
123 
124 //===----------------------------------------------------------------------===//
125 // Builtins.
126 //===----------------------------------------------------------------------===//
127 
128 template <typename SourceOp, spirv::BuiltIn builtin>
129 LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
130     SourceOp op, typename SourceOp::Adaptor adaptor,
131     ConversionPatternRewriter &rewriter) const {
132   auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
133   auto indexType = typeConverter->getIndexType();
134 
135   // SPIR-V invocation builtin variables are a vector of type <3xi32>
136   auto spirvBuiltin =
137       spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
138   rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
139       op, indexType, spirvBuiltin,
140       rewriter.getI32ArrayAttr({static_cast<int32_t>(op.dimension())}));
141   return success();
142 }
143 
144 template <typename SourceOp, spirv::BuiltIn builtin>
145 LogicalResult
146 SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
147     SourceOp op, typename SourceOp::Adaptor adaptor,
148     ConversionPatternRewriter &rewriter) const {
149   auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
150   auto indexType = typeConverter->getIndexType();
151 
152   auto spirvBuiltin =
153       spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
154   rewriter.replaceOp(op, spirvBuiltin);
155   return success();
156 }
157 
158 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
159     gpu::BlockDimOp op, OpAdaptor adaptor,
160     ConversionPatternRewriter &rewriter) const {
161   auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
162   auto val = workGroupSizeAttr
163                  .getValues<int32_t>()[static_cast<int32_t>(op.dimension())];
164   auto convertedType =
165       getTypeConverter()->convertType(op.getResult().getType());
166   if (!convertedType)
167     return failure();
168   rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
169       op, convertedType, IntegerAttr::get(convertedType, val));
170   return success();
171 }
172 
173 //===----------------------------------------------------------------------===//
174 // GPUFuncOp
175 //===----------------------------------------------------------------------===//
176 
177 // Legalizes a GPU function as an entry SPIR-V function.
178 static spirv::FuncOp
179 lowerAsEntryFunction(gpu::GPUFuncOp funcOp, TypeConverter &typeConverter,
180                      ConversionPatternRewriter &rewriter,
181                      spirv::EntryPointABIAttr entryPointInfo,
182                      ArrayRef<spirv::InterfaceVarABIAttr> argABIInfo) {
183   auto fnType = funcOp.getType();
184   if (fnType.getNumResults()) {
185     funcOp.emitError("SPIR-V lowering only supports entry functions"
186                      "with no return values right now");
187     return nullptr;
188   }
189   if (!argABIInfo.empty() && fnType.getNumInputs() != argABIInfo.size()) {
190     funcOp.emitError(
191         "lowering as entry functions requires ABI info for all arguments "
192         "or none of them");
193     return nullptr;
194   }
195   // Update the signature to valid SPIR-V types and add the ABI
196   // attributes. These will be "materialized" by using the
197   // LowerABIAttributesPass.
198   TypeConverter::SignatureConversion signatureConverter(fnType.getNumInputs());
199   {
200     for (const auto &argType : enumerate(funcOp.getType().getInputs())) {
201       auto convertedType = typeConverter.convertType(argType.value());
202       signatureConverter.addInputs(argType.index(), convertedType);
203     }
204   }
205   auto newFuncOp = rewriter.create<spirv::FuncOp>(
206       funcOp.getLoc(), funcOp.getName(),
207       rewriter.getFunctionType(signatureConverter.getConvertedTypes(),
208                                llvm::None));
209   for (const auto &namedAttr : funcOp->getAttrs()) {
210     if (namedAttr.getName() == FunctionOpInterface::getTypeAttrName() ||
211         namedAttr.getName() == SymbolTable::getSymbolAttrName())
212       continue;
213     newFuncOp->setAttr(namedAttr.getName(), namedAttr.getValue());
214   }
215 
216   rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
217                               newFuncOp.end());
218   if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), typeConverter,
219                                          &signatureConverter)))
220     return nullptr;
221   rewriter.eraseOp(funcOp);
222 
223   // Set the attributes for argument and the function.
224   StringRef argABIAttrName = spirv::getInterfaceVarABIAttrName();
225   for (auto argIndex : llvm::seq<unsigned>(0, argABIInfo.size())) {
226     newFuncOp.setArgAttr(argIndex, argABIAttrName, argABIInfo[argIndex]);
227   }
228   newFuncOp->setAttr(spirv::getEntryPointABIAttrName(), entryPointInfo);
229 
230   return newFuncOp;
231 }
232 
233 /// Populates `argABI` with spv.interface_var_abi attributes for lowering
234 /// gpu.func to spv.func if no arguments have the attributes set
235 /// already. Returns failure if any argument has the ABI attribute set already.
236 static LogicalResult
237 getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
238                    SmallVectorImpl<spirv::InterfaceVarABIAttr> &argABI) {
239   spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(funcOp);
240   if (!spirv::needsInterfaceVarABIAttrs(targetEnv))
241     return success();
242 
243   for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
244     if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
245             argIndex, spirv::getInterfaceVarABIAttrName()))
246       return failure();
247     // Vulkan's interface variable requirements needs scalars to be wrapped in a
248     // struct. The struct held in storage buffer.
249     Optional<spirv::StorageClass> sc;
250     if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
251       sc = spirv::StorageClass::StorageBuffer;
252     argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
253   }
254   return success();
255 }
256 
257 LogicalResult GPUFuncOpConversion::matchAndRewrite(
258     gpu::GPUFuncOp funcOp, OpAdaptor adaptor,
259     ConversionPatternRewriter &rewriter) const {
260   if (!gpu::GPUDialect::isKernel(funcOp))
261     return failure();
262 
263   SmallVector<spirv::InterfaceVarABIAttr, 4> argABI;
264   if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
265     argABI.clear();
266     for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
267       // If the ABI is already specified, use it.
268       auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
269           argIndex, spirv::getInterfaceVarABIAttrName());
270       if (!abiAttr) {
271         funcOp.emitRemark(
272             "match failure: missing 'spv.interface_var_abi' attribute at "
273             "argument ")
274             << argIndex;
275         return failure();
276       }
277       argABI.push_back(abiAttr);
278     }
279   }
280 
281   auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
282   if (!entryPointAttr) {
283     funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
284     return failure();
285   }
286   spirv::FuncOp newFuncOp = lowerAsEntryFunction(
287       funcOp, *getTypeConverter(), rewriter, entryPointAttr, argABI);
288   if (!newFuncOp)
289     return failure();
290   newFuncOp->removeAttr(
291       rewriter.getStringAttr(gpu::GPUDialect::getKernelFuncAttrName()));
292   return success();
293 }
294 
295 //===----------------------------------------------------------------------===//
296 // ModuleOp with gpu.module.
297 //===----------------------------------------------------------------------===//
298 
299 LogicalResult GPUModuleConversion::matchAndRewrite(
300     gpu::GPUModuleOp moduleOp, OpAdaptor adaptor,
301     ConversionPatternRewriter &rewriter) const {
302   spirv::TargetEnvAttr targetEnv = spirv::lookupTargetEnvOrDefault(moduleOp);
303   spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv);
304   FailureOr<spirv::MemoryModel> memoryModel = spirv::getMemoryModel(targetEnv);
305   if (failed(memoryModel))
306     return moduleOp.emitRemark("match failure: could not selected memory model "
307                                "based on 'spv.target_env'");
308 
309   // Add a keyword to the module name to avoid symbolic conflict.
310   std::string spvModuleName = (kSPIRVModule + moduleOp.getName()).str();
311   auto spvModule = rewriter.create<spirv::ModuleOp>(
312       moduleOp.getLoc(), addressingModel, memoryModel.getValue(), llvm::None,
313       StringRef(spvModuleName));
314 
315   // Move the region from the module op into the SPIR-V module.
316   Region &spvModuleRegion = spvModule.getRegion();
317   rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
318                               spvModuleRegion.begin());
319   // The spv.module build method adds a block. Remove that.
320   rewriter.eraseBlock(&spvModuleRegion.back());
321   rewriter.eraseOp(moduleOp);
322   return success();
323 }
324 
325 //===----------------------------------------------------------------------===//
326 // GPU return inside kernel functions to SPIR-V return.
327 //===----------------------------------------------------------------------===//
328 
329 LogicalResult GPUReturnOpConversion::matchAndRewrite(
330     gpu::ReturnOp returnOp, OpAdaptor adaptor,
331     ConversionPatternRewriter &rewriter) const {
332   if (!adaptor.getOperands().empty())
333     return failure();
334 
335   rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
336   return success();
337 }
338 
339 //===----------------------------------------------------------------------===//
340 // Barrier.
341 //===----------------------------------------------------------------------===//
342 
343 LogicalResult GPUBarrierConversion::matchAndRewrite(
344     gpu::BarrierOp barrierOp, OpAdaptor adaptor,
345     ConversionPatternRewriter &rewriter) const {
346   MLIRContext *context = getContext();
347   // Both execution and memory scope should be workgroup.
348   auto scope = spirv::ScopeAttr::get(context, spirv::Scope::Workgroup);
349   // Require acquire and release memory semantics for workgroup memory.
350   auto memorySemantics = spirv::MemorySemanticsAttr::get(
351       context, spirv::MemorySemantics::WorkgroupMemory |
352                    spirv::MemorySemantics::AcquireRelease);
353   rewriter.replaceOpWithNewOp<spirv::ControlBarrierOp>(barrierOp, scope, scope,
354                                                        memorySemantics);
355   return success();
356 }
357 
358 //===----------------------------------------------------------------------===//
359 // GPU To SPIRV Patterns.
360 //===----------------------------------------------------------------------===//
361 
362 void mlir::populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
363                                       RewritePatternSet &patterns) {
364   patterns.add<
365       GPUBarrierConversion, GPUFuncOpConversion, GPUModuleConversion,
366       GPUModuleEndConversion, GPUReturnOpConversion,
367       LaunchConfigConversion<gpu::BlockIdOp, spirv::BuiltIn::WorkgroupId>,
368       LaunchConfigConversion<gpu::GridDimOp, spirv::BuiltIn::NumWorkgroups>,
369       LaunchConfigConversion<gpu::ThreadIdOp,
370                              spirv::BuiltIn::LocalInvocationId>,
371       SingleDimLaunchConfigConversion<gpu::SubgroupIdOp,
372                                       spirv::BuiltIn::SubgroupId>,
373       SingleDimLaunchConfigConversion<gpu::NumSubgroupsOp,
374                                       spirv::BuiltIn::NumSubgroups>,
375       SingleDimLaunchConfigConversion<gpu::SubgroupSizeOp,
376                                       spirv::BuiltIn::SubgroupSize>,
377       WorkGroupSizeConversion>(typeConverter, patterns.getContext());
378 }
379