1 //===- SCFToSPIRV.cpp - SCF 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 // This file implements patterns to convert SCF dialect to SPIR-V dialect.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
14 #include "mlir/Dialect/SCF/SCF.h"
15 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
16 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
17 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
18 #include "mlir/IR/BuiltinOps.h"
19 #include "mlir/Transforms/DialectConversion.h"
20 
21 using namespace mlir;
22 
23 //===----------------------------------------------------------------------===//
24 // Context
25 //===----------------------------------------------------------------------===//
26 
27 namespace mlir {
28 struct ScfToSPIRVContextImpl {
29   // Map between the spirv region control flow operation (spv.loop or
30   // spv.selection) to the VariableOp created to store the region results. The
31   // order of the VariableOp matches the order of the results.
32   DenseMap<Operation *, SmallVector<spirv::VariableOp, 8>> outputVars;
33 };
34 } // namespace mlir
35 
36 /// We use ScfToSPIRVContext to store information about the lowering of the scf
37 /// region that need to be used later on. When we lower scf.for/scf.if we create
38 /// VariableOp to store the results. We need to keep track of the VariableOp
39 /// created as we need to insert stores into them when lowering Yield. Those
40 /// StoreOp cannot be created earlier as they may use a different type than
41 /// yield operands.
42 ScfToSPIRVContext::ScfToSPIRVContext() {
43   impl = std::make_unique<ScfToSPIRVContextImpl>();
44 }
45 
46 ScfToSPIRVContext::~ScfToSPIRVContext() = default;
47 
48 //===----------------------------------------------------------------------===//
49 // Pattern Declarations
50 //===----------------------------------------------------------------------===//
51 
52 namespace {
53 /// Common class for all vector to GPU patterns.
54 template <typename OpTy>
55 class SCFToSPIRVPattern : public OpConversionPattern<OpTy> {
56 public:
57   SCFToSPIRVPattern<OpTy>(MLIRContext *context, SPIRVTypeConverter &converter,
58                           ScfToSPIRVContextImpl *scfToSPIRVContext)
59       : OpConversionPattern<OpTy>::OpConversionPattern(context),
60         scfToSPIRVContext(scfToSPIRVContext), typeConverter(converter) {}
61 
62 protected:
63   ScfToSPIRVContextImpl *scfToSPIRVContext;
64   // FIXME: We explicitly keep a reference of the type converter here instead of
65   // passing it to OpConversionPattern during construction. This effectively
66   // bypasses the conversion framework's automation on type conversion. This is
67   // needed right now because the conversion framework will unconditionally
68   // legalize all types used by SCF ops upon discovering them, for example, the
69   // types of loop carried values. We use SPIR-V variables for those loop
70   // carried values. Depending on the available capabilities, the SPIR-V
71   // variable can be different, for example, cooperative matrix or normal
72   // variable. We'd like to detach the conversion of the loop carried values
73   // from the SCF ops (which is mainly a region). So we need to "mark" types
74   // used by SCF ops as legal, if to use the conversion framework for type
75   // conversion. There isn't a straightforward way to do that yet, as when
76   // converting types, ops aren't taken into consideration. Therefore, we just
77   // bypass the framework's type conversion for now.
78   SPIRVTypeConverter &typeConverter;
79 };
80 
81 /// Pattern to convert a scf::ForOp within kernel functions into spirv::LoopOp.
82 class ForOpConversion final : public SCFToSPIRVPattern<scf::ForOp> {
83 public:
84   using SCFToSPIRVPattern<scf::ForOp>::SCFToSPIRVPattern;
85 
86   LogicalResult
87   matchAndRewrite(scf::ForOp forOp, ArrayRef<Value> operands,
88                   ConversionPatternRewriter &rewriter) const override;
89 };
90 
91 /// Pattern to convert a scf::IfOp within kernel functions into
92 /// spirv::SelectionOp.
93 class IfOpConversion final : public SCFToSPIRVPattern<scf::IfOp> {
94 public:
95   using SCFToSPIRVPattern<scf::IfOp>::SCFToSPIRVPattern;
96 
97   LogicalResult
98   matchAndRewrite(scf::IfOp ifOp, ArrayRef<Value> operands,
99                   ConversionPatternRewriter &rewriter) const override;
100 };
101 
102 class TerminatorOpConversion final : public SCFToSPIRVPattern<scf::YieldOp> {
103 public:
104   using SCFToSPIRVPattern<scf::YieldOp>::SCFToSPIRVPattern;
105 
106   LogicalResult
107   matchAndRewrite(scf::YieldOp terminatorOp, ArrayRef<Value> operands,
108                   ConversionPatternRewriter &rewriter) const override;
109 };
110 } // namespace
111 
112 /// Helper function to replaces SCF op outputs with SPIR-V variable loads.
113 /// We create VariableOp to handle the results value of the control flow region.
114 /// spv.loop/spv.selection currently don't yield value. Right after the loop
115 /// we load the value from the allocation and use it as the SCF op result.
116 template <typename ScfOp, typename OpTy>
117 static void replaceSCFOutputValue(ScfOp scfOp, OpTy newOp,
118                                   ConversionPatternRewriter &rewriter,
119                                   ScfToSPIRVContextImpl *scfToSPIRVContext,
120                                   ArrayRef<Type> returnTypes) {
121 
122   Location loc = scfOp.getLoc();
123   auto &allocas = scfToSPIRVContext->outputVars[newOp];
124   // Clearing the allocas is necessary in case a dialect conversion path failed
125   // previously, and this is the second attempt of this conversion.
126   allocas.clear();
127   SmallVector<Value, 8> resultValue;
128   for (Type convertedType : returnTypes) {
129     auto pointerType =
130         spirv::PointerType::get(convertedType, spirv::StorageClass::Function);
131     rewriter.setInsertionPoint(newOp);
132     auto alloc = rewriter.create<spirv::VariableOp>(
133         loc, pointerType, spirv::StorageClass::Function,
134         /*initializer=*/nullptr);
135     allocas.push_back(alloc);
136     rewriter.setInsertionPointAfter(newOp);
137     Value loadResult = rewriter.create<spirv::LoadOp>(loc, alloc);
138     resultValue.push_back(loadResult);
139   }
140   rewriter.replaceOp(scfOp, resultValue);
141 }
142 
143 //===----------------------------------------------------------------------===//
144 // scf::ForOp
145 //===----------------------------------------------------------------------===//
146 
147 LogicalResult
148 ForOpConversion::matchAndRewrite(scf::ForOp forOp, ArrayRef<Value> operands,
149                                  ConversionPatternRewriter &rewriter) const {
150   // scf::ForOp can be lowered to the structured control flow represented by
151   // spirv::LoopOp by making the continue block of the spirv::LoopOp the loop
152   // latch and the merge block the exit block. The resulting spirv::LoopOp has a
153   // single back edge from the continue to header block, and a single exit from
154   // header to merge.
155   scf::ForOpAdaptor forOperands(operands);
156   auto loc = forOp.getLoc();
157   auto loopControl = rewriter.getI32IntegerAttr(
158       static_cast<uint32_t>(spirv::LoopControl::None));
159   auto loopOp = rewriter.create<spirv::LoopOp>(loc, loopControl);
160   loopOp.addEntryAndMergeBlock();
161 
162   OpBuilder::InsertionGuard guard(rewriter);
163   // Create the block for the header.
164   auto *header = new Block();
165   // Insert the header.
166   loopOp.body().getBlocks().insert(std::next(loopOp.body().begin(), 1), header);
167 
168   // Create the new induction variable to use.
169   BlockArgument newIndVar =
170       header->addArgument(forOperands.lowerBound().getType());
171   for (Value arg : forOperands.initArgs())
172     header->addArgument(arg.getType());
173   Block *body = forOp.getBody();
174 
175   // Apply signature conversion to the body of the forOp. It has a single block,
176   // with argument which is the induction variable. That has to be replaced with
177   // the new induction variable.
178   TypeConverter::SignatureConversion signatureConverter(
179       body->getNumArguments());
180   signatureConverter.remapInput(0, newIndVar);
181   for (unsigned i = 1, e = body->getNumArguments(); i < e; i++)
182     signatureConverter.remapInput(i, header->getArgument(i));
183   body = rewriter.applySignatureConversion(&forOp.getLoopBody(),
184                                            signatureConverter);
185 
186   // Move the blocks from the forOp into the loopOp. This is the body of the
187   // loopOp.
188   rewriter.inlineRegionBefore(forOp->getRegion(0), loopOp.body(),
189                               std::next(loopOp.body().begin(), 2));
190 
191   SmallVector<Value, 8> args(1, forOperands.lowerBound());
192   args.append(forOperands.initArgs().begin(), forOperands.initArgs().end());
193   // Branch into it from the entry.
194   rewriter.setInsertionPointToEnd(&(loopOp.body().front()));
195   rewriter.create<spirv::BranchOp>(loc, header, args);
196 
197   // Generate the rest of the loop header.
198   rewriter.setInsertionPointToEnd(header);
199   auto *mergeBlock = loopOp.getMergeBlock();
200   auto cmpOp = rewriter.create<spirv::SLessThanOp>(
201       loc, rewriter.getI1Type(), newIndVar, forOperands.upperBound());
202 
203   rewriter.create<spirv::BranchConditionalOp>(
204       loc, cmpOp, body, ArrayRef<Value>(), mergeBlock, ArrayRef<Value>());
205 
206   // Generate instructions to increment the step of the induction variable and
207   // branch to the header.
208   Block *continueBlock = loopOp.getContinueBlock();
209   rewriter.setInsertionPointToEnd(continueBlock);
210 
211   // Add the step to the induction variable and branch to the header.
212   Value updatedIndVar = rewriter.create<spirv::IAddOp>(
213       loc, newIndVar.getType(), newIndVar, forOperands.step());
214   rewriter.create<spirv::BranchOp>(loc, header, updatedIndVar);
215 
216   // Infer the return types from the init operands. Vector type may get
217   // converted to CooperativeMatrix or to Vector type, to avoid having complex
218   // extra logic to figure out the right type we just infer it from the Init
219   // operands.
220   SmallVector<Type, 8> initTypes;
221   for (auto arg : forOperands.initArgs())
222     initTypes.push_back(arg.getType());
223   replaceSCFOutputValue(forOp, loopOp, rewriter, scfToSPIRVContext, initTypes);
224   return success();
225 }
226 
227 //===----------------------------------------------------------------------===//
228 // scf::IfOp
229 //===----------------------------------------------------------------------===//
230 
231 LogicalResult
232 IfOpConversion::matchAndRewrite(scf::IfOp ifOp, ArrayRef<Value> operands,
233                                 ConversionPatternRewriter &rewriter) const {
234   // When lowering `scf::IfOp` we explicitly create a selection header block
235   // before the control flow diverges and a merge block where control flow
236   // subsequently converges.
237   scf::IfOpAdaptor ifOperands(operands);
238   auto loc = ifOp.getLoc();
239 
240   // Create `spv.selection` operation, selection header block and merge block.
241   auto selectionControl = rewriter.getI32IntegerAttr(
242       static_cast<uint32_t>(spirv::SelectionControl::None));
243   auto selectionOp = rewriter.create<spirv::SelectionOp>(loc, selectionControl);
244   auto *mergeBlock =
245       rewriter.createBlock(&selectionOp.body(), selectionOp.body().end());
246   rewriter.create<spirv::MergeOp>(loc);
247 
248   OpBuilder::InsertionGuard guard(rewriter);
249   auto *selectionHeaderBlock =
250       rewriter.createBlock(&selectionOp.body().front());
251 
252   // Inline `then` region before the merge block and branch to it.
253   auto &thenRegion = ifOp.thenRegion();
254   auto *thenBlock = &thenRegion.front();
255   rewriter.setInsertionPointToEnd(&thenRegion.back());
256   rewriter.create<spirv::BranchOp>(loc, mergeBlock);
257   rewriter.inlineRegionBefore(thenRegion, mergeBlock);
258 
259   auto *elseBlock = mergeBlock;
260   // If `else` region is not empty, inline that region before the merge block
261   // and branch to it.
262   if (!ifOp.elseRegion().empty()) {
263     auto &elseRegion = ifOp.elseRegion();
264     elseBlock = &elseRegion.front();
265     rewriter.setInsertionPointToEnd(&elseRegion.back());
266     rewriter.create<spirv::BranchOp>(loc, mergeBlock);
267     rewriter.inlineRegionBefore(elseRegion, mergeBlock);
268   }
269 
270   // Create a `spv.BranchConditional` operation for selection header block.
271   rewriter.setInsertionPointToEnd(selectionHeaderBlock);
272   rewriter.create<spirv::BranchConditionalOp>(loc, ifOperands.condition(),
273                                               thenBlock, ArrayRef<Value>(),
274                                               elseBlock, ArrayRef<Value>());
275 
276   SmallVector<Type, 8> returnTypes;
277   for (auto result : ifOp.results()) {
278     auto convertedType = typeConverter.convertType(result.getType());
279     returnTypes.push_back(convertedType);
280   }
281   replaceSCFOutputValue(ifOp, selectionOp, rewriter, scfToSPIRVContext,
282                         returnTypes);
283   return success();
284 }
285 
286 //===----------------------------------------------------------------------===//
287 // scf::YieldOp
288 //===----------------------------------------------------------------------===//
289 
290 /// Yield is lowered to stores to the VariableOp created during lowering of the
291 /// parent region. For loops we also need to update the branch looping back to
292 /// the header with the loop carried values.
293 LogicalResult TerminatorOpConversion::matchAndRewrite(
294     scf::YieldOp terminatorOp, ArrayRef<Value> operands,
295     ConversionPatternRewriter &rewriter) const {
296   // If the region is return values, store each value into the associated
297   // VariableOp created during lowering of the parent region.
298   if (!operands.empty()) {
299     auto loc = terminatorOp.getLoc();
300     auto &allocas = scfToSPIRVContext->outputVars[terminatorOp->getParentOp()];
301     assert(allocas.size() == operands.size());
302     for (unsigned i = 0, e = operands.size(); i < e; i++)
303       rewriter.create<spirv::StoreOp>(loc, allocas[i], operands[i]);
304     if (isa<spirv::LoopOp>(terminatorOp->getParentOp())) {
305       // For loops we also need to update the branch jumping back to the header.
306       auto br =
307           cast<spirv::BranchOp>(rewriter.getInsertionBlock()->getTerminator());
308       SmallVector<Value, 8> args(br.getBlockArguments());
309       args.append(operands.begin(), operands.end());
310       rewriter.setInsertionPoint(br);
311       rewriter.create<spirv::BranchOp>(terminatorOp.getLoc(), br.getTarget(),
312                                        args);
313       rewriter.eraseOp(br);
314     }
315   }
316   rewriter.eraseOp(terminatorOp);
317   return success();
318 }
319 
320 //===----------------------------------------------------------------------===//
321 // Hooks
322 //===----------------------------------------------------------------------===//
323 
324 void mlir::populateSCFToSPIRVPatterns(MLIRContext *context,
325                                       SPIRVTypeConverter &typeConverter,
326                                       ScfToSPIRVContext &scfToSPIRVContext,
327                                       OwningRewritePatternList &patterns) {
328   patterns.insert<ForOpConversion, IfOpConversion, TerminatorOpConversion>(
329       context, typeConverter, scfToSPIRVContext.getImpl());
330 }
331