570 lines · cpp
1//===- AffineToStandard.cpp - Lower affine constructs to primitives -------===//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-exception6//7//===----------------------------------------------------------------------===//8//9// This file lowers affine constructs (If and For statements, AffineApply10// operations) within a function into their standard If and For equivalent ops.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"15 16#include "mlir/Dialect/Affine/IR/AffineOps.h"17#include "mlir/Dialect/Affine/Transforms/Transforms.h"18#include "mlir/Dialect/Affine/Utils.h"19#include "mlir/Dialect/MemRef/IR/MemRef.h"20#include "mlir/Dialect/SCF/IR/SCF.h"21#include "mlir/Dialect/Vector/IR/VectorOps.h"22#include "mlir/IR/IntegerSet.h"23#include "mlir/IR/MLIRContext.h"24#include "mlir/Transforms/DialectConversion.h"25#include "mlir/Transforms/Passes.h"26 27namespace mlir {28#define GEN_PASS_DEF_LOWERAFFINEPASS29#include "mlir/Conversion/Passes.h.inc"30} // namespace mlir31 32using namespace mlir;33using namespace mlir::affine;34using namespace mlir::vector;35 36/// Given a range of values, emit the code that reduces them with "min" or "max"37/// depending on the provided comparison predicate, sgt for max and slt for min.38///39/// Multiple values are scanned in a linear sequence. This creates a data40/// dependences that wouldn't exist in a tree reduction, but is easier to41/// recognize as a reduction by the subsequent passes.42static Value buildMinMaxReductionSeq(Location loc,43 arith::CmpIPredicate predicate,44 ValueRange values, OpBuilder &builder) {45 assert(!values.empty() && "empty min/max chain");46 assert(predicate == arith::CmpIPredicate::sgt ||47 predicate == arith::CmpIPredicate::slt);48 49 auto valueIt = values.begin();50 Value value = *valueIt++;51 for (; valueIt != values.end(); ++valueIt) {52 if (predicate == arith::CmpIPredicate::sgt)53 value = arith::MaxSIOp::create(builder, loc, value, *valueIt);54 else55 value = arith::MinSIOp::create(builder, loc, value, *valueIt);56 }57 58 return value;59}60 61/// Emit instructions that correspond to computing the maximum value among the62/// values of a (potentially) multi-output affine map applied to `operands`.63static Value lowerAffineMapMax(OpBuilder &builder, Location loc, AffineMap map,64 ValueRange operands) {65 if (auto values = expandAffineMap(builder, loc, map, operands))66 return buildMinMaxReductionSeq(loc, arith::CmpIPredicate::sgt, *values,67 builder);68 return nullptr;69}70 71/// Emit instructions that correspond to computing the minimum value among the72/// values of a (potentially) multi-output affine map applied to `operands`.73static Value lowerAffineMapMin(OpBuilder &builder, Location loc, AffineMap map,74 ValueRange operands) {75 if (auto values = expandAffineMap(builder, loc, map, operands))76 return buildMinMaxReductionSeq(loc, arith::CmpIPredicate::slt, *values,77 builder);78 return nullptr;79}80 81/// Emit instructions that correspond to the affine map in the upper bound82/// applied to the respective operands, and compute the minimum value across83/// the results.84Value mlir::lowerAffineUpperBound(AffineForOp op, OpBuilder &builder) {85 return lowerAffineMapMin(builder, op.getLoc(), op.getUpperBoundMap(),86 op.getUpperBoundOperands());87}88 89/// Emit instructions that correspond to the affine map in the lower bound90/// applied to the respective operands, and compute the maximum value across91/// the results.92Value mlir::lowerAffineLowerBound(AffineForOp op, OpBuilder &builder) {93 return lowerAffineMapMax(builder, op.getLoc(), op.getLowerBoundMap(),94 op.getLowerBoundOperands());95}96 97namespace {98class AffineMinLowering : public OpRewritePattern<AffineMinOp> {99public:100 using OpRewritePattern<AffineMinOp>::OpRewritePattern;101 102 LogicalResult matchAndRewrite(AffineMinOp op,103 PatternRewriter &rewriter) const override {104 Value reduced =105 lowerAffineMapMin(rewriter, op.getLoc(), op.getMap(), op.getOperands());106 if (!reduced)107 return failure();108 109 rewriter.replaceOp(op, reduced);110 return success();111 }112};113 114class AffineMaxLowering : public OpRewritePattern<AffineMaxOp> {115public:116 using OpRewritePattern<AffineMaxOp>::OpRewritePattern;117 118 LogicalResult matchAndRewrite(AffineMaxOp op,119 PatternRewriter &rewriter) const override {120 Value reduced =121 lowerAffineMapMax(rewriter, op.getLoc(), op.getMap(), op.getOperands());122 if (!reduced)123 return failure();124 125 rewriter.replaceOp(op, reduced);126 return success();127 }128};129 130/// Affine yields ops are removed.131class AffineYieldOpLowering : public OpRewritePattern<AffineYieldOp> {132public:133 using OpRewritePattern<AffineYieldOp>::OpRewritePattern;134 135 LogicalResult matchAndRewrite(AffineYieldOp op,136 PatternRewriter &rewriter) const override {137 if (isa<scf::ParallelOp>(op->getParentOp())) {138 // Terminator is rewritten as part of the "affine.parallel" lowering139 // pattern.140 return failure();141 }142 rewriter.replaceOpWithNewOp<scf::YieldOp>(op, op.getOperands());143 return success();144 }145};146 147class AffineForLowering : public OpRewritePattern<AffineForOp> {148public:149 using OpRewritePattern<AffineForOp>::OpRewritePattern;150 151 LogicalResult matchAndRewrite(AffineForOp op,152 PatternRewriter &rewriter) const override {153 Location loc = op.getLoc();154 Value lowerBound = lowerAffineLowerBound(op, rewriter);155 Value upperBound = lowerAffineUpperBound(op, rewriter);156 Value step =157 arith::ConstantIndexOp::create(rewriter, loc, op.getStepAsInt());158 auto scfForOp = scf::ForOp::create(rewriter, loc, lowerBound, upperBound,159 step, op.getInits());160 rewriter.eraseBlock(scfForOp.getBody());161 rewriter.inlineRegionBefore(op.getRegion(), scfForOp.getRegion(),162 scfForOp.getRegion().end());163 rewriter.replaceOp(op, scfForOp.getResults());164 return success();165 }166};167 168/// Convert an `affine.parallel` (loop nest) operation into a `scf.parallel`169/// operation.170class AffineParallelLowering : public OpRewritePattern<AffineParallelOp> {171public:172 using OpRewritePattern<AffineParallelOp>::OpRewritePattern;173 174 LogicalResult matchAndRewrite(AffineParallelOp op,175 PatternRewriter &rewriter) const override {176 Location loc = op.getLoc();177 SmallVector<Value, 8> steps;178 SmallVector<Value, 8> upperBoundTuple;179 SmallVector<Value, 8> lowerBoundTuple;180 SmallVector<Value, 8> identityVals;181 // Emit IR computing the lower and upper bound by expanding the map182 // expression.183 lowerBoundTuple.reserve(op.getNumDims());184 upperBoundTuple.reserve(op.getNumDims());185 for (unsigned i = 0, e = op.getNumDims(); i < e; ++i) {186 Value lower = lowerAffineMapMax(rewriter, loc, op.getLowerBoundMap(i),187 op.getLowerBoundsOperands());188 if (!lower)189 return rewriter.notifyMatchFailure(op, "couldn't convert lower bounds");190 lowerBoundTuple.push_back(lower);191 192 Value upper = lowerAffineMapMin(rewriter, loc, op.getUpperBoundMap(i),193 op.getUpperBoundsOperands());194 if (!upper)195 return rewriter.notifyMatchFailure(op, "couldn't convert upper bounds");196 upperBoundTuple.push_back(upper);197 }198 steps.reserve(op.getSteps().size());199 for (int64_t step : op.getSteps())200 steps.push_back(arith::ConstantIndexOp::create(rewriter, loc, step));201 202 // Get the terminator op.203 auto affineParOpTerminator =204 cast<AffineYieldOp>(op.getBody()->getTerminator());205 scf::ParallelOp parOp;206 if (op.getResults().empty()) {207 // Case with no reduction operations/return values.208 parOp = scf::ParallelOp::create(rewriter, loc, lowerBoundTuple,209 upperBoundTuple, steps,210 /*bodyBuilderFn=*/nullptr);211 rewriter.eraseBlock(parOp.getBody());212 rewriter.inlineRegionBefore(op.getRegion(), parOp.getRegion(),213 parOp.getRegion().end());214 rewriter.replaceOp(op, parOp.getResults());215 rewriter.setInsertionPoint(affineParOpTerminator);216 rewriter.replaceOpWithNewOp<scf::ReduceOp>(affineParOpTerminator);217 return success();218 }219 // Case with affine.parallel with reduction operations/return values.220 // scf.parallel handles the reduction operation differently unlike221 // affine.parallel.222 ArrayRef<Attribute> reductions = op.getReductions().getValue();223 for (auto pair : llvm::zip(reductions, op.getResultTypes())) {224 // For each of the reduction operations get the identity values for225 // initialization of the result values.226 Attribute reduction = std::get<0>(pair);227 Type resultType = std::get<1>(pair);228 std::optional<arith::AtomicRMWKind> reductionOp =229 arith::symbolizeAtomicRMWKind(230 static_cast<uint64_t>(cast<IntegerAttr>(reduction).getInt()));231 assert(reductionOp && "Reduction operation cannot be of None Type");232 arith::AtomicRMWKind reductionOpValue = *reductionOp;233 identityVals.push_back(234 arith::getIdentityValue(reductionOpValue, resultType, rewriter, loc));235 }236 parOp = scf::ParallelOp::create(rewriter, loc, lowerBoundTuple,237 upperBoundTuple, steps, identityVals,238 /*bodyBuilderFn=*/nullptr);239 240 // Copy the body of the affine.parallel op.241 rewriter.eraseBlock(parOp.getBody());242 rewriter.inlineRegionBefore(op.getRegion(), parOp.getRegion(),243 parOp.getRegion().end());244 assert(reductions.size() == affineParOpTerminator->getNumOperands() &&245 "Unequal number of reductions and operands.");246 247 // Emit new "scf.reduce" terminator.248 rewriter.setInsertionPoint(affineParOpTerminator);249 auto reduceOp = rewriter.replaceOpWithNewOp<scf::ReduceOp>(250 affineParOpTerminator, affineParOpTerminator->getOperands());251 for (unsigned i = 0, end = reductions.size(); i < end; i++) {252 // For each of the reduction operations get the respective mlir::Value.253 std::optional<arith::AtomicRMWKind> reductionOp =254 arith::symbolizeAtomicRMWKind(255 cast<IntegerAttr>(reductions[i]).getInt());256 assert(reductionOp && "Reduction Operation cannot be of None Type");257 arith::AtomicRMWKind reductionOpValue = *reductionOp;258 rewriter.setInsertionPoint(&parOp.getBody()->back());259 Block &reductionBody = reduceOp.getReductions()[i].front();260 rewriter.setInsertionPointToEnd(&reductionBody);261 Value reductionResult = arith::getReductionOp(262 reductionOpValue, rewriter, loc, reductionBody.getArgument(0),263 reductionBody.getArgument(1));264 scf::ReduceReturnOp::create(rewriter, loc, reductionResult);265 }266 rewriter.replaceOp(op, parOp.getResults());267 return success();268 }269};270 271class AffineIfLowering : public OpRewritePattern<AffineIfOp> {272public:273 using OpRewritePattern<AffineIfOp>::OpRewritePattern;274 275 LogicalResult matchAndRewrite(AffineIfOp op,276 PatternRewriter &rewriter) const override {277 auto loc = op.getLoc();278 279 // Now we just have to handle the condition logic.280 auto integerSet = op.getIntegerSet();281 Value zeroConstant = arith::ConstantIndexOp::create(rewriter, loc, 0);282 SmallVector<Value, 8> operands(op.getOperands());283 auto operandsRef = llvm::ArrayRef(operands);284 285 // Calculate cond as a conjunction without short-circuiting.286 Value cond = nullptr;287 for (unsigned i = 0, e = integerSet.getNumConstraints(); i < e; ++i) {288 AffineExpr constraintExpr = integerSet.getConstraint(i);289 bool isEquality = integerSet.isEq(i);290 291 // Build and apply an affine expression292 auto numDims = integerSet.getNumDims();293 Value affResult = expandAffineExpr(rewriter, loc, constraintExpr,294 operandsRef.take_front(numDims),295 operandsRef.drop_front(numDims));296 if (!affResult)297 return failure();298 auto pred =299 isEquality ? arith::CmpIPredicate::eq : arith::CmpIPredicate::sge;300 Value cmpVal =301 arith::CmpIOp::create(rewriter, loc, pred, affResult, zeroConstant);302 cond =303 cond ? arith::AndIOp::create(rewriter, loc, cond, cmpVal).getResult()304 : cmpVal;305 }306 cond = cond ? cond307 : arith::ConstantIntOp::create(rewriter, loc, /*value=*/1,308 /*width=*/1);309 310 bool hasElseRegion = !op.getElseRegion().empty();311 auto ifOp = scf::IfOp::create(rewriter, loc, op.getResultTypes(), cond,312 hasElseRegion);313 rewriter.inlineRegionBefore(op.getThenRegion(),314 &ifOp.getThenRegion().back());315 rewriter.eraseBlock(&ifOp.getThenRegion().back());316 if (hasElseRegion) {317 rewriter.inlineRegionBefore(op.getElseRegion(),318 &ifOp.getElseRegion().back());319 rewriter.eraseBlock(&ifOp.getElseRegion().back());320 }321 322 // Replace the Affine IfOp finally.323 rewriter.replaceOp(op, ifOp.getResults());324 return success();325 }326};327 328/// Convert an "affine.apply" operation into a sequence of arithmetic329/// operations using the StandardOps dialect.330class AffineApplyLowering : public OpRewritePattern<AffineApplyOp> {331public:332 using OpRewritePattern<AffineApplyOp>::OpRewritePattern;333 334 LogicalResult matchAndRewrite(AffineApplyOp op,335 PatternRewriter &rewriter) const override {336 auto maybeExpandedMap =337 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(),338 llvm::to_vector<8>(op.getOperands()));339 if (!maybeExpandedMap)340 return failure();341 rewriter.replaceOp(op, *maybeExpandedMap);342 return success();343 }344};345 346/// Apply the affine map from an 'affine.load' operation to its operands, and347/// feed the results to a newly created 'memref.load' operation (which replaces348/// the original 'affine.load').349class AffineLoadLowering : public OpRewritePattern<AffineLoadOp> {350public:351 using OpRewritePattern<AffineLoadOp>::OpRewritePattern;352 353 LogicalResult matchAndRewrite(AffineLoadOp op,354 PatternRewriter &rewriter) const override {355 // Expand affine map from 'affineLoadOp'.356 SmallVector<Value, 8> indices(op.getMapOperands());357 auto resultOperands =358 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);359 if (!resultOperands)360 return failure();361 362 // Build vector.load memref[expandedMap.results].363 rewriter.replaceOpWithNewOp<memref::LoadOp>(op, op.getMemRef(),364 *resultOperands);365 return success();366 }367};368 369/// Apply the affine map from an 'affine.prefetch' operation to its operands,370/// and feed the results to a newly created 'memref.prefetch' operation (which371/// replaces the original 'affine.prefetch').372class AffinePrefetchLowering : public OpRewritePattern<AffinePrefetchOp> {373public:374 using OpRewritePattern<AffinePrefetchOp>::OpRewritePattern;375 376 LogicalResult matchAndRewrite(AffinePrefetchOp op,377 PatternRewriter &rewriter) const override {378 // Expand affine map from 'affinePrefetchOp'.379 SmallVector<Value, 8> indices(op.getMapOperands());380 auto resultOperands =381 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);382 if (!resultOperands)383 return failure();384 385 // Build memref.prefetch memref[expandedMap.results].386 rewriter.replaceOpWithNewOp<memref::PrefetchOp>(387 op, op.getMemref(), *resultOperands, op.getIsWrite(),388 op.getLocalityHint(), op.getIsDataCache());389 return success();390 }391};392 393/// Apply the affine map from an 'affine.store' operation to its operands, and394/// feed the results to a newly created 'memref.store' operation (which replaces395/// the original 'affine.store').396class AffineStoreLowering : public OpRewritePattern<AffineStoreOp> {397public:398 using OpRewritePattern<AffineStoreOp>::OpRewritePattern;399 400 LogicalResult matchAndRewrite(AffineStoreOp op,401 PatternRewriter &rewriter) const override {402 // Expand affine map from 'affineStoreOp'.403 SmallVector<Value, 8> indices(op.getMapOperands());404 auto maybeExpandedMap =405 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);406 if (!maybeExpandedMap)407 return failure();408 409 // Build memref.store valueToStore, memref[expandedMap.results].410 rewriter.replaceOpWithNewOp<memref::StoreOp>(411 op, op.getValueToStore(), op.getMemRef(), *maybeExpandedMap);412 return success();413 }414};415 416/// Apply the affine maps from an 'affine.dma_start' operation to each of their417/// respective map operands, and feed the results to a newly created418/// 'memref.dma_start' operation (which replaces the original419/// 'affine.dma_start').420class AffineDmaStartLowering : public OpRewritePattern<AffineDmaStartOp> {421public:422 using OpRewritePattern<AffineDmaStartOp>::OpRewritePattern;423 424 LogicalResult matchAndRewrite(AffineDmaStartOp op,425 PatternRewriter &rewriter) const override {426 SmallVector<Value, 8> operands(op.getOperands());427 auto operandsRef = llvm::ArrayRef(operands);428 429 // Expand affine map for DMA source memref.430 auto maybeExpandedSrcMap = expandAffineMap(431 rewriter, op.getLoc(), op.getSrcMap(),432 operandsRef.drop_front(op.getSrcMemRefOperandIndex() + 1));433 if (!maybeExpandedSrcMap)434 return failure();435 // Expand affine map for DMA destination memref.436 auto maybeExpandedDstMap = expandAffineMap(437 rewriter, op.getLoc(), op.getDstMap(),438 operandsRef.drop_front(op.getDstMemRefOperandIndex() + 1));439 if (!maybeExpandedDstMap)440 return failure();441 // Expand affine map for DMA tag memref.442 auto maybeExpandedTagMap = expandAffineMap(443 rewriter, op.getLoc(), op.getTagMap(),444 operandsRef.drop_front(op.getTagMemRefOperandIndex() + 1));445 if (!maybeExpandedTagMap)446 return failure();447 448 // Build memref.dma_start operation with affine map results.449 rewriter.replaceOpWithNewOp<memref::DmaStartOp>(450 op, op.getSrcMemRef(), *maybeExpandedSrcMap, op.getDstMemRef(),451 *maybeExpandedDstMap, op.getNumElements(), op.getTagMemRef(),452 *maybeExpandedTagMap, op.getStride(), op.getNumElementsPerStride());453 return success();454 }455};456 457/// Apply the affine map from an 'affine.dma_wait' operation tag memref,458/// and feed the results to a newly created 'memref.dma_wait' operation (which459/// replaces the original 'affine.dma_wait').460class AffineDmaWaitLowering : public OpRewritePattern<AffineDmaWaitOp> {461public:462 using OpRewritePattern<AffineDmaWaitOp>::OpRewritePattern;463 464 LogicalResult matchAndRewrite(AffineDmaWaitOp op,465 PatternRewriter &rewriter) const override {466 // Expand affine map for DMA tag memref.467 SmallVector<Value, 8> indices(op.getTagIndices());468 auto maybeExpandedTagMap =469 expandAffineMap(rewriter, op.getLoc(), op.getTagMap(), indices);470 if (!maybeExpandedTagMap)471 return failure();472 473 // Build memref.dma_wait operation with affine map results.474 rewriter.replaceOpWithNewOp<memref::DmaWaitOp>(475 op, op.getTagMemRef(), *maybeExpandedTagMap, op.getNumElements());476 return success();477 }478};479 480/// Apply the affine map from an 'affine.vector_load' operation to its operands,481/// and feed the results to a newly created 'vector.load' operation (which482/// replaces the original 'affine.vector_load').483class AffineVectorLoadLowering : public OpRewritePattern<AffineVectorLoadOp> {484public:485 using OpRewritePattern<AffineVectorLoadOp>::OpRewritePattern;486 487 LogicalResult matchAndRewrite(AffineVectorLoadOp op,488 PatternRewriter &rewriter) const override {489 // Expand affine map from 'affineVectorLoadOp'.490 SmallVector<Value, 8> indices(op.getMapOperands());491 auto resultOperands =492 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);493 if (!resultOperands)494 return failure();495 496 // Build vector.load memref[expandedMap.results].497 rewriter.replaceOpWithNewOp<vector::LoadOp>(498 op, op.getVectorType(), op.getMemRef(), *resultOperands);499 return success();500 }501};502 503/// Apply the affine map from an 'affine.vector_store' operation to its504/// operands, and feed the results to a newly created 'vector.store' operation505/// (which replaces the original 'affine.vector_store').506class AffineVectorStoreLowering : public OpRewritePattern<AffineVectorStoreOp> {507public:508 using OpRewritePattern<AffineVectorStoreOp>::OpRewritePattern;509 510 LogicalResult matchAndRewrite(AffineVectorStoreOp op,511 PatternRewriter &rewriter) const override {512 // Expand affine map from 'affineVectorStoreOp'.513 SmallVector<Value, 8> indices(op.getMapOperands());514 auto maybeExpandedMap =515 expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);516 if (!maybeExpandedMap)517 return failure();518 519 rewriter.replaceOpWithNewOp<vector::StoreOp>(520 op, op.getValueToStore(), op.getMemRef(), *maybeExpandedMap);521 return success();522 }523};524 525} // namespace526 527void mlir::populateAffineToStdConversionPatterns(RewritePatternSet &patterns) {528 // clang-format off529 patterns.add<530 AffineApplyLowering,531 AffineDmaStartLowering,532 AffineDmaWaitLowering,533 AffineLoadLowering,534 AffineMinLowering,535 AffineMaxLowering,536 AffineParallelLowering,537 AffinePrefetchLowering,538 AffineStoreLowering,539 AffineForLowering,540 AffineIfLowering,541 AffineYieldOpLowering>(patterns.getContext());542 // clang-format on543}544 545void mlir::populateAffineToVectorConversionPatterns(546 RewritePatternSet &patterns) {547 // clang-format off548 patterns.add<549 AffineVectorLoadLowering,550 AffineVectorStoreLowering>(patterns.getContext());551 // clang-format on552}553 554namespace {555class LowerAffine : public impl::LowerAffinePassBase<LowerAffine> {556 void runOnOperation() override {557 RewritePatternSet patterns(&getContext());558 populateAffineToStdConversionPatterns(patterns);559 populateAffineToVectorConversionPatterns(patterns);560 populateAffineExpandIndexOpsPatterns(patterns);561 ConversionTarget target(getContext());562 target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect,563 scf::SCFDialect, VectorDialect>();564 if (failed(applyPartialConversion(getOperation(), target,565 std::move(patterns))))566 signalPassFailure();567 }568};569} // namespace570