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