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 =
45         static_cast<spirv::StorageClass>(abiInfo.storage_class().getInt());
46     varType =
47         spirv::PointerType::get(spirv::StructType::get(varType), storageClass);
48   }
49   auto varPtrType = varType.cast<spirv::PointerType>();
50   auto varPointeeType = varPtrType.getPointeeType().cast<spirv::StructType>();
51 
52   // Set the offset information.
53   varPointeeType =
54       VulkanLayoutUtils::decorateType(varPointeeType).cast<spirv::StructType>();
55   varType =
56       spirv::PointerType::get(varPointeeType, varPtrType.getStorageClass());
57 
58   return builder.create<spirv::GlobalVariableOp>(
59       funcOp.getLoc(), varType, varName, abiInfo.descriptor_set().getInt(),
60       abiInfo.binding().getInt());
61 }
62 
63 /// Gets the global variables that need to be specified as interface variable
64 /// with an spv.EntryPointOp. Traverses the body of a entry function to do so.
65 static LogicalResult
66 getInterfaceVariables(spirv::FuncOp funcOp,
67                       SmallVectorImpl<Attribute> &interfaceVars) {
68   auto module = funcOp.getParentOfType<spirv::ModuleOp>();
69   if (!module) {
70     return failure();
71   }
72   llvm::SetVector<Operation *> interfaceVarSet;
73 
74   // TODO(ravishankarm) : This should in reality traverse the entry function
75   // call graph and collect all the interfaces. For now, just traverse the
76   // instructions in this function.
77   funcOp.walk([&](spirv::AddressOfOp addressOfOp) {
78     auto var =
79         module.lookupSymbol<spirv::GlobalVariableOp>(addressOfOp.variable());
80     // TODO(antiagainst): Per SPIR-V spec: "Before version 1.4, the interface’s
81     // storage classes are limited to the Input and Output storage classes.
82     // Starting with version 1.4, the interface’s storage classes are all
83     // storage classes used in declaring all global variables referenced by the
84     // entry point’s call tree." We should consider the target environment here.
85     switch (var.type().cast<spirv::PointerType>().getStorageClass()) {
86     case spirv::StorageClass::Input:
87     case spirv::StorageClass::Output:
88       interfaceVarSet.insert(var.getOperation());
89       break;
90     default:
91       break;
92     }
93   });
94   for (auto &var : interfaceVarSet) {
95     interfaceVars.push_back(SymbolRefAttr::get(
96         cast<spirv::GlobalVariableOp>(var).sym_name(), funcOp.getContext()));
97   }
98   return success();
99 }
100 
101 /// Lowers the entry point attribute.
102 static LogicalResult lowerEntryPointABIAttr(spirv::FuncOp funcOp,
103                                             OpBuilder &builder) {
104   auto entryPointAttrName = spirv::getEntryPointABIAttrName();
105   auto entryPointAttr =
106       funcOp.getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName);
107   if (!entryPointAttr) {
108     return failure();
109   }
110 
111   OpBuilder::InsertionGuard moduleInsertionGuard(builder);
112   auto spirvModule = funcOp.getParentOfType<spirv::ModuleOp>();
113   builder.setInsertionPoint(spirvModule.body().front().getTerminator());
114 
115   // Adds the spv.EntryPointOp after collecting all the interface variables
116   // needed.
117   SmallVector<Attribute, 1> interfaceVars;
118   if (failed(getInterfaceVariables(funcOp, interfaceVars))) {
119     return failure();
120   }
121   builder.create<spirv::EntryPointOp>(
122       funcOp.getLoc(), spirv::ExecutionModel::GLCompute, funcOp, interfaceVars);
123   // Specifies the spv.ExecutionModeOp.
124   auto localSizeAttr = entryPointAttr.local_size();
125   SmallVector<int32_t, 3> localSize(localSizeAttr.getValues<int32_t>());
126   builder.create<spirv::ExecutionModeOp>(
127       funcOp.getLoc(), funcOp, spirv::ExecutionMode::LocalSize, localSize);
128   funcOp.removeAttr(entryPointAttrName);
129   return success();
130 }
131 
132 namespace {
133 /// A pattern to convert function signature according to interface variable ABI
134 /// attributes.
135 ///
136 /// Specifically, this pattern creates global variables according to interface
137 /// variable ABI attributes attached to function arguments and converts all
138 /// function argument uses to those global variables. This is necessary because
139 /// Vulkan requires all shader entry points to be of void(void) type.
140 class ProcessInterfaceVarABI final : public SPIRVOpLowering<spirv::FuncOp> {
141 public:
142   using SPIRVOpLowering<spirv::FuncOp>::SPIRVOpLowering;
143   LogicalResult
144   matchAndRewrite(spirv::FuncOp funcOp, ArrayRef<Value> operands,
145                   ConversionPatternRewriter &rewriter) const override;
146 };
147 
148 /// Pass to implement the ABI information specified as attributes.
149 class LowerABIAttributesPass final
150     : public SPIRVLowerABIAttributesBase<LowerABIAttributesPass> {
151   void runOnOperation() override;
152 };
153 } // namespace
154 
155 LogicalResult ProcessInterfaceVarABI::matchAndRewrite(
156     spirv::FuncOp funcOp, ArrayRef<Value> operands,
157     ConversionPatternRewriter &rewriter) const {
158   if (!funcOp.getAttrOfType<spirv::EntryPointABIAttr>(
159           spirv::getEntryPointABIAttrName())) {
160     // TODO(ravishankarm) : Non-entry point functions are not handled.
161     return failure();
162   }
163   TypeConverter::SignatureConversion signatureConverter(
164       funcOp.getType().getNumInputs());
165 
166   auto attrName = spirv::getInterfaceVarABIAttrName();
167   for (auto argType : llvm::enumerate(funcOp.getType().getInputs())) {
168     auto abiInfo = funcOp.getArgAttrOfType<spirv::InterfaceVarABIAttr>(
169         argType.index(), attrName);
170     if (!abiInfo) {
171       // TODO(ravishankarm) : For non-entry point functions, it should be legal
172       // to pass around scalar/vector values and return a scalar/vector. For now
173       // non-entry point functions are not handled in this ABI lowering and will
174       // produce an error.
175       return failure();
176     }
177     spirv::GlobalVariableOp var = createGlobalVarForEntryPointArgument(
178         rewriter, funcOp, argType.index(), abiInfo);
179     if (!var)
180       return failure();
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 (argType.value().cast<spirv::SPIRVType>().isScalarOrVector()) {
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 success();
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   spirv::TargetEnv targetEnv(spirv::lookupTargetEnv(module));
220 
221   SPIRVTypeConverter typeConverter(targetEnv);
222   OwningRewritePatternList patterns;
223   patterns.insert<ProcessInterfaceVarABI>(context, typeConverter);
224 
225   ConversionTarget target(*context);
226   // "Legal" function ops should have no interface variable ABI attributes.
227   target.addDynamicallyLegalOp<spirv::FuncOp>([&](spirv::FuncOp op) {
228     StringRef attrName = spirv::getInterfaceVarABIAttrName();
229     for (unsigned i = 0, e = op.getNumArguments(); i < e; ++i)
230       if (op.getArgAttr(i, attrName))
231         return false;
232     return true;
233   });
234   // All other SPIR-V ops are legal.
235   target.markUnknownOpDynamicallyLegal([](Operation *op) {
236     return op->getDialect()->getNamespace() ==
237            spirv::SPIRVDialect::getDialectNamespace();
238   });
239   if (failed(
240           applyPartialConversion(module, target, patterns, &typeConverter))) {
241     return signalPassFailure();
242   }
243 
244   // Walks over all the FuncOps in spirv::ModuleOp to lower the entry point
245   // attributes.
246   OpBuilder builder(context);
247   SmallVector<spirv::FuncOp, 1> entryPointFns;
248   auto entryPointAttrName = spirv::getEntryPointABIAttrName();
249   module.walk([&](spirv::FuncOp funcOp) {
250     if (funcOp.getAttrOfType<spirv::EntryPointABIAttr>(entryPointAttrName)) {
251       entryPointFns.push_back(funcOp);
252     }
253   });
254   for (auto fn : entryPointFns) {
255     if (failed(lowerEntryPointABIAttr(fn, builder))) {
256       return signalPassFailure();
257     }
258   }
259 }
260 
261 std::unique_ptr<OperationPass<spirv::ModuleOp>>
262 mlir::spirv::createLowerABIAttributesPass() {
263   return std::make_unique<LowerABIAttributesPass>();
264 }
265