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