MLIR  18.0.0git
LowerABIAttributesPass.cpp
Go to the documentation of this file.
1 //===- LowerABIAttributesPass.cpp - Decorate composite type ---------------===//
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 a pass to lower attributes that specify the shader ABI
10 // for the functions in the generated SPIR-V module.
11 //
12 //===----------------------------------------------------------------------===//
13 
15 
25 #include "llvm/ADT/SetVector.h"
26 
27 namespace mlir {
28 namespace spirv {
29 #define GEN_PASS_DEF_SPIRVLOWERABIATTRIBUTESPASS
30 #include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
31 } // namespace spirv
32 } // namespace mlir
33 
34 using namespace mlir;
35 
36 /// Creates a global variable for an argument based on the ABI info.
37 static spirv::GlobalVariableOp
38 createGlobalVarForEntryPointArgument(OpBuilder &builder, spirv::FuncOp funcOp,
39  unsigned argIndex,
41  auto spirvModule = funcOp->getParentOfType<spirv::ModuleOp>();
42  if (!spirvModule)
43  return nullptr;
44 
45  OpBuilder::InsertionGuard moduleInsertionGuard(builder);
46  builder.setInsertionPoint(funcOp.getOperation());
47  std::string varName =
48  funcOp.getName().str() + "_arg_" + std::to_string(argIndex);
49 
50  // Get the type of variable. If this is a scalar/vector type and has an ABI
51  // info create a variable of type !spirv.ptr<!spirv.struct<elementType>>. If
52  // not it must already be a !spirv.ptr<!spirv.struct<...>>.
53  auto varType = funcOp.getFunctionType().getInput(argIndex);
54  if (cast<spirv::SPIRVType>(varType).isScalarOrVector()) {
55  auto storageClass = abiInfo.getStorageClass();
56  if (!storageClass)
57  return nullptr;
58  varType =
59  spirv::PointerType::get(spirv::StructType::get(varType), *storageClass);
60  }
61  auto varPtrType = cast<spirv::PointerType>(varType);
62  auto varPointeeType = cast<spirv::StructType>(varPtrType.getPointeeType());
63 
64  // Set the offset information.
65  varPointeeType =
66  cast<spirv::StructType>(VulkanLayoutUtils::decorateType(varPointeeType));
67 
68  if (!varPointeeType)
69  return nullptr;
70 
71  varType =
72  spirv::PointerType::get(varPointeeType, varPtrType.getStorageClass());
73 
74  return builder.create<spirv::GlobalVariableOp>(
75  funcOp.getLoc(), varType, varName, abiInfo.getDescriptorSet(),
76  abiInfo.getBinding());
77 }
78 
79 /// Gets the global variables that need to be specified as interface variable
80 /// with an spirv.EntryPointOp. Traverses the body of a entry function to do so.
81 static LogicalResult
82 getInterfaceVariables(spirv::FuncOp funcOp,
83  SmallVectorImpl<Attribute> &interfaceVars) {
84  auto module = funcOp->getParentOfType<spirv::ModuleOp>();
85  if (!module) {
86  return failure();
87  }
88  SetVector<Operation *> interfaceVarSet;
89 
90  // TODO: This should in reality traverse the entry function
91  // call graph and collect all the interfaces. For now, just traverse the
92  // instructions in this function.
93  funcOp.walk([&](spirv::AddressOfOp addressOfOp) {
94  auto var =
95  module.lookupSymbol<spirv::GlobalVariableOp>(addressOfOp.getVariable());
96  // TODO: Per SPIR-V spec: "Before version 1.4, the interface’s
97  // storage classes are limited to the Input and Output storage classes.
98  // Starting with version 1.4, the interface’s storage classes are all
99  // storage classes used in declaring all global variables referenced by the
100  // entry point’s call tree." We should consider the target environment here.
101  switch (cast<spirv::PointerType>(var.getType()).getStorageClass()) {
102  case spirv::StorageClass::Input:
103  case spirv::StorageClass::Output:
104  interfaceVarSet.insert(var.getOperation());
105  break;
106  default:
107  break;
108  }
109  });
110  for (auto &var : interfaceVarSet) {
111  interfaceVars.push_back(SymbolRefAttr::get(
112  funcOp.getContext(), cast<spirv::GlobalVariableOp>(var).getSymName()));
113  }
114  return success();
115 }
116 
117 /// Lowers the entry point attribute.
118 static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
119  OpBuilder &builder) {
120  auto entryPointAttrName = spirv::getEntryPointABIAttrName();
121  auto entryPointAttr =
122  funcOp->getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName);
123  if (!entryPointAttr) {
124  return failure();
125  }
126 
127  OpBuilder::InsertionGuard moduleInsertionGuard(builder);
128  auto spirvModule = funcOp->getParentOfType<spirv::ModuleOp>();
129  builder.setInsertionPointToEnd(spirvModule.getBody());
130 
131  // Adds the spirv.EntryPointOp after collecting all the interface variables
132  // needed.
133  SmallVector<Attribute, 1> interfaceVars;
134  if (failed(getInterfaceVariables(funcOp, interfaceVars))) {
135  return failure();
136  }
137 
138  spirv::TargetEnvAttr targetEnvAttr = spirv::lookupTargetEnv(funcOp);
139  spirv::TargetEnv targetEnv(targetEnvAttr);
140  FailureOr<spirv::ExecutionModel> executionModel =
141  spirv::getExecutionModel(targetEnvAttr);
142  if (failed(executionModel))
143  return funcOp.emitRemark("lower entry point failure: could not select "
144  "execution model based on 'spirv.target_env'");
145 
146  builder.create<spirv::EntryPointOp>(funcOp.getLoc(), *executionModel, funcOp,
147  interfaceVars);
148 
149  // Specifies the spirv.ExecutionModeOp.
150  if (DenseI32ArrayAttr workgroupSizeAttr = entryPointAttr.getWorkgroupSize()) {
151  std::optional<ArrayRef<spirv::Capability>> caps =
152  spirv::getCapabilities(spirv::ExecutionMode::LocalSize);
153  if (!caps || targetEnv.allows(*caps)) {
154  builder.create<spirv::ExecutionModeOp>(funcOp.getLoc(), funcOp,
155  spirv::ExecutionMode::LocalSize,
156  workgroupSizeAttr.asArrayRef());
157  // Erase workgroup size.
158  entryPointAttr = spirv::EntryPointABIAttr::get(
159  entryPointAttr.getContext(), DenseI32ArrayAttr(),
160  entryPointAttr.getSubgroupSize());
161  }
162  }
163  if (std::optional<int> subgroupSize = entryPointAttr.getSubgroupSize()) {
164  std::optional<ArrayRef<spirv::Capability>> caps =
165  spirv::getCapabilities(spirv::ExecutionMode::SubgroupSize);
166  if (!caps || targetEnv.allows(*caps)) {
167  builder.create<spirv::ExecutionModeOp>(funcOp.getLoc(), funcOp,
168  spirv::ExecutionMode::SubgroupSize,
169  *subgroupSize);
170  // Erase subgroup size.
171  entryPointAttr = spirv::EntryPointABIAttr::get(
172  entryPointAttr.getContext(), entryPointAttr.getWorkgroupSize(),
173  std::nullopt);
174  }
175  }
176  if (entryPointAttr.getWorkgroupSize() || entryPointAttr.getSubgroupSize())
177  funcOp->setAttr(entryPointAttrName, entryPointAttr);
178  else
179  funcOp->removeAttr(entryPointAttrName);
180  return success();
181 }
182 
183 namespace {
184 /// A pattern to convert function signature according to interface variable ABI
185 /// attributes.
186 ///
187 /// Specifically, this pattern creates global variables according to interface
188 /// variable ABI attributes attached to function arguments and converts all
189 /// function argument uses to those global variables. This is necessary because
190 /// Vulkan requires all shader entry points to be of void(void) type.
191 class ProcessInterfaceVarABI final : public OpConversionPattern<spirv::FuncOp> {
192 public:
194 
196  matchAndRewrite(spirv::FuncOp funcOp, OpAdaptor adaptor,
197  ConversionPatternRewriter &rewriter) const override;
198 };
199 
200 /// Pass to implement the ABI information specified as attributes.
201 class LowerABIAttributesPass final
202  : public spirv::impl::SPIRVLowerABIAttributesPassBase<
203  LowerABIAttributesPass> {
204  void runOnOperation() override;
205 };
206 } // namespace
207 
208 LogicalResult ProcessInterfaceVarABI::matchAndRewrite(
209  spirv::FuncOp funcOp, OpAdaptor adaptor,
210  ConversionPatternRewriter &rewriter) const {
211  if (!funcOp->getAttrOfType<spirv::EntryPointABIAttr>(
213  // TODO: Non-entry point functions are not handled.
214  return failure();
215  }
216  TypeConverter::SignatureConversion signatureConverter(
217  funcOp.getFunctionType().getNumInputs());
218 
219  auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
220  auto indexType = typeConverter.getIndexType();
221 
222  auto attrName = spirv::getInterfaceVarABIAttrName();
223  for (const auto &argType :
224  llvm::enumerate(funcOp.getFunctionType().getInputs())) {
225  auto abiInfo = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
226  argType.index(), attrName);
227  if (!abiInfo) {
228  // TODO: For non-entry point functions, it should be legal
229  // to pass around scalar/vector values and return a scalar/vector. For now
230  // non-entry point functions are not handled in this ABI lowering and will
231  // produce an error.
232  return failure();
233  }
234  spirv::GlobalVariableOp var = createGlobalVarForEntryPointArgument(
235  rewriter, funcOp, argType.index(), abiInfo);
236  if (!var)
237  return failure();
238 
239  OpBuilder::InsertionGuard funcInsertionGuard(rewriter);
240  rewriter.setInsertionPointToStart(&funcOp.front());
241  // Insert spirv::AddressOf and spirv::AccessChain operations.
242  Value replacement =
243  rewriter.create<spirv::AddressOfOp>(funcOp.getLoc(), var);
244  // Check if the arg is a scalar or vector type. In that case, the value
245  // needs to be loaded into registers.
246  // TODO: This is loading value of the scalar into registers
247  // at the start of the function. It is probably better to do the load just
248  // before the use. There might be multiple loads and currently there is no
249  // easy way to replace all uses with a sequence of operations.
250  if (cast<spirv::SPIRVType>(argType.value()).isScalarOrVector()) {
251  auto zero =
252  spirv::ConstantOp::getZero(indexType, funcOp.getLoc(), rewriter);
253  auto loadPtr = rewriter.create<spirv::AccessChainOp>(
254  funcOp.getLoc(), replacement, zero.getConstant());
255  replacement = rewriter.create<spirv::LoadOp>(funcOp.getLoc(), loadPtr);
256  }
257  signatureConverter.remapInput(argType.index(), replacement);
258  }
259  if (failed(rewriter.convertRegionTypes(&funcOp.getBody(), *getTypeConverter(),
260  &signatureConverter)))
261  return failure();
262 
263  // Creates a new function with the update signature.
264  rewriter.updateRootInPlace(funcOp, [&] {
265  funcOp.setType(rewriter.getFunctionType(
266  signatureConverter.getConvertedTypes(), std::nullopt));
267  });
268  return success();
269 }
270 
271 void LowerABIAttributesPass::runOnOperation() {
272  // Uses the signature conversion methodology of the dialect conversion
273  // framework to implement the conversion.
274  spirv::ModuleOp module = getOperation();
275  MLIRContext *context = &getContext();
276 
277  spirv::TargetEnvAttr targetEnvAttr = spirv::lookupTargetEnv(module);
278  if (!targetEnvAttr) {
279  module->emitOpError("missing SPIR-V target env attribute");
280  return signalPassFailure();
281  }
282  spirv::TargetEnv targetEnv(targetEnvAttr);
283 
284  SPIRVTypeConverter typeConverter(targetEnv);
285 
286  // Insert a bitcast in the case of a pointer type change.
287  typeConverter.addSourceMaterialization([](OpBuilder &builder,
288  spirv::PointerType type,
289  ValueRange inputs, Location loc) {
290  if (inputs.size() != 1 || !isa<spirv::PointerType>(inputs[0].getType()))
291  return Value();
292  return builder.create<spirv::BitcastOp>(loc, type, inputs[0]).getResult();
293  });
294 
295  RewritePatternSet patterns(context);
296  patterns.add<ProcessInterfaceVarABI>(typeConverter, context);
297 
298  ConversionTarget target(*context);
299  // "Legal" function ops should have no interface variable ABI attributes.
300  target.addDynamicallyLegalOp<spirv::FuncOp>([&](spirv::FuncOp op) {
301  StringRef attrName = spirv::getInterfaceVarABIAttrName();
302  for (unsigned i = 0, e = op.getNumArguments(); i < e; ++i)
303  if (op.getArgAttr(i, attrName))
304  return false;
305  return true;
306  });
307  // All other SPIR-V ops are legal.
308  target.markUnknownOpDynamicallyLegal([](Operation *op) {
309  return op->getDialect()->getNamespace() ==
310  spirv::SPIRVDialect::getDialectNamespace();
311  });
312  if (failed(applyPartialConversion(module, target, std::move(patterns))))
313  return signalPassFailure();
314 
315  // Walks over all the FuncOps in spirv::ModuleOp to lower the entry point
316  // attributes.
317  OpBuilder builder(context);
318  SmallVector<spirv::FuncOp, 1> entryPointFns;
319  auto entryPointAttrName = spirv::getEntryPointABIAttrName();
320  module.walk([&](spirv::FuncOp funcOp) {
321  if (funcOp->getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName)) {
322  entryPointFns.push_back(funcOp);
323  }
324  });
325  for (auto fn : entryPointFns) {
326  if (failed(lowerEntryPointABIAttr(fn, builder))) {
327  return signalPassFailure();
328  }
329  }
330 }
static Value getZero(OpBuilder &b, Location loc, Type elementType)
Get zero value for an element type.
static MLIRContext * getContext(OpFoldResult val)
static spirv::GlobalVariableOp createGlobalVarForEntryPointArgument(OpBuilder &builder, spirv::FuncOp funcOp, unsigned argIndex, spirv::InterfaceVarABIAttr abiInfo)
Creates a global variable for an argument based on the ABI info.
static LogicalResult getInterfaceVariables(spirv::FuncOp funcOp, SmallVectorImpl< Attribute > &interfaceVars)
Gets the global variables that need to be specified as interface variable with an spirv....
static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp, OpBuilder &builder)
Lowers the entry point attribute.
FunctionType getFunctionType(TypeRange inputs, TypeRange results)
Definition: Builders.cpp:96
This class implements a pattern rewriter for use with ConversionPatterns.
FailureOr< Block * > convertRegionTypes(Region *region, const TypeConverter &converter, TypeConverter::SignatureConversion *entryConversion=nullptr)
Convert the types of block arguments within the given region.
This class describes a specific conversion target.
StringRef getNamespace() const
Definition: Dialect.h:57
This class provides support for representing a failure result, or a valid value of type T.
Definition: LogicalResult.h:78
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:63
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:333
This class helps build Operations.
Definition: Builders.h:206
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:416
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:383
void setInsertionPointToEnd(Block *block)
Sets the insertion point to the end of the specified block.
Definition: Builders.h:421
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:446
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Dialect * getDialect()
Return the dialect this operation is associated with, or nullptr if the associated dialect is not loa...
Definition: Operation.h:220
void updateRootInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around a root update of an operation.
Definition: PatternMatch.h:606
Type conversion from builtin types to SPIR-V types for shader interface.
This class provides all of the information necessary to convert a type signature.
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:378
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
static spirv::StructType decorateType(spirv::StructType structType)
Returns a new StructType with layout decoration.
Definition: LayoutUtils.cpp:21
An attribute that specifies the information regarding the interface variable: descriptor set,...
uint32_t getBinding()
Returns binding.
uint32_t getDescriptorSet()
Returns descriptor set.
std::optional< StorageClass > getStorageClass()
Returns spirv::StorageClass.
static PointerType get(Type pointeeType, StorageClass storageClass)
Definition: SPIRVTypes.cpp:546
static StructType get(ArrayRef< Type > memberTypes, ArrayRef< OffsetInfo > offsetInfo={}, ArrayRef< MemberDecorationInfo > memberDecorations={})
Construct a literal StructType with at least one member.
An attribute that specifies the target version, allowed extensions and capabilities,...
A wrapper class around a spirv::TargetEnvAttr to provide query methods for allowed version/capabiliti...
Definition: TargetAndABI.h:29
bool allows(Capability) const
Returns true if the given capability is allowed.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
StringRef getInterfaceVarABIAttrName()
Returns the attribute name for specifying argument ABI information.
TargetEnvAttr lookupTargetEnv(Operation *op)
Queries the target environment recursively from enclosing symbol table ops containing the given op.
FailureOr< ExecutionModel > getExecutionModel(TargetEnvAttr targetAttr)
Returns execution model selected based on target environment.
StringRef getEntryPointABIAttrName()
Returns the attribute name for specifying entry point information.
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult applyPartialConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, DenseSet< Operation * > *unconvertedOps=nullptr)
Below we define several entry points for operation conversion.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
Definition: LogicalResult.h:72
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26