|
| 1 | +//===- SCFToAffine.cpp - SCF to Affine conversion -------------------------===// |
| 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 raise scf.for, scf.if and loop.terminator |
| 10 | +// ops into affine ops. |
| 11 | +// |
| 12 | +//===----------------------------------------------------------------------===// |
| 13 | + |
| 14 | +#include "mlir/Conversion/SCFToAffine/SCFToAffine.h" |
| 15 | +#include "mlir/Dialect/Affine/IR/AffineOps.h" |
| 16 | +#include "mlir/Dialect/SCF/IR/SCF.h" |
| 17 | +#include "mlir/IR/Verifier.h" |
| 18 | +#include "mlir/Transforms/DialectConversion.h" |
| 19 | +#include "mlir/Transforms/Passes.h" |
| 20 | + |
| 21 | +namespace mlir { |
| 22 | +#define GEN_PASS_DEF_RAISESCFTOAFFINEPASS |
| 23 | +#include "mlir/Conversion/Passes.h.inc" |
| 24 | +} // namespace mlir |
| 25 | + |
| 26 | +using namespace mlir; |
| 27 | + |
| 28 | +namespace { |
| 29 | + |
| 30 | +struct SCFToAffinePass |
| 31 | + : public impl::RaiseSCFToAffinePassBase<SCFToAffinePass> { |
| 32 | + void runOnOperation() override; |
| 33 | +}; |
| 34 | + |
| 35 | +bool canRaiseToAffine(scf::ForOp op) { |
| 36 | + return affine::isValidDim(op.getLowerBound()) && |
| 37 | + affine::isValidDim(op.getUpperBound()) && |
| 38 | + affine::isValidSymbol(op.getStep()); |
| 39 | +} |
| 40 | + |
| 41 | +struct ForOpRewrite : public OpRewritePattern<scf::ForOp> { |
| 42 | + using OpRewritePattern<scf::ForOp>::OpRewritePattern; |
| 43 | + |
| 44 | + std::pair<affine::AffineForOp, Value> |
| 45 | + createAffineFor(scf::ForOp op, PatternRewriter &rewriter) const { |
| 46 | + if ( auto constantStep = op. getStep(). getDefiningOp<arith::ConstantOp>()) { |
| 47 | + int64_t step = cast<IntegerAttr>(constantStep.getValue()).getInt(); |
| 48 | + if (step > 0) |
| 49 | + return positiveConstantStep(op, step, rewriter); |
| 50 | + } |
| 51 | + return genericBounds(op, rewriter); |
| 52 | + } |
| 53 | + |
| 54 | + std::pair<affine::AffineForOp, Value> |
| 55 | + positiveConstantStep(scf::ForOp op, int64_t step, |
| 56 | + PatternRewriter &rewriter) const { |
| 57 | + auto affineFor = affine::AffineForOp::create( |
| 58 | + rewriter, op.getLoc(), ValueRange(op.getLowerBound()), |
| 59 | + AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), |
| 60 | + ValueRange(op.getUpperBound()), |
| 61 | + AffineMap::get(1, 0, rewriter.getAffineDimExpr(0)), step, |
| 62 | + op.getInits()); |
| 63 | + return std::make_pair(affineFor, affineFor.getInductionVar()); |
| 64 | + } |
| 65 | + |
| 66 | + std::pair<affine::AffineForOp, Value> |
| 67 | + genericBounds(scf::ForOp op, PatternRewriter &rewriter) const { |
| 68 | + Value lower = op.getLowerBound(); |
| 69 | + Value upper = op.getUpperBound(); |
| 70 | + Value step = op.getStep(); |
| 71 | + AffineExpr lowerExpr = rewriter.getAffineDimExpr(0); |
| 72 | + AffineExpr upperExpr = rewriter.getAffineDimExpr(1); |
| 73 | + AffineExpr stepExpr = rewriter.getAffineSymbolExpr(0); |
| 74 | + auto affineFor = affine::AffineForOp::create( |
| 75 | + rewriter, op.getLoc(), ValueRange(), rewriter.getConstantAffineMap(0), |
| 76 | + ValueRange({lower, upper, step}), |
| 77 | + AffineMap::get( |
| 78 | + 2, 1, (upperExpr - lowerExpr + stepExpr - 1).floorDiv(stepExpr)), |
| 79 | + 1, op.getInits()); |
| 80 | + |
| 81 | + rewriter.setInsertionPointToStart(affineFor.getBody()); |
| 82 | + auto actualIndexMap = AffineMap::get( |
| 83 | + 2, 1, lowerExpr + rewriter.getAffineDimExpr(1) * stepExpr); |
| 84 | + auto actualIndex = affine::AffineApplyOp::create( |
| 85 | + rewriter, op.getLoc(), actualIndexMap, |
| 86 | + ValueRange({lower, affineFor.getInductionVar(), step})); |
| 87 | + return std::make_pair(affineFor, actualIndex.getResult()); |
| 88 | + } |
| 89 | + |
| 90 | + LogicalResult matchAndRewrite(scf::ForOp op, |
| 91 | + PatternRewriter &rewriter) const override { |
| 92 | + if (!canRaiseToAffine(op)) |
| 93 | + return failure(); |
| 94 | + |
| 95 | + auto [affineFor, actualIndex] = createAffineFor(op, rewriter); |
| 96 | + Block *affineBody = affineFor.getBody(); |
| 97 | + |
| 98 | + if (affineBody->mightHaveTerminator()) |
| 99 | + rewriter.eraseOp(affineBody->getTerminator()); |
| 100 | + |
| 101 | + SmallVector<Value> argValues; |
| 102 | + argValues.push_back(actualIndex); |
| 103 | + llvm::append_range(argValues, affineFor.getRegionIterArgs()); |
| 104 | + rewriter.inlineBlockBefore(op.getBody(), affineBody, affineBody->end(), |
| 105 | + argValues); |
| 106 | + |
| 107 | + auto scfYieldOp = cast<scf::YieldOp>(affineBody->getTerminator()); |
| 108 | + rewriter.setInsertionPointToEnd(affineBody); |
| 109 | + rewriter. replaceOpWithNewOp<affine::AffineYieldOp>( |
| 110 | + scfYieldOp, scfYieldOp->getOperands()); |
| 111 | + |
| 112 | + rewriter.replaceOp(op, affineFor); |
| 113 | + return success(); |
| 114 | + } |
| 115 | +}; |
| 116 | + |
| 117 | +} // namespace |
| 118 | + |
| 119 | +void mlir::populateSCFToAffineConversionPatterns(RewritePatternSet &patterns) { |
| 120 | + patterns.add<ForOpRewrite>(patterns.getContext()); |
| 121 | +} |
| 122 | + |
| 123 | +void SCFToAffinePass::runOnOperation() { |
| 124 | + MLIRContext &ctx = getContext(); |
| 125 | + RewritePatternSet patterns(&ctx); |
| 126 | + populateSCFToAffineConversionPatterns(patterns); |
| 127 | + |
| 128 | + // Configure conversion to raise SCF operations. |
| 129 | + ConversionTarget target(ctx); |
| 130 | + target.addDynamicallyLegalOp<scf::ForOp>( |
| 131 | + [](scf::ForOp op) { return !canRaiseToAffine(op); }); |
| 132 | + target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); |
| 133 | + if (failed( |
| 134 | + applyPartialConversion(getOperation(), target, std::move(patterns)))) |
| 135 | + signalPassFailure(); |
| 136 | +} |
0 commit comments