1 //===- LinalgToSPIRV.cpp - Linalg 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 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRV.h"
10 #include "mlir/Dialect/Linalg/IR/Linalg.h"
11 #include "mlir/Dialect/Linalg/Utils/Utils.h"
12 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
13 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
14 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
15 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
16 #include "mlir/IR/AffineExpr.h"
17 #include "mlir/Transforms/DialectConversion.h"
18 
19 using namespace mlir;
20 
21 //===----------------------------------------------------------------------===//
22 // Utilities
23 //===----------------------------------------------------------------------===//
24 
25 /// Returns a `Value` containing the `dim`-th dimension's size of SPIR-V
26 /// location invocation ID. This function will create necessary operations with
27 /// `builder` at the proper region containing `op`.
28 static Value getLocalInvocationDimSize(Operation *op, int dim, Type integerType,
29                                        Location loc, OpBuilder *builder) {
30   assert(dim >= 0 && dim < 3 && "local invocation only has three dimensions");
31   Value invocation = spirv::getBuiltinVariableValue(
32       op, spirv::BuiltIn::LocalInvocationId, integerType, *builder);
33   Type xType = invocation.getType().cast<ShapedType>().getElementType();
34   return builder->create<spirv::CompositeExtractOp>(
35       loc, xType, invocation, builder->getI32ArrayAttr({dim}));
36 }
37 
38 //===----------------------------------------------------------------------===//
39 // Reduction (single workgroup)
40 //===----------------------------------------------------------------------===//
41 
42 namespace {
43 
44 /// A pattern to convert a linalg.generic op to SPIR-V ops under the condition
45 /// that the linalg.generic op is performing reduction with a workload size that
46 /// can fit in one workgroup.
47 struct SingleWorkgroupReduction final
48     : public OpConversionPattern<linalg::GenericOp> {
49   using OpConversionPattern::OpConversionPattern;
50 
51   /// Matches the given linalg.generic op as performing reduction and returns
52   /// the binary op kind if successful.
53   static Optional<linalg::RegionMatcher::BinaryOpKind>
54   matchAsPerformingReduction(linalg::GenericOp genericOp);
55 
56   LogicalResult
57   matchAndRewrite(linalg::GenericOp genericOp, OpAdaptor adaptor,
58                   ConversionPatternRewriter &rewriter) const override;
59 };
60 
61 } // namespace
62 
63 Optional<linalg::RegionMatcher::BinaryOpKind>
64 SingleWorkgroupReduction::matchAsPerformingReduction(
65     linalg::GenericOp genericOp) {
66   Operation *op = genericOp.getOperation();
67 
68   // Make sure the linalg.generic is working on memrefs.
69   if (!genericOp.hasBufferSemantics())
70     return llvm::None;
71 
72   // Make sure this is reduction with one input and one output.
73   if (genericOp.getNumInputs() != 1 || genericOp.getNumOutputs() != 1)
74     return llvm::None;
75 
76   auto originalInputType = op->getOperand(0).getType().cast<MemRefType>();
77   auto originalOutputType = op->getOperand(1).getType().cast<MemRefType>();
78 
79   // Make sure the original input has one dimension.
80   if (!originalInputType.hasStaticShape() || originalInputType.getRank() != 1)
81     return llvm::None;
82   // Make sure the original output has one element.
83   if (!originalOutputType.hasStaticShape() ||
84       originalOutputType.getNumElements() != 1)
85     return llvm::None;
86 
87   if (!genericOp.hasSingleReductionLoop())
88     return llvm::None;
89 
90   if (genericOp.indexing_maps().getValue().size() != 2)
91     return llvm::None;
92 
93   // TODO: create utility functions for these checks in Linalg
94   // and use them.
95   auto inputMap = genericOp.indexing_maps().getValue()[0].cast<AffineMapAttr>();
96   auto outputMap =
97       genericOp.indexing_maps().getValue()[1].cast<AffineMapAttr>();
98   // The indexing map for the input should be `(i) -> (i)`.
99   if (inputMap.getValue() !=
100       AffineMap::get(1, 0, getAffineDimExpr(0, op->getContext())))
101     return llvm::None;
102   // The indexing map for the input should be `(i) -> (0)`.
103   if (outputMap.getValue() !=
104       AffineMap::get(1, 0, getAffineConstantExpr(0, op->getContext())))
105     return llvm::None;
106 
107   return linalg::RegionMatcher::matchAsScalarBinaryOp(genericOp);
108 }
109 
110 LogicalResult SingleWorkgroupReduction::matchAndRewrite(
111     linalg::GenericOp genericOp, OpAdaptor adaptor,
112     ConversionPatternRewriter &rewriter) const {
113   Operation *op = genericOp.getOperation();
114   auto originalInputType = op->getOperand(0).getType().cast<MemRefType>();
115   auto originalOutputType = op->getOperand(1).getType().cast<MemRefType>();
116 
117   auto binaryOpKind = matchAsPerformingReduction(genericOp);
118   if (!binaryOpKind)
119     return failure();
120 
121   // Query the shader interface for local workgroup size to make sure the
122   // invocation configuration fits with the input memref's shape.
123   DenseIntElementsAttr localSize = spirv::lookupLocalWorkGroupSize(genericOp);
124   if (!localSize)
125     return failure();
126 
127   if ((*localSize.begin()).getSExtValue() != originalInputType.getDimSize(0))
128     return failure();
129   if (llvm::any_of(llvm::drop_begin(localSize.getValues<APInt>(), 1),
130                    [](const APInt &size) { return !size.isOneValue(); }))
131     return failure();
132 
133   // TODO: Query the target environment to make sure the current
134   // workload fits in a local workgroup.
135 
136   Value convertedInput = adaptor.getOperands()[0];
137   Value convertedOutput = adaptor.getOperands()[1];
138   Location loc = genericOp.getLoc();
139 
140   auto *typeConverter = getTypeConverter<SPIRVTypeConverter>();
141   auto indexType = typeConverter->getIndexType();
142 
143   // Get the invocation ID.
144   Value x = getLocalInvocationDimSize(genericOp, /*dim=*/0, indexType, loc,
145                                       &rewriter);
146 
147   // TODO: Load to Workgroup storage class first.
148 
149   // Get the input element accessed by this invocation.
150   Value inputElementPtr = spirv::getElementPtr(
151       *typeConverter, originalInputType, convertedInput, {x}, loc, rewriter);
152   Value inputElement = rewriter.create<spirv::LoadOp>(loc, inputElementPtr);
153 
154   // Perform the group reduction operation.
155   Value groupOperation;
156 #define CREATE_GROUP_NON_UNIFORM_BIN_OP(opKind, spvOp)                         \
157   case linalg::RegionMatcher::BinaryOpKind::opKind: {                          \
158     groupOperation = rewriter.create<spirv::spvOp>(                            \
159         loc, originalInputType.getElementType(), spirv::Scope::Subgroup,       \
160         spirv::GroupOperation::Reduce, inputElement,                           \
161         /*cluster_size=*/nullptr);                                             \
162   } break
163   switch (*binaryOpKind) {
164     CREATE_GROUP_NON_UNIFORM_BIN_OP(IAdd, GroupNonUniformIAddOp);
165   }
166 #undef CREATE_GROUP_NON_UNIFORM_BIN_OP
167 
168   // Get the output element accessed by this reduction.
169   Value zero = spirv::ConstantOp::getZero(indexType, loc, rewriter);
170   SmallVector<Value, 1> zeroIndices(originalOutputType.getRank(), zero);
171   Value outputElementPtr =
172       spirv::getElementPtr(*typeConverter, originalOutputType, convertedOutput,
173                            zeroIndices, loc, rewriter);
174 
175   // Write out the final reduction result. This should be only conducted by one
176   // invocation. We use spv.GroupNonUniformElect to find the invocation with the
177   // lowest ID.
178   //
179   // ```
180   // if (spv.GroupNonUniformElect) { output = ... }
181   // ```
182 
183   Value condition = rewriter.create<spirv::GroupNonUniformElectOp>(
184       loc, spirv::Scope::Subgroup);
185 
186   auto createAtomicOp = [&](OpBuilder &builder) {
187 #define CREATE_ATOMIC_BIN_OP(opKind, spvOp)                                    \
188   case linalg::RegionMatcher::BinaryOpKind::opKind: {                          \
189     builder.create<spirv::spvOp>(loc, outputElementPtr, spirv::Scope::Device,  \
190                                  spirv::MemorySemantics::AcquireRelease,       \
191                                  groupOperation);                              \
192   } break
193     switch (*binaryOpKind) { CREATE_ATOMIC_BIN_OP(IAdd, AtomicIAddOp); }
194 #undef CREATE_ATOMIC_BIN_OP
195   };
196 
197   spirv::SelectionOp::createIfThen(loc, condition, createAtomicOp, rewriter);
198 
199   rewriter.eraseOp(genericOp);
200   return success();
201 }
202 
203 //===----------------------------------------------------------------------===//
204 // Pattern population
205 //===----------------------------------------------------------------------===//
206 
207 void mlir::populateLinalgToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
208                                          RewritePatternSet &patterns) {
209   patterns.add<SingleWorkgroupReduction>(typeConverter, patterns.getContext());
210 }
211