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/IR/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 | #include "llvm/Support/FormatVariadic.h" |
21 | |
22 | using namespace mlir; |
23 | |
24 | //===----------------------------------------------------------------------===// |
25 | // Context |
26 | //===----------------------------------------------------------------------===// |
27 | |
28 | namespace mlir { |
29 | struct ScfToSPIRVContextImpl { |
30 | // Map between the spirv region control flow operation (spirv.mlir.loop or |
31 | // spirv.mlir.selection) to the VariableOp created to store the region |
32 | // results. The order of the VariableOp matches the order of the results. |
33 | DenseMap<Operation *, SmallVector<spirv::VariableOp, 8>> outputVars; |
34 | }; |
35 | } // namespace mlir |
36 | |
37 | /// We use ScfToSPIRVContext to store information about the lowering of the scf |
38 | /// region that need to be used later on. When we lower scf.for/scf.if we create |
39 | /// VariableOp to store the results. We need to keep track of the VariableOp |
40 | /// created as we need to insert stores into them when lowering Yield. Those |
41 | /// StoreOp cannot be created earlier as they may use a different type than |
42 | /// yield operands. |
43 | ScfToSPIRVContext::ScfToSPIRVContext() { |
44 | impl = std::make_unique<::ScfToSPIRVContextImpl>(); |
45 | } |
46 | |
47 | ScfToSPIRVContext::~ScfToSPIRVContext() = default; |
48 | |
49 | namespace { |
50 | |
51 | //===----------------------------------------------------------------------===// |
52 | // Helper Functions |
53 | //===----------------------------------------------------------------------===// |
54 | |
55 | /// Replaces SCF op outputs with SPIR-V variable loads. |
56 | /// We create VariableOp to handle the results value of the control flow region. |
57 | /// spirv.mlir.loop/spirv.mlir.selection currently don't yield value. Right |
58 | /// after the loop we load the value from the allocation and use it as the SCF |
59 | /// op result. |
60 | template <typename ScfOp, typename OpTy> |
61 | void replaceSCFOutputValue(ScfOp scfOp, OpTy newOp, |
62 | ConversionPatternRewriter &rewriter, |
63 | ScfToSPIRVContextImpl *scfToSPIRVContext, |
64 | ArrayRef<Type> returnTypes) { |
65 | |
66 | Location loc = scfOp.getLoc(); |
67 | auto &allocas = scfToSPIRVContext->outputVars[newOp]; |
68 | // Clearing the allocas is necessary in case a dialect conversion path failed |
69 | // previously, and this is the second attempt of this conversion. |
70 | allocas.clear(); |
71 | SmallVector<Value, 8> resultValue; |
72 | for (Type convertedType : returnTypes) { |
73 | auto pointerType = |
74 | spirv::PointerType::get(convertedType, spirv::StorageClass::Function); |
75 | rewriter.setInsertionPoint(newOp); |
76 | auto alloc = rewriter.create<spirv::VariableOp>( |
77 | loc, pointerType, spirv::StorageClass::Function, |
78 | /*initializer=*/nullptr); |
79 | allocas.push_back(alloc); |
80 | rewriter.setInsertionPointAfter(newOp); |
81 | Value loadResult = rewriter.create<spirv::LoadOp>(loc, alloc); |
82 | resultValue.push_back(Elt: loadResult); |
83 | } |
84 | rewriter.replaceOp(scfOp, resultValue); |
85 | } |
86 | |
87 | Region::iterator getBlockIt(Region ®ion, unsigned index) { |
88 | return std::next(x: region.begin(), n: index); |
89 | } |
90 | |
91 | //===----------------------------------------------------------------------===// |
92 | // Conversion Patterns |
93 | //===----------------------------------------------------------------------===// |
94 | |
95 | /// Common class for all vector to GPU patterns. |
96 | template <typename OpTy> |
97 | class SCFToSPIRVPattern : public OpConversionPattern<OpTy> { |
98 | public: |
99 | SCFToSPIRVPattern(MLIRContext *context, const SPIRVTypeConverter &converter, |
100 | ScfToSPIRVContextImpl *scfToSPIRVContext) |
101 | : OpConversionPattern<OpTy>::OpConversionPattern(converter, context), |
102 | scfToSPIRVContext(scfToSPIRVContext), typeConverter(converter) {} |
103 | |
104 | protected: |
105 | ScfToSPIRVContextImpl *scfToSPIRVContext; |
106 | // FIXME: We explicitly keep a reference of the type converter here instead of |
107 | // passing it to OpConversionPattern during construction. This effectively |
108 | // bypasses the conversion framework's automation on type conversion. This is |
109 | // needed right now because the conversion framework will unconditionally |
110 | // legalize all types used by SCF ops upon discovering them, for example, the |
111 | // types of loop carried values. We use SPIR-V variables for those loop |
112 | // carried values. Depending on the available capabilities, the SPIR-V |
113 | // variable can be different, for example, cooperative matrix or normal |
114 | // variable. We'd like to detach the conversion of the loop carried values |
115 | // from the SCF ops (which is mainly a region). So we need to "mark" types |
116 | // used by SCF ops as legal, if to use the conversion framework for type |
117 | // conversion. There isn't a straightforward way to do that yet, as when |
118 | // converting types, ops aren't taken into consideration. Therefore, we just |
119 | // bypass the framework's type conversion for now. |
120 | const SPIRVTypeConverter &typeConverter; |
121 | }; |
122 | |
123 | //===----------------------------------------------------------------------===// |
124 | // scf::ForOp |
125 | //===----------------------------------------------------------------------===// |
126 | |
127 | /// Pattern to convert a scf::ForOp within kernel functions into spirv::LoopOp. |
128 | struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> { |
129 | using SCFToSPIRVPattern::SCFToSPIRVPattern; |
130 | |
131 | LogicalResult |
132 | matchAndRewrite(scf::ForOp forOp, OpAdaptor adaptor, |
133 | ConversionPatternRewriter &rewriter) const override { |
134 | // scf::ForOp can be lowered to the structured control flow represented by |
135 | // spirv::LoopOp by making the continue block of the spirv::LoopOp the loop |
136 | // latch and the merge block the exit block. The resulting spirv::LoopOp has |
137 | // a single back edge from the continue to header block, and a single exit |
138 | // from header to merge. |
139 | auto loc = forOp.getLoc(); |
140 | auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None); |
141 | loopOp.addEntryAndMergeBlock(rewriter); |
142 | |
143 | OpBuilder::InsertionGuard guard(rewriter); |
144 | // Create the block for the header. |
145 | Block * = rewriter.createBlock(&loopOp.getBody(), |
146 | getBlockIt(loopOp.getBody(), 1)); |
147 | rewriter.setInsertionPointAfter(loopOp); |
148 | |
149 | // Create the new induction variable to use. |
150 | Value adapLowerBound = adaptor.getLowerBound(); |
151 | BlockArgument newIndVar = |
152 | header->addArgument(type: adapLowerBound.getType(), loc: adapLowerBound.getLoc()); |
153 | for (Value arg : adaptor.getInitArgs()) |
154 | header->addArgument(arg.getType(), arg.getLoc()); |
155 | Block *body = forOp.getBody(); |
156 | |
157 | // Apply signature conversion to the body of the forOp. It has a single |
158 | // block, with argument which is the induction variable. That has to be |
159 | // replaced with the new induction variable. |
160 | TypeConverter::SignatureConversion signatureConverter( |
161 | body->getNumArguments()); |
162 | signatureConverter.remapInput(origInputNo: 0, replacements: newIndVar); |
163 | for (unsigned i = 1, e = body->getNumArguments(); i < e; i++) |
164 | signatureConverter.remapInput(origInputNo: i, replacements: header->getArgument(i)); |
165 | body = rewriter.applySignatureConversion(block: &forOp.getRegion().front(), |
166 | conversion&: signatureConverter); |
167 | |
168 | // Move the blocks from the forOp into the loopOp. This is the body of the |
169 | // loopOp. |
170 | rewriter.inlineRegionBefore(forOp->getRegion(0), loopOp.getBody(), |
171 | getBlockIt(loopOp.getBody(), 2)); |
172 | |
173 | SmallVector<Value, 8> args(1, adaptor.getLowerBound()); |
174 | args.append(adaptor.getInitArgs().begin(), adaptor.getInitArgs().end()); |
175 | // Branch into it from the entry. |
176 | rewriter.setInsertionPointToEnd(&(loopOp.getBody().front())); |
177 | rewriter.create<spirv::BranchOp>(loc, header, args); |
178 | |
179 | // Generate the rest of the loop header. |
180 | rewriter.setInsertionPointToEnd(header); |
181 | auto *mergeBlock = loopOp.getMergeBlock(); |
182 | auto cmpOp = rewriter.create<spirv::SLessThanOp>( |
183 | loc, rewriter.getI1Type(), newIndVar, adaptor.getUpperBound()); |
184 | |
185 | rewriter.create<spirv::BranchConditionalOp>( |
186 | loc, cmpOp, body, ArrayRef<Value>(), mergeBlock, ArrayRef<Value>()); |
187 | |
188 | // Generate instructions to increment the step of the induction variable and |
189 | // branch to the header. |
190 | Block *continueBlock = loopOp.getContinueBlock(); |
191 | rewriter.setInsertionPointToEnd(continueBlock); |
192 | |
193 | // Add the step to the induction variable and branch to the header. |
194 | Value updatedIndVar = rewriter.create<spirv::IAddOp>( |
195 | loc, newIndVar.getType(), newIndVar, adaptor.getStep()); |
196 | rewriter.create<spirv::BranchOp>(loc, header, updatedIndVar); |
197 | |
198 | // Infer the return types from the init operands. Vector type may get |
199 | // converted to CooperativeMatrix or to Vector type, to avoid having complex |
200 | // extra logic to figure out the right type we just infer it from the Init |
201 | // operands. |
202 | SmallVector<Type, 8> initTypes; |
203 | for (auto arg : adaptor.getInitArgs()) |
204 | initTypes.push_back(arg.getType()); |
205 | replaceSCFOutputValue(forOp, loopOp, rewriter, scfToSPIRVContext, |
206 | initTypes); |
207 | return success(); |
208 | } |
209 | }; |
210 | |
211 | //===----------------------------------------------------------------------===// |
212 | // scf::IfOp |
213 | //===----------------------------------------------------------------------===// |
214 | |
215 | /// Pattern to convert a scf::IfOp within kernel functions into |
216 | /// spirv::SelectionOp. |
217 | struct IfOpConversion : SCFToSPIRVPattern<scf::IfOp> { |
218 | using SCFToSPIRVPattern::SCFToSPIRVPattern; |
219 | |
220 | LogicalResult |
221 | matchAndRewrite(scf::IfOp ifOp, OpAdaptor adaptor, |
222 | ConversionPatternRewriter &rewriter) const override { |
223 | // When lowering `scf::IfOp` we explicitly create a selection header block |
224 | // before the control flow diverges and a merge block where control flow |
225 | // subsequently converges. |
226 | auto loc = ifOp.getLoc(); |
227 | |
228 | // Compute return types. |
229 | SmallVector<Type, 8> returnTypes; |
230 | for (auto result : ifOp.getResults()) { |
231 | auto convertedType = typeConverter.convertType(result.getType()); |
232 | if (!convertedType) |
233 | return rewriter.notifyMatchFailure( |
234 | loc, |
235 | llvm::formatv("failed to convert type '{0}'" , result.getType())); |
236 | |
237 | returnTypes.push_back(convertedType); |
238 | } |
239 | |
240 | // Create `spirv.selection` operation, selection header block and merge |
241 | // block. |
242 | auto selectionOp = |
243 | rewriter.create<spirv::SelectionOp>(loc, spirv::SelectionControl::None); |
244 | auto *mergeBlock = rewriter.createBlock(&selectionOp.getBody(), |
245 | selectionOp.getBody().end()); |
246 | rewriter.create<spirv::MergeOp>(loc); |
247 | |
248 | OpBuilder::InsertionGuard guard(rewriter); |
249 | auto * = |
250 | rewriter.createBlock(&selectionOp.getBody().front()); |
251 | |
252 | // Inline `then` region before the merge block and branch to it. |
253 | auto &thenRegion = ifOp.getThenRegion(); |
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.getElseRegion().empty()) { |
263 | auto &elseRegion = ifOp.getElseRegion(); |
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 `spirv.BranchConditional` operation for selection header block. |
271 | rewriter.setInsertionPointToEnd(selectionHeaderBlock); |
272 | rewriter.create<spirv::BranchConditionalOp>(loc, adaptor.getCondition(), |
273 | thenBlock, ArrayRef<Value>(), |
274 | elseBlock, ArrayRef<Value>()); |
275 | |
276 | replaceSCFOutputValue(ifOp, selectionOp, rewriter, scfToSPIRVContext, |
277 | returnTypes); |
278 | return success(); |
279 | } |
280 | }; |
281 | |
282 | //===----------------------------------------------------------------------===// |
283 | // scf::YieldOp |
284 | //===----------------------------------------------------------------------===// |
285 | |
286 | struct TerminatorOpConversion final : SCFToSPIRVPattern<scf::YieldOp> { |
287 | public: |
288 | using SCFToSPIRVPattern::SCFToSPIRVPattern; |
289 | |
290 | LogicalResult |
291 | matchAndRewrite(scf::YieldOp terminatorOp, OpAdaptor adaptor, |
292 | ConversionPatternRewriter &rewriter) const override { |
293 | ValueRange operands = adaptor.getOperands(); |
294 | |
295 | Operation *parent = terminatorOp->getParentOp(); |
296 | |
297 | // TODO: Implement conversion for the remaining `scf` ops. |
298 | if (parent->getDialect()->getNamespace() == |
299 | scf::SCFDialect::getDialectNamespace() && |
300 | !isa<scf::IfOp, scf::ForOp, scf::WhileOp>(parent)) |
301 | return rewriter.notifyMatchFailure( |
302 | terminatorOp, |
303 | llvm::formatv(Fmt: "conversion not supported for parent op: '{0}'" , |
304 | Vals: parent->getName())); |
305 | |
306 | // If the region return values, store each value into the associated |
307 | // VariableOp created during lowering of the parent region. |
308 | if (!operands.empty()) { |
309 | auto &allocas = scfToSPIRVContext->outputVars[parent]; |
310 | if (allocas.size() != operands.size()) |
311 | return failure(); |
312 | |
313 | auto loc = terminatorOp.getLoc(); |
314 | for (unsigned i = 0, e = operands.size(); i < e; i++) |
315 | rewriter.create<spirv::StoreOp>(loc, allocas[i], operands[i]); |
316 | if (isa<spirv::LoopOp>(parent)) { |
317 | // For loops we also need to update the branch jumping back to the |
318 | // header. |
319 | auto br = cast<spirv::BranchOp>( |
320 | rewriter.getInsertionBlock()->getTerminator()); |
321 | SmallVector<Value, 8> args(br.getBlockArguments()); |
322 | args.append(in_start: operands.begin(), in_end: operands.end()); |
323 | rewriter.setInsertionPoint(br); |
324 | rewriter.create<spirv::BranchOp>(terminatorOp.getLoc(), br.getTarget(), |
325 | args); |
326 | rewriter.eraseOp(op: br); |
327 | } |
328 | } |
329 | rewriter.eraseOp(op: terminatorOp); |
330 | return success(); |
331 | } |
332 | }; |
333 | |
334 | //===----------------------------------------------------------------------===// |
335 | // scf::WhileOp |
336 | //===----------------------------------------------------------------------===// |
337 | |
338 | struct WhileOpConversion final : SCFToSPIRVPattern<scf::WhileOp> { |
339 | using SCFToSPIRVPattern::SCFToSPIRVPattern; |
340 | |
341 | LogicalResult |
342 | matchAndRewrite(scf::WhileOp whileOp, OpAdaptor adaptor, |
343 | ConversionPatternRewriter &rewriter) const override { |
344 | auto loc = whileOp.getLoc(); |
345 | auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None); |
346 | loopOp.addEntryAndMergeBlock(rewriter); |
347 | |
348 | Region &beforeRegion = whileOp.getBefore(); |
349 | Region &afterRegion = whileOp.getAfter(); |
350 | |
351 | if (failed(rewriter.convertRegionTypes(&beforeRegion, typeConverter)) || |
352 | failed(rewriter.convertRegionTypes(&afterRegion, typeConverter))) |
353 | return rewriter.notifyMatchFailure(whileOp, |
354 | "Failed to convert region types" ); |
355 | |
356 | OpBuilder::InsertionGuard guard(rewriter); |
357 | |
358 | Block &entryBlock = *loopOp.getEntryBlock(); |
359 | Block &beforeBlock = beforeRegion.front(); |
360 | Block &afterBlock = afterRegion.front(); |
361 | Block &mergeBlock = *loopOp.getMergeBlock(); |
362 | |
363 | auto cond = cast<scf::ConditionOp>(beforeBlock.getTerminator()); |
364 | SmallVector<Value> condArgs; |
365 | if (failed(rewriter.getRemappedValues(keys: cond.getArgs(), results&: condArgs))) |
366 | return failure(); |
367 | |
368 | Value conditionVal = rewriter.getRemappedValue(key: cond.getCondition()); |
369 | if (!conditionVal) |
370 | return failure(); |
371 | |
372 | auto yield = cast<scf::YieldOp>(afterBlock.getTerminator()); |
373 | SmallVector<Value> yieldArgs; |
374 | if (failed(rewriter.getRemappedValues(keys: yield.getResults(), results&: yieldArgs))) |
375 | return failure(); |
376 | |
377 | // Move the while before block as the initial loop header block. |
378 | rewriter.inlineRegionBefore(beforeRegion, loopOp.getBody(), |
379 | getBlockIt(loopOp.getBody(), 1)); |
380 | |
381 | // Move the while after block as the initial loop body block. |
382 | rewriter.inlineRegionBefore(afterRegion, loopOp.getBody(), |
383 | getBlockIt(loopOp.getBody(), 2)); |
384 | |
385 | // Jump from the loop entry block to the loop header block. |
386 | rewriter.setInsertionPointToEnd(&entryBlock); |
387 | rewriter.create<spirv::BranchOp>(loc, &beforeBlock, adaptor.getInits()); |
388 | |
389 | auto condLoc = cond.getLoc(); |
390 | |
391 | SmallVector<Value> resultValues(condArgs.size()); |
392 | |
393 | // For other SCF ops, the scf.yield op yields the value for the whole SCF |
394 | // op. So we use the scf.yield op as the anchor to create/load/store SPIR-V |
395 | // local variables. But for the scf.while op, the scf.yield op yields a |
396 | // value for the before region, which may not matching the whole op's |
397 | // result. Instead, the scf.condition op returns values matching the whole |
398 | // op's results. So we need to create/load/store variables according to |
399 | // that. |
400 | for (const auto &it : llvm::enumerate(First&: condArgs)) { |
401 | auto res = it.value(); |
402 | auto i = it.index(); |
403 | auto pointerType = |
404 | spirv::PointerType::get(res.getType(), spirv::StorageClass::Function); |
405 | |
406 | // Create local variables before the scf.while op. |
407 | rewriter.setInsertionPoint(loopOp); |
408 | auto alloc = rewriter.create<spirv::VariableOp>( |
409 | condLoc, pointerType, spirv::StorageClass::Function, |
410 | /*initializer=*/nullptr); |
411 | |
412 | // Load the final result values after the scf.while op. |
413 | rewriter.setInsertionPointAfter(loopOp); |
414 | auto loadResult = rewriter.create<spirv::LoadOp>(condLoc, alloc); |
415 | resultValues[i] = loadResult; |
416 | |
417 | // Store the current iteration's result value. |
418 | rewriter.setInsertionPointToEnd(&beforeBlock); |
419 | rewriter.create<spirv::StoreOp>(condLoc, alloc, res); |
420 | } |
421 | |
422 | rewriter.setInsertionPointToEnd(&beforeBlock); |
423 | rewriter.replaceOpWithNewOp<spirv::BranchConditionalOp>( |
424 | cond, conditionVal, &afterBlock, condArgs, &mergeBlock, std::nullopt); |
425 | |
426 | // Convert the scf.yield op to a branch back to the header block. |
427 | rewriter.setInsertionPointToEnd(&afterBlock); |
428 | rewriter.replaceOpWithNewOp<spirv::BranchOp>(yield, &beforeBlock, |
429 | yieldArgs); |
430 | |
431 | rewriter.replaceOp(whileOp, resultValues); |
432 | return success(); |
433 | } |
434 | }; |
435 | } // namespace |
436 | |
437 | //===----------------------------------------------------------------------===// |
438 | // Public API |
439 | //===----------------------------------------------------------------------===// |
440 | |
441 | void mlir::populateSCFToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, |
442 | ScfToSPIRVContext &scfToSPIRVContext, |
443 | RewritePatternSet &patterns) { |
444 | patterns.add<ForOpConversion, IfOpConversion, TerminatorOpConversion, |
445 | WhileOpConversion>(arg: patterns.getContext(), args: typeConverter, |
446 | args: scfToSPIRVContext.getImpl()); |
447 | } |
448 | |