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 
14 #include "mlir/Dialect/SPIRV/LayoutUtils.h"
15 #include "mlir/Dialect/SPIRV/Passes.h"
16 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
17 #include "mlir/Dialect/SPIRV/SPIRVLowering.h"
18 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
19 #include "mlir/Transforms/DialectConversion.h"
20 #include "llvm/ADT/SetVector.h"
21 
22 using namespace mlir;
23 
24 /// Checks if the `type` is a scalar or vector type. It is assumed that they are
25 /// valid for SPIR-V dialect already.
26 static bool isScalarOrVectorType(Type type) {
27   return spirv::SPIRVDialect::isValidScalarType(type) || type.isa<VectorType>();
28 }
29 
30 /// Creates a global variable for an argument based on the ABI info.
31 static spirv::GlobalVariableOp
32 createGlobalVariableForArg(spirv::FuncOp funcOp, OpBuilder &builder,
33                            unsigned argNum,
34                            spirv::InterfaceVarABIAttr abiInfo) {
35   auto spirvModule = funcOp.getParentOfType<spirv::ModuleOp>();
36   if (!spirvModule) {
37     return nullptr;
38   }
39   OpBuilder::InsertionGuard moduleInsertionGuard(builder);
40   builder.setInsertionPoint(funcOp.getOperation());
41   std::string varName =
42       funcOp.getName().str() + "_arg_" + std::to_string(argNum);
43 
44   // Get the type of variable. If this is a scalar/vector type and has an ABI
45   // info create a variable of type !spv.ptr<!spv.struct<elementTYpe>>. If not
46   // it must already be a !spv.ptr<!spv.struct<...>>.
47   auto varType = funcOp.getType().getInput(argNum);
48   auto storageClass =
49       static_cast<spirv::StorageClass>(abiInfo.storage_class().getInt());
50   if (isScalarOrVectorType(varType)) {
51     varType =
52         spirv::PointerType::get(spirv::StructType::get(varType), storageClass);
53   }
54   auto varPtrType = varType.cast<spirv::PointerType>();
55   auto varPointeeType = varPtrType.getPointeeType().cast<spirv::StructType>();
56 
57   // Set the offset information.
58   VulkanLayoutUtils::Size size = 0, alignment = 0;
59   varPointeeType =
60       VulkanLayoutUtils::decorateType(varPointeeType, size, alignment)
61           .cast<spirv::StructType>();
62   varType =
63       spirv::PointerType::get(varPointeeType, varPtrType.getStorageClass());
64 
65   return builder.create<spirv::GlobalVariableOp>(
66       funcOp.getLoc(), varType, varName, abiInfo.descriptor_set().getInt(),
67       abiInfo.binding().getInt());
68 }
69 
70 /// Gets the global variables that need to be specified as interface variable
71 /// with an spv.EntryPointOp. Traverses the body of a entry function to do so.
72 static LogicalResult
73 getInterfaceVariables(spirv::FuncOp funcOp,
74                       SmallVectorImpl<Attribute> &interfaceVars) {
75   auto module = funcOp.getParentOfType<spirv::ModuleOp>();
76   if (!module) {
77     return failure();
78   }
79   llvm::SetVector<Operation *> interfaceVarSet;
80 
81   // TODO(ravishankarm) : This should in reality traverse the entry function
82   // call graph and collect all the interfaces. For now, just traverse the
83   // instructions in this function.
84   funcOp.walk([&](spirv::AddressOfOp addressOfOp) {
85     auto var =
86         module.lookupSymbol<spirv::GlobalVariableOp>(addressOfOp.variable());
87     if (var.type().cast<spirv::PointerType>().getStorageClass() !=
88         spirv::StorageClass::StorageBuffer) {
89       interfaceVarSet.insert(var.getOperation());
90     }
91   });
92   for (auto &var : interfaceVarSet) {
93     interfaceVars.push_back(SymbolRefAttr::get(
94         cast<spirv::GlobalVariableOp>(var).sym_name(), funcOp.getContext()));
95   }
96   return success();
97 }
98 
99 /// Lowers the entry point attribute.
100 static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
101                                             OpBuilder &builder) {
102   auto entryPointAttrName = spirv::getEntryPointABIAttrName();
103   auto entryPointAttr =
104       funcOp.getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName);
105   if (!entryPointAttr) {
106     return failure();
107   }
108 
109   OpBuilder::InsertionGuard moduleInsertionGuard(builder);
110   auto spirvModule = funcOp.getParentOfType<spirv::ModuleOp>();
111   builder.setInsertionPoint(spirvModule.body().front().getTerminator());
112 
113   // Adds the spv.EntryPointOp after collecting all the interface variables
114   // needed.
115   SmallVector<Attribute, 1> interfaceVars;
116   if (failed(getInterfaceVariables(funcOp, interfaceVars))) {
117     return failure();
118   }
119   builder.create<spirv::EntryPointOp>(
120       funcOp.getLoc(), spirv::ExecutionModel::GLCompute, funcOp, interfaceVars);
121   // Specifies the spv.ExecutionModeOp.
122   auto localSizeAttr = entryPointAttr.local_size();
123   SmallVector<int32_t, 3> localSize(localSizeAttr.getValues<int32_t>());
124   builder.create<spirv::ExecutionModeOp>(
125       funcOp.getLoc(), funcOp, spirv::ExecutionMode::LocalSize, localSize);
126   funcOp.removeAttr(entryPointAttrName);
127   return success();
128 }
129 
130 namespace {
131 /// A pattern to convert function signature according to interface variable ABI
132 /// attributes.
133 ///
134 /// Specifically, this pattern creates global variables according to interface
135 /// variable ABI attributes attached to function arguments and converts all
136 /// function argument uses to those global variables. This is necessary because
137 /// Vulkan requires all shader entry points to be of void(void) type.
138 class ProcessInterfaceVarABI final : public SPIRVOpLowering<spirv::FuncOp> {
139 public:
140   using SPIRVOpLowering<spirv::FuncOp>::SPIRVOpLowering;
141   PatternMatchResult
142   matchAndRewrite(spirv::FuncOp funcOp, ArrayRef<Value> operands,
143                   ConversionPatternRewriter &rewriter) const override;
144 };
145 
146 /// Pass to implement the ABI information specified as attributes.
147 class LowerABIAttributesPass final
148     : public OperationPass<LowerABIAttributesPass, spirv::ModuleOp> {
149 private:
150   void runOnOperation() override;
151 };
152 } // namespace
153 
154 PatternMatchResult ProcessInterfaceVarABI::matchAndRewrite(
155     spirv::FuncOp funcOp, ArrayRef<Value> operands,
156     ConversionPatternRewriter &rewriter) const {
157   if (!funcOp.getAttrOfType<spirv::EntryPointABIAttr>(
158           spirv::getEntryPointABIAttrName())) {
159     // TODO(ravishankarm) : Non-entry point functions are not handled.
160     return matchFailure();
161   }
162   TypeConverter::SignatureConversion signatureConverter(
163       funcOp.getType().getNumInputs());
164 
165   auto attrName = spirv::getInterfaceVarABIAttrName();
166   for (auto argType : llvm::enumerate(funcOp.getType().getInputs())) {
167     auto abiInfo = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
168         argType.index(), attrName);
169     if (!abiInfo) {
170       // TODO(ravishankarm) : For non-entry point functions, it should be legal
171       // to pass around scalar/vector values and return a scalar/vector. For now
172       // non-entry point functions are not handled in this ABI lowering and will
173       // produce an error.
174       return matchFailure();
175     }
176     auto var =
177         createGlobalVariableForArg(funcOp, rewriter, argType.index(), abiInfo);
178     if (!var) {
179       return matchFailure();
180     }
181 
182     OpBuilder::InsertionGuard funcInsertionGuard(rewriter);
183     rewriter.setInsertionPointToStart(&funcOp.front());
184     // Insert spirv::AddressOf and spirv::AccessChain operations.
185     Value replacement =
186         rewriter.create<spirv::AddressOfOp>(funcOp.getLoc(), var);
187     // Check if the arg is a scalar or vector type. In that case, the value
188     // needs to be loaded into registers.
189     // TODO(ravishankarm) : This is loading value of the scalar into registers
190     // at the start of the function. It is probably better to do the load just
191     // before the use. There might be multiple loads and currently there is no
192     // easy way to replace all uses with a sequence of operations.
193     if (isScalarOrVectorType(argType.value())) {
194       auto indexType = SPIRVTypeConverter::getIndexType(funcOp.getContext());
195       auto zero =
196           spirv::ConstantOp::getZero(indexType, funcOp.getLoc(), &rewriter);
197       auto loadPtr = rewriter.create<spirv::AccessChainOp>(
198           funcOp.getLoc(), replacement, zero.constant());
199       replacement = rewriter.create<spirv::LoadOp>(funcOp.getLoc(), loadPtr);
200     }
201     signatureConverter.remapInput(argType.index(), replacement);
202   }
203 
204   // Creates a new function with the update signature.
205   rewriter.updateRootInPlace(funcOp, [&] {
206     funcOp.setType(rewriter.getFunctionType(
207         signatureConverter.getConvertedTypes(), llvm::None));
208     rewriter.applySignatureConversion(&funcOp.getBody(), signatureConverter);
209   });
210   return matchSuccess();
211 }
212 
213 void LowerABIAttributesPass::runOnOperation() {
214   // Uses the signature conversion methodology of the dialect conversion
215   // framework to implement the conversion.
216   spirv::ModuleOp module = getOperation();
217   MLIRContext *context = &getContext();
218 
219   SPIRVTypeConverter typeConverter;
220   OwningRewritePatternList patterns;
221   patterns.insert<ProcessInterfaceVarABI>(context, typeConverter);
222 
223   ConversionTarget target(*context);
224   // "Legal" function ops should have no interface variable ABI attributes.
225   target.addDynamicallyLegalOp<spirv::FuncOp>([&](spirv::FuncOp op) {
226     StringRef attrName = spirv::getInterfaceVarABIAttrName();
227     for (unsigned i = 0, e = op.getNumArguments(); i < e; ++i)
228       if (op.getArgAttr(i, attrName))
229         return false;
230     return true;
231   });
232   // All other SPIR-V ops are legal.
233   target.markUnknownOpDynamicallyLegal([](Operation *op) {
234     return op->getDialect()->getNamespace() ==
235            spirv::SPIRVDialect::getDialectNamespace();
236   });
237   if (failed(
238           applyPartialConversion(module, target, patterns, &typeConverter))) {
239     return signalPassFailure();
240   }
241 
242   // Walks over all the FuncOps in spirv::ModuleOp to lower the entry point
243   // attributes.
244   OpBuilder builder(context);
245   SmallVector<spirv::FuncOp, 1> entryPointFns;
246   auto entryPointAttrName = spirv::getEntryPointABIAttrName();
247   module.walk([&](spirv::FuncOp funcOp) {
248     if (funcOp.getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName)) {
249       entryPointFns.push_back(funcOp);
250     }
251   });
252   for (auto fn : entryPointFns) {
253     if (failed(lowerEntryPointABIAttr(fn, builder))) {
254       return signalPassFailure();
255     }
256   }
257 }
258 
259 std::unique_ptr<OpPassBase<spirv::ModuleOp>>
260 mlir::spirv::createLowerABIAttributesPass() {
261   return std::make_unique<LowerABIAttributesPass>();
262 }
263 
264 static PassRegistration<LowerABIAttributesPass>
265     pass("spirv-lower-abi-attrs", "Lower SPIR-V ABI Attributes");
266