MLIR  14.0.0git
GPUToSPIRV.cpp
Go to the documentation of this file.
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 
19 #include "mlir/IR/BuiltinOps.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:
34 
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:
46 
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:
59 
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:
69 
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:
82 
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:
92 
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:
106 
108  matchAndRewrite(gpu::ReturnOp returnOp, OpAdaptor adaptor,
109  ConversionPatternRewriter &rewriter) const override;
110 };
111 
112 } // namespace
113 
114 //===----------------------------------------------------------------------===//
115 // Builtins.
116 //===----------------------------------------------------------------------===//
117 
118 template <typename SourceOp, spirv::BuiltIn builtin>
119 LogicalResult LaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
120  SourceOp op, typename SourceOp::Adaptor adaptor,
121  ConversionPatternRewriter &rewriter) const {
122  auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
123  auto indexType = typeConverter->getIndexType();
124 
125  // SPIR-V invocation builtin variables are a vector of type <3xi32>
126  auto spirvBuiltin =
127  spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
128  rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
129  op, indexType, spirvBuiltin,
130  rewriter.getI32ArrayAttr({static_cast<int32_t>(op.dimension())}));
131  return success();
132 }
133 
134 template <typename SourceOp, spirv::BuiltIn builtin>
136 SingleDimLaunchConfigConversion<SourceOp, builtin>::matchAndRewrite(
137  SourceOp op, typename SourceOp::Adaptor adaptor,
138  ConversionPatternRewriter &rewriter) const {
139  auto *typeConverter = this->template getTypeConverter<SPIRVTypeConverter>();
140  auto indexType = typeConverter->getIndexType();
141 
142  auto spirvBuiltin =
143  spirv::getBuiltinVariableValue(op, builtin, indexType, rewriter);
144  rewriter.replaceOp(op, spirvBuiltin);
145  return success();
146 }
147 
148 LogicalResult WorkGroupSizeConversion::matchAndRewrite(
149  gpu::BlockDimOp op, OpAdaptor adaptor,
150  ConversionPatternRewriter &rewriter) const {
151  auto workGroupSizeAttr = spirv::lookupLocalWorkGroupSize(op);
152  auto val = workGroupSizeAttr
153  .getValues<int32_t>()[static_cast<int32_t>(op.dimension())];
154  auto convertedType =
155  getTypeConverter()->convertType(op.getResult().getType());
156  if (!convertedType)
157  return failure();
158  rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
159  op, convertedType, IntegerAttr::get(convertedType, val));
160  return success();
161 }
162 
163 //===----------------------------------------------------------------------===//
164 // GPUFuncOp
165 //===----------------------------------------------------------------------===//
166 
167 // Legalizes a GPU function as an entry SPIR-V function.
168 static spirv::FuncOp
169 lowerAsEntryFunction(gpu::GPUFuncOp funcOp, TypeConverter &typeConverter,
170  ConversionPatternRewriter &rewriter,
171  spirv::EntryPointABIAttr entryPointInfo,
173  auto fnType = funcOp.getType();
174  if (fnType.getNumResults()) {
175  funcOp.emitError("SPIR-V lowering only supports entry functions"
176  "with no return values right now");
177  return nullptr;
178  }
179  if (!argABIInfo.empty() && fnType.getNumInputs() != argABIInfo.size()) {
180  funcOp.emitError(
181  "lowering as entry functions requires ABI info for all arguments "
182  "or none of them");
183  return nullptr;
184  }
185  // Update the signature to valid SPIR-V types and add the ABI
186  // attributes. These will be "materialized" by using the
187  // LowerABIAttributesPass.
188  TypeConverter::SignatureConversion signatureConverter(fnType.getNumInputs());
189  {
190  for (const auto &argType : enumerate(funcOp.getType().getInputs())) {
191  auto convertedType = typeConverter.convertType(argType.value());
192  signatureConverter.addInputs(argType.index(), convertedType);
193  }
194  }
195  auto newFuncOp = rewriter.create<spirv::FuncOp>(
196  funcOp.getLoc(), funcOp.getName(),
197  rewriter.getFunctionType(signatureConverter.getConvertedTypes(),
198  llvm::None));
199  for (const auto &namedAttr : funcOp->getAttrs()) {
200  if (namedAttr.getName() == FunctionOpInterface::getTypeAttrName() ||
201  namedAttr.getName() == SymbolTable::getSymbolAttrName())
202  continue;
203  newFuncOp->setAttr(namedAttr.getName(), namedAttr.getValue());
204  }
205 
206  rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(),
207  newFuncOp.end());
208  if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), typeConverter,
209  &signatureConverter)))
210  return nullptr;
211  rewriter.eraseOp(funcOp);
212 
213  // Set the attributes for argument and the function.
214  StringRef argABIAttrName = spirv::getInterfaceVarABIAttrName();
215  for (auto argIndex : llvm::seq<unsigned>(0, argABIInfo.size())) {
216  newFuncOp.setArgAttr(argIndex, argABIAttrName, argABIInfo[argIndex]);
217  }
218  newFuncOp->setAttr(spirv::getEntryPointABIAttrName(), entryPointInfo);
219 
220  return newFuncOp;
221 }
222 
223 /// Populates `argABI` with spv.interface_var_abi attributes for lowering
224 /// gpu.func to spv.func if no arguments have the attributes set
225 /// already. Returns failure if any argument has the ABI attribute set already.
226 static LogicalResult
227 getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp,
230  if (!spirv::needsInterfaceVarABIAttrs(targetEnv))
231  return success();
232 
233  for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
234  if (funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
236  return failure();
237  // Vulkan's interface variable requirements needs scalars to be wrapped in a
238  // struct. The struct held in storage buffer.
240  if (funcOp.getArgument(argIndex).getType().isIntOrIndexOrFloat())
241  sc = spirv::StorageClass::StorageBuffer;
242  argABI.push_back(spirv::getInterfaceVarABIAttr(0, argIndex, sc, context));
243  }
244  return success();
245 }
246 
247 LogicalResult GPUFuncOpConversion::matchAndRewrite(
248  gpu::GPUFuncOp funcOp, OpAdaptor adaptor,
249  ConversionPatternRewriter &rewriter) const {
250  if (!gpu::GPUDialect::isKernel(funcOp))
251  return failure();
252 
254  if (failed(getDefaultABIAttrs(rewriter.getContext(), funcOp, argABI))) {
255  argABI.clear();
256  for (auto argIndex : llvm::seq<unsigned>(0, funcOp.getNumArguments())) {
257  // If the ABI is already specified, use it.
258  auto abiAttr = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
260  if (!abiAttr) {
261  funcOp.emitRemark(
262  "match failure: missing 'spv.interface_var_abi' attribute at "
263  "argument ")
264  << argIndex;
265  return failure();
266  }
267  argABI.push_back(abiAttr);
268  }
269  }
270 
271  auto entryPointAttr = spirv::lookupEntryPointABI(funcOp);
272  if (!entryPointAttr) {
273  funcOp.emitRemark("match failure: missing 'spv.entry_point_abi' attribute");
274  return failure();
275  }
276  spirv::FuncOp newFuncOp = lowerAsEntryFunction(
277  funcOp, *getTypeConverter(), rewriter, entryPointAttr, argABI);
278  if (!newFuncOp)
279  return failure();
280  newFuncOp->removeAttr(
281  rewriter.getStringAttr(gpu::GPUDialect::getKernelFuncAttrName()));
282  return success();
283 }
284 
285 //===----------------------------------------------------------------------===//
286 // ModuleOp with gpu.module.
287 //===----------------------------------------------------------------------===//
288 
289 LogicalResult GPUModuleConversion::matchAndRewrite(
290  gpu::GPUModuleOp moduleOp, OpAdaptor adaptor,
291  ConversionPatternRewriter &rewriter) const {
293  spirv::AddressingModel addressingModel = spirv::getAddressingModel(targetEnv);
294  FailureOr<spirv::MemoryModel> memoryModel = spirv::getMemoryModel(targetEnv);
295  if (failed(memoryModel))
296  return moduleOp.emitRemark("match failure: could not selected memory model "
297  "based on 'spv.target_env'");
298 
299  // Add a keyword to the module name to avoid symbolic conflict.
300  std::string spvModuleName = (kSPIRVModule + moduleOp.getName()).str();
301  auto spvModule = rewriter.create<spirv::ModuleOp>(
302  moduleOp.getLoc(), addressingModel, memoryModel.getValue(), llvm::None,
303  StringRef(spvModuleName));
304 
305  // Move the region from the module op into the SPIR-V module.
306  Region &spvModuleRegion = spvModule.getRegion();
307  rewriter.inlineRegionBefore(moduleOp.body(), spvModuleRegion,
308  spvModuleRegion.begin());
309  // The spv.module build method adds a block. Remove that.
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, OpAdaptor adaptor,
321  ConversionPatternRewriter &rewriter) const {
322  if (!adaptor.getOperands().empty())
323  return failure();
324 
325  rewriter.replaceOpWithNewOp<spirv::ReturnOp>(returnOp);
326  return success();
327 }
328 
329 //===----------------------------------------------------------------------===//
330 // GPU To SPIRV Patterns.
331 //===----------------------------------------------------------------------===//
332 
334  RewritePatternSet &patterns) {
335  patterns.add<
336  GPUFuncOpConversion, GPUModuleConversion, GPUModuleEndConversion,
337  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, patterns.getContext());
349 }
TargetEnvAttr lookupTargetEnvOrDefault(Operation *op)
Queries the target environment recursively from enclosing symbol table ops containing the given op or...
Include the generated interface declarations.
DenseIntElementsAttr lookupLocalWorkGroupSize(Operation *op)
Queries the local workgroup size from entry point ABI on the nearest function-like op containing the ...
OpTy create(Location location, Args &&...args)
Create an operation of specific op type at the current insertion point.
Definition: Builders.h:430
static LogicalResult getDefaultABIAttrs(MLIRContext *context, gpu::GPUFuncOp funcOp, SmallVectorImpl< spirv::InterfaceVarABIAttr > &argABI)
Populates argABI with spv.interface_var_abi attributes for lowering gpu.func to spv.func if no arguments have the attributes set already.
Definition: GPUToSPIRV.cpp:227
This class contains a list of basic blocks and a link to the parent operation it is attached to...
Definition: Region.h:26
static StringRef getSymbolAttrName()
Return the name of the attribute used for symbol names.
Definition: SymbolTable.h:55
MLIRContext * getContext() const
Definition: Builders.h:54
FailureOr< MemoryModel > getMemoryModel(TargetEnvAttr targetAttr)
Returns memory model selected based on target environment.
LogicalResult convertType(Type t, SmallVectorImpl< Type > &results)
Convert the given type.
An attribute that specifies the information regarding the interface variable: descriptor set...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
StringRef getEntryPointABIAttrName()
Returns the attribute name for specifying entry point information.
StringRef getTypeAttrName()
Return the name of the attribute used for function types.
bool needsInterfaceVarABIAttrs(TargetEnvAttr targetAttr)
Returns whether the given SPIR-V target (described by TargetEnvAttr) needs ABI attributes for interfa...
static spirv::FuncOp lowerAsEntryFunction(gpu::GPUFuncOp funcOp, TypeConverter &typeConverter, ConversionPatternRewriter &rewriter, spirv::EntryPointABIAttr entryPointInfo, ArrayRef< spirv::InterfaceVarABIAttr > argABIInfo)
Definition: GPUToSPIRV.cpp:169
EntryPointABIAttr lookupEntryPointABI(Operation *op)
Queries the entry point ABI on the nearest function-like op containing the given op.
void replaceOp(Operation *op, ValueRange newValues) override
PatternRewriter hook for replacing the results of an operation.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
This class provides all of the information necessary to convert a type signature. ...
This class provides support for representing a failure result, or a valid value of type T...
Definition: LogicalResult.h:77
Value getBuiltinVariableValue(Operation *op, BuiltIn builtin, Type integerType, OpBuilder &builder)
Returns the value for the given builtin variable.
void eraseBlock(Block *block) override
PatternRewriter hook for erase all operations in a block.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:206
StringRef getInterfaceVarABIAttrName()
Returns the attribute name for specifying argument ABI information.
ArrayAttr getI32ArrayAttr(ArrayRef< int32_t > values)
Definition: Builders.cpp:215
OpConversionPattern(MLIRContext *context, PatternBenefit benefit=1)
InterfaceVarABIAttr getInterfaceVarABIAttr(unsigned descriptorSet, unsigned binding, Optional< StorageClass > storageClass, MLIRContext *context)
Gets the InterfaceVarABIAttr given its fields.
AddressingModel getAddressingModel(TargetEnvAttr targetAttr)
Returns addressing model selected based on target environment.
void inlineRegionBefore(Region &region, Region &parent, Region::iterator before) override
PatternRewriter hook for moving blocks out of a region.
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
static constexpr const char kSPIRVModule[]
Definition: GPUToSPIRV.cpp:25
OpTy replaceOpWithNewOp(Operation *op, Args &&... args)
Replaces the result op with a new op that is created without verification.
Definition: PatternMatch.h:741
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
Definition: PatternMatch.h:930
Type conversion class.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
void populateGPUToSPIRVPatterns(SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns)
Appends to a pattern list additional patterns for translating GPU Ops to SPIR-V ops.
Definition: GPUToSPIRV.cpp:333
This class implements a pattern rewriter for use with ConversionPatterns.
void eraseOp(Operation *op) override
PatternRewriter hook for erasing a dead operation.
FunctionType getFunctionType(TypeRange inputs, TypeRange results)
Definition: Builders.cpp:67
StringAttr getStringAttr(const Twine &bytes)
Definition: Builders.cpp:201
MLIRContext * getContext() const
Definition: PatternMatch.h:906
An attribute that specifies the target version, allowed extensions and capabilities, and resource limits.
Type conversion from builtin types to SPIR-V types for shader interface.
FailureOr< Block * > convertRegionTypes(Region *region, TypeConverter &converter, TypeConverter::SignatureConversion *entryConversion=nullptr)
Convert the types of block arguments within the given region.