649 lines · cpp
1//===-- AffinePromotion.cpp -----------------------------------------------===//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 transformation is a prototype that promote FIR loops operations10// to affine dialect operations.11// It is not part of the production pipeline and would need more work in order12// to be used in production.13// More information can be found in this presentation:14// https://slides.com/rajanwalia/deck15//16//===----------------------------------------------------------------------===//17 18#include "flang/Optimizer/Dialect/FIRDialect.h"19#include "flang/Optimizer/Dialect/FIROps.h"20#include "flang/Optimizer/Dialect/FIRType.h"21#include "flang/Optimizer/Transforms/Passes.h"22#include "mlir/Dialect/Affine/IR/AffineOps.h"23#include "mlir/Dialect/Func/IR/FuncOps.h"24#include "mlir/Dialect/SCF/IR/SCF.h"25#include "mlir/IR/BuiltinAttributes.h"26#include "mlir/IR/IntegerSet.h"27#include "mlir/IR/Visitors.h"28#include "mlir/Transforms/WalkPatternRewriteDriver.h"29#include "llvm/ADT/DenseMap.h"30#include "llvm/Support/Debug.h"31#include <optional>32 33namespace fir {34#define GEN_PASS_DEF_AFFINEDIALECTPROMOTION35#include "flang/Optimizer/Transforms/Passes.h.inc"36} // namespace fir37 38#define DEBUG_TYPE "flang-affine-promotion"39 40using namespace fir;41using namespace mlir;42 43namespace {44struct AffineLoopAnalysis;45struct AffineIfAnalysis;46 47/// Stores analysis objects for all loops and if operations inside a function48/// these analysis are used twice, first for marking operations for rewrite and49/// second when doing rewrite.50struct AffineFunctionAnalysis {51 explicit AffineFunctionAnalysis(mlir::func::FuncOp funcOp) {52 funcOp->walk([&](fir::DoLoopOp doloop) {53 loopAnalysisMap.try_emplace(doloop, doloop, *this);54 });55 }56 57 AffineLoopAnalysis getChildLoopAnalysis(fir::DoLoopOp op) const;58 59 AffineIfAnalysis getChildIfAnalysis(fir::IfOp op) const;60 61 llvm::DenseMap<mlir::Operation *, AffineLoopAnalysis> loopAnalysisMap;62 llvm::DenseMap<mlir::Operation *, AffineIfAnalysis> ifAnalysisMap;63};64} // namespace65 66static bool analyzeCoordinate(mlir::Value coordinate, mlir::Operation *op) {67 if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(coordinate)) {68 if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()))69 return true;70 LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a "71 "loop induction variable (owner not loopOp)\n";72 op->dump());73 return false;74 }75 LLVM_DEBUG(76 llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a loop "77 "induction variable (not a block argument)\n";78 op->dump(); coordinate.getDefiningOp()->dump());79 return false;80}81 82namespace {83struct AffineLoopAnalysis {84 AffineLoopAnalysis() = default;85 86 explicit AffineLoopAnalysis(fir::DoLoopOp op, AffineFunctionAnalysis &afa)87 : legality(analyzeLoop(op, afa)) {}88 89 bool canPromoteToAffine() { return legality; }90 91private:92 bool analyzeBody(fir::DoLoopOp loopOperation,93 AffineFunctionAnalysis &functionAnalysis) {94 for (auto loopOp : loopOperation.getOps<fir::DoLoopOp>()) {95 auto analysis = functionAnalysis.loopAnalysisMap96 .try_emplace(loopOp, loopOp, functionAnalysis)97 .first->getSecond();98 if (!analysis.canPromoteToAffine())99 return false;100 }101 for (auto ifOp : loopOperation.getOps<fir::IfOp>())102 functionAnalysis.ifAnalysisMap.try_emplace(ifOp, ifOp, functionAnalysis);103 return true;104 }105 106 bool analysisResults(fir::DoLoopOp loopOperation) {107 if (loopOperation.getFinalValue() &&108 !loopOperation.getResult(0).use_empty()) {109 LLVM_DEBUG(110 llvm::dbgs()111 << "AffineLoopAnalysis: cannot promote loop final value\n";);112 return false;113 }114 115 return true;116 }117 118 bool analyzeLoop(fir::DoLoopOp loopOperation,119 AffineFunctionAnalysis &functionAnalysis) {120 LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: \n"; loopOperation.dump(););121 return analyzeMemoryAccess(loopOperation) &&122 analysisResults(loopOperation) &&123 analyzeBody(loopOperation, functionAnalysis);124 }125 126 bool analyzeReference(mlir::Value memref, mlir::Operation *op) {127 if (auto acoOp = memref.getDefiningOp<ArrayCoorOp>()) {128 if (mlir::isa<fir::BoxType>(acoOp.getMemref().getType())) {129 // TODO: Look if and how fir.box can be promoted to affine.130 LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: cannot promote loop, "131 "array memory operation uses fir.box\n";132 op->dump(); acoOp.dump(););133 return false;134 }135 bool canPromote = true;136 for (auto coordinate : acoOp.getIndices())137 canPromote = canPromote && analyzeCoordinate(coordinate, op);138 return canPromote;139 }140 if (auto coOp = memref.getDefiningOp<CoordinateOp>()) {141 LLVM_DEBUG(llvm::dbgs()142 << "AffineLoopAnalysis: cannot promote loop, "143 "array memory operation uses non ArrayCoorOp\n";144 op->dump(); coOp.dump(););145 146 return false;147 }148 LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: unknown type of memory "149 "reference for array load\n";150 op->dump(););151 return false;152 }153 154 bool analyzeMemoryAccess(fir::DoLoopOp loopOperation) {155 for (auto loadOp : loopOperation.getOps<fir::LoadOp>())156 if (!analyzeReference(loadOp.getMemref(), loadOp))157 return false;158 for (auto storeOp : loopOperation.getOps<fir::StoreOp>())159 if (!analyzeReference(storeOp.getMemref(), storeOp))160 return false;161 return true;162 }163 164 bool legality{};165};166} // namespace167 168AffineLoopAnalysis169AffineFunctionAnalysis::getChildLoopAnalysis(fir::DoLoopOp op) const {170 auto it = loopAnalysisMap.find_as(op);171 if (it == loopAnalysisMap.end()) {172 LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n";173 op.dump(););174 op.emitError("error in fetching loop analysis in AffineFunctionAnalysis\n");175 return {};176 }177 return it->getSecond();178}179 180namespace {181/// Calculates arguments for creating an IntegerSet. symCount, dimCount are the182/// final number of symbols and dimensions of the affine map. Integer set if183/// possible is in Optional IntegerSet.184struct AffineIfCondition {185 using MaybeAffineExpr = std::optional<mlir::AffineExpr>;186 187 explicit AffineIfCondition(mlir::Value fc) : firCondition(fc) {188 if (auto condDef = firCondition.getDefiningOp<mlir::arith::CmpIOp>())189 fromCmpIOp(condDef);190 }191 192 bool hasIntegerSet() const { return integerSet.has_value(); }193 194 mlir::IntegerSet getIntegerSet() const {195 assert(hasIntegerSet() && "integer set is missing");196 return *integerSet;197 }198 199 mlir::ValueRange getAffineArgs() const { return affineArgs; }200 201private:202 MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, mlir::Value lhs,203 mlir::Value rhs) {204 return affineBinaryOp(kind, toAffineExpr(lhs), toAffineExpr(rhs));205 }206 207 MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, MaybeAffineExpr lhs,208 MaybeAffineExpr rhs) {209 if (lhs && rhs)210 return mlir::getAffineBinaryOpExpr(kind, *lhs, *rhs);211 return {};212 }213 214 MaybeAffineExpr toAffineExpr(MaybeAffineExpr e) { return e; }215 216 MaybeAffineExpr toAffineExpr(int64_t value) {217 return {mlir::getAffineConstantExpr(value, firCondition.getContext())};218 }219 220 /// Returns an AffineExpr if it is a result of operations that can be done221 /// in an affine expression, this includes -, +, *, rem, constant.222 /// block arguments of a loopOp or forOp are used as dimensions223 MaybeAffineExpr toAffineExpr(mlir::Value value) {224 if (auto op = value.getDefiningOp<mlir::arith::SubIOp>())225 return affineBinaryOp(226 mlir::AffineExprKind::Add, toAffineExpr(op.getLhs()),227 affineBinaryOp(mlir::AffineExprKind::Mul, toAffineExpr(op.getRhs()),228 toAffineExpr(-1)));229 if (auto op = value.getDefiningOp<mlir::arith::AddIOp>())230 return affineBinaryOp(mlir::AffineExprKind::Add, op.getLhs(),231 op.getRhs());232 if (auto op = value.getDefiningOp<mlir::arith::MulIOp>())233 return affineBinaryOp(mlir::AffineExprKind::Mul, op.getLhs(),234 op.getRhs());235 if (auto op = value.getDefiningOp<mlir::arith::RemUIOp>())236 return affineBinaryOp(mlir::AffineExprKind::Mod, op.getLhs(),237 op.getRhs());238 if (auto op = value.getDefiningOp<mlir::arith::ConstantOp>())239 if (auto intConstant = mlir::dyn_cast<IntegerAttr>(op.getValue()))240 return toAffineExpr(intConstant.getInt());241 if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(value)) {242 affineArgs.push_back(value);243 if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()) ||244 isa<mlir::affine::AffineForOp>(blockArg.getOwner()->getParentOp()))245 return {mlir::getAffineDimExpr(dimCount++, value.getContext())};246 return {mlir::getAffineSymbolExpr(symCount++, value.getContext())};247 }248 return {};249 }250 251 void fromCmpIOp(mlir::arith::CmpIOp cmpOp) {252 auto lhsAffine = toAffineExpr(cmpOp.getLhs());253 auto rhsAffine = toAffineExpr(cmpOp.getRhs());254 if (!lhsAffine || !rhsAffine)255 return;256 auto constraintPair =257 constraint(cmpOp.getPredicate(), *rhsAffine - *lhsAffine);258 if (!constraintPair)259 return;260 integerSet = mlir::IntegerSet::get(261 dimCount, symCount, {constraintPair->first}, {constraintPair->second});262 }263 264 std::optional<std::pair<AffineExpr, bool>>265 constraint(mlir::arith::CmpIPredicate predicate, mlir::AffineExpr basic) {266 switch (predicate) {267 case mlir::arith::CmpIPredicate::slt:268 return {std::make_pair(basic - 1, false)};269 case mlir::arith::CmpIPredicate::sle:270 return {std::make_pair(basic, false)};271 case mlir::arith::CmpIPredicate::sgt:272 return {std::make_pair(1 - basic, false)};273 case mlir::arith::CmpIPredicate::sge:274 return {std::make_pair(0 - basic, false)};275 case mlir::arith::CmpIPredicate::eq:276 return {std::make_pair(basic, true)};277 default:278 return {};279 }280 }281 282 llvm::SmallVector<mlir::Value> affineArgs;283 std::optional<mlir::IntegerSet> integerSet;284 mlir::Value firCondition;285 unsigned symCount{0u};286 unsigned dimCount{0u};287};288} // namespace289 290namespace {291/// Analysis for affine promotion of fir.if292struct AffineIfAnalysis {293 AffineIfAnalysis() = default;294 295 explicit AffineIfAnalysis(fir::IfOp op, AffineFunctionAnalysis &afa)296 : legality(analyzeIf(op, afa)) {}297 298 bool canPromoteToAffine() { return legality; }299 300private:301 bool analyzeIf(fir::IfOp op, AffineFunctionAnalysis &afa) {302 if (op.getNumResults() == 0)303 return true;304 LLVM_DEBUG(llvm::dbgs()305 << "AffineIfAnalysis: not promoting as op has results\n";);306 return false;307 }308 309 bool legality{};310};311} // namespace312 313AffineIfAnalysis314AffineFunctionAnalysis::getChildIfAnalysis(fir::IfOp op) const {315 auto it = ifAnalysisMap.find_as(op);316 if (it == ifAnalysisMap.end()) {317 LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n";318 op.dump(););319 op.emitError("error in fetching if analysis in AffineFunctionAnalysis\n");320 return {};321 }322 return it->getSecond();323}324 325/// AffineMap rewriting fir.array_coor operation to affine apply,326/// %dim = fir.gendim %lowerBound, %upperBound, %stride327/// %a = fir.array_coor %arr(%dim) %i328/// returning affineMap = affine_map<(i)[lb, ub, st] -> (i*st - lb)>329static mlir::AffineMap createArrayIndexAffineMap(unsigned dimensions,330 MLIRContext *context) {331 auto index = mlir::getAffineConstantExpr(0, context);332 auto accuExtent = mlir::getAffineConstantExpr(1, context);333 for (unsigned i = 0; i < dimensions; ++i) {334 mlir::AffineExpr idx = mlir::getAffineDimExpr(i, context),335 lowerBound = mlir::getAffineSymbolExpr(i * 3, context),336 currentExtent =337 mlir::getAffineSymbolExpr(i * 3 + 1, context),338 stride = mlir::getAffineSymbolExpr(i * 3 + 2, context),339 currentPart = (idx * stride - lowerBound) * accuExtent;340 index = currentPart + index;341 accuExtent = accuExtent * currentExtent;342 }343 return mlir::AffineMap::get(dimensions, dimensions * 3, index);344}345 346static std::optional<int64_t> constantIntegerLike(const mlir::Value value) {347 if (auto definition = value.getDefiningOp<mlir::arith::ConstantOp>())348 if (auto stepAttr = mlir::dyn_cast<IntegerAttr>(definition.getValue()))349 return stepAttr.getInt();350 return {};351}352 353static mlir::Type coordinateArrayElement(fir::ArrayCoorOp op) {354 if (auto refType =355 mlir::dyn_cast_or_null<ReferenceType>(op.getMemref().getType())) {356 if (auto seqType =357 mlir::dyn_cast_or_null<SequenceType>(refType.getEleTy())) {358 return seqType.getEleTy();359 }360 }361 op.emitError(362 "AffineLoopConversion: array type in coordinate operation not valid\n");363 return mlir::Type();364}365 366static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeOp shape,367 SmallVectorImpl<mlir::Value> &indexArgs,368 mlir::PatternRewriter &rewriter) {369 auto one = mlir::arith::ConstantOp::create(rewriter, acoOp.getLoc(),370 rewriter.getIndexType(),371 rewriter.getIndexAttr(1));372 auto extents = shape.getExtents();373 for (auto i = extents.begin(); i < extents.end(); i++) {374 indexArgs.push_back(one);375 indexArgs.push_back(*i);376 indexArgs.push_back(one);377 }378}379 380static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeShiftOp shape,381 SmallVectorImpl<mlir::Value> &indexArgs,382 mlir::PatternRewriter &rewriter) {383 auto one = mlir::arith::ConstantOp::create(rewriter, acoOp.getLoc(),384 rewriter.getIndexType(),385 rewriter.getIndexAttr(1));386 auto extents = shape.getPairs();387 for (auto i = extents.begin(); i < extents.end();) {388 indexArgs.push_back(*i++);389 indexArgs.push_back(*i++);390 indexArgs.push_back(one);391 }392}393 394static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::SliceOp slice,395 SmallVectorImpl<mlir::Value> &indexArgs,396 mlir::PatternRewriter &rewriter) {397 auto extents = slice.getTriples();398 for (auto i = extents.begin(); i < extents.end();) {399 indexArgs.push_back(*i++);400 indexArgs.push_back(*i++);401 indexArgs.push_back(*i++);402 }403}404 405static void populateIndexArgs(fir::ArrayCoorOp acoOp,406 SmallVectorImpl<mlir::Value> &indexArgs,407 mlir::PatternRewriter &rewriter) {408 if (auto shape = acoOp.getShape().getDefiningOp<ShapeOp>())409 return populateIndexArgs(acoOp, shape, indexArgs, rewriter);410 if (auto shapeShift = acoOp.getShape().getDefiningOp<ShapeShiftOp>())411 return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter);412 if (auto slice = acoOp.getShape().getDefiningOp<SliceOp>())413 return populateIndexArgs(acoOp, slice, indexArgs, rewriter);414}415 416/// Returns affine.apply and fir.convert from array_coor and gendims417static std::pair<affine::AffineApplyOp, fir::ConvertOp>418createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) {419 auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>();420 auto affineMap =421 createArrayIndexAffineMap(acoOp.getIndices().size(), acoOp.getContext());422 SmallVector<mlir::Value> indexArgs;423 indexArgs.append(acoOp.getIndices().begin(), acoOp.getIndices().end());424 425 populateIndexArgs(acoOp, indexArgs, rewriter);426 427 auto affineApply = affine::AffineApplyOp::create(rewriter, acoOp.getLoc(),428 affineMap, indexArgs);429 auto arrayElementType = coordinateArrayElement(acoOp);430 auto newType =431 mlir::MemRefType::get({mlir::ShapedType::kDynamic}, arrayElementType);432 auto arrayConvert = fir::ConvertOp::create(rewriter, acoOp.getLoc(), newType,433 acoOp.getMemref());434 return std::make_pair(affineApply, arrayConvert);435}436 437static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) {438 rewriter.setInsertionPoint(loadOp);439 auto affineOps = createAffineOps(loadOp.getMemref(), rewriter);440 rewriter.replaceOpWithNewOp<affine::AffineLoadOp>(441 loadOp, affineOps.second.getResult(), affineOps.first.getResult());442}443 444static void rewriteStore(fir::StoreOp storeOp,445 mlir::PatternRewriter &rewriter) {446 rewriter.setInsertionPoint(storeOp);447 auto affineOps = createAffineOps(storeOp.getMemref(), rewriter);448 rewriter.replaceOpWithNewOp<affine::AffineStoreOp>(449 storeOp, storeOp.getValue(), affineOps.second.getResult(),450 affineOps.first.getResult());451}452 453static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) {454 for (auto &bodyOp : llvm::make_early_inc_range(block->getOperations())) {455 if (isa<fir::LoadOp>(bodyOp))456 rewriteLoad(cast<fir::LoadOp>(bodyOp), rewriter);457 else if (isa<fir::StoreOp>(bodyOp))458 rewriteStore(cast<fir::StoreOp>(bodyOp), rewriter);459 }460}461 462namespace {463/// Convert `fir.do_loop` to `affine.for`, creates fir.convert for arrays to464/// memref, rewrites array_coor to affine.apply with affine_map. Rewrites fir465/// loads and stores to affine.466class AffineLoopConversion : public mlir::OpRewritePattern<fir::DoLoopOp> {467public:468 using OpRewritePattern::OpRewritePattern;469 AffineLoopConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)470 : OpRewritePattern(context), functionAnalysis(afa) {}471 472 llvm::LogicalResult473 matchAndRewrite(fir::DoLoopOp loop,474 mlir::PatternRewriter &rewriter) const override {475 LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: rewriting loop:\n";476 loop.dump(););477 [[maybe_unused]] auto loopAnalysis =478 functionAnalysis.getChildLoopAnalysis(loop);479 if (!loopAnalysis.canPromoteToAffine())480 return rewriter.notifyMatchFailure(loop, "cannot promote to affine");481 auto &loopOps = loop.getBody()->getOperations();482 auto resultOp = cast<fir::ResultOp>(loop.getBody()->getTerminator());483 auto results = resultOp.getOperands();484 auto loopResults = loop->getResults();485 auto loopAndIndex = createAffineFor(loop, rewriter);486 auto affineFor = loopAndIndex.first;487 auto inductionVar = loopAndIndex.second;488 489 if (loop.getFinalValue()) {490 results = results.drop_front();491 loopResults = loopResults.drop_front();492 }493 494 rewriter.startOpModification(affineFor.getOperation());495 affineFor.getBody()->getOperations().splice(496 std::prev(affineFor.getBody()->end()), loopOps, loopOps.begin(),497 std::prev(loopOps.end()));498 rewriter.replaceAllUsesWith(loop.getRegionIterArgs(),499 affineFor.getRegionIterArgs());500 if (!results.empty()) {501 rewriter.setInsertionPointToEnd(affineFor.getBody());502 affine::AffineYieldOp::create(rewriter, resultOp->getLoc(), results);503 }504 rewriter.finalizeOpModification(affineFor.getOperation());505 506 rewriter.startOpModification(loop.getOperation());507 loop.getInductionVar().replaceAllUsesWith(inductionVar);508 rewriter.finalizeOpModification(loop.getOperation());509 510 rewriteMemoryOps(affineFor.getBody(), rewriter);511 512 LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: loop rewriten to:\n";513 affineFor.dump(););514 rewriter.replaceAllUsesWith(loopResults, affineFor->getResults());515 rewriter.eraseOp(loop);516 return success();517 }518 519private:520 std::pair<affine::AffineForOp, mlir::Value>521 createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {522 if (auto constantStep = constantIntegerLike(op.getStep()))523 if (*constantStep > 0)524 return positiveConstantStep(op, *constantStep, rewriter);525 return genericBounds(op, rewriter);526 }527 528 // when step for the loop is positive compile time constant529 std::pair<affine::AffineForOp, mlir::Value>530 positiveConstantStep(fir::DoLoopOp op, int64_t step,531 mlir::PatternRewriter &rewriter) const {532 auto affineFor = affine::AffineForOp::create(533 rewriter, op.getLoc(), ValueRange(op.getLowerBound()),534 mlir::AffineMap::get(0, 1,535 mlir::getAffineSymbolExpr(0, op.getContext())),536 ValueRange(op.getUpperBound()),537 mlir::AffineMap::get(0, 1,538 1 + mlir::getAffineSymbolExpr(0, op.getContext())),539 step, op.getIterOperands());540 return std::make_pair(affineFor, affineFor.getInductionVar());541 }542 543 std::pair<affine::AffineForOp, mlir::Value>544 genericBounds(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {545 auto lowerBound = mlir::getAffineSymbolExpr(0, op.getContext());546 auto upperBound = mlir::getAffineSymbolExpr(1, op.getContext());547 auto step = mlir::getAffineSymbolExpr(2, op.getContext());548 mlir::AffineMap upperBoundMap = mlir::AffineMap::get(549 0, 3, (upperBound - lowerBound + step).floorDiv(step));550 auto genericUpperBound = affine::AffineApplyOp::create(551 rewriter, op.getLoc(), upperBoundMap,552 ValueRange({op.getLowerBound(), op.getUpperBound(), op.getStep()}));553 auto actualIndexMap = mlir::AffineMap::get(554 1, 2,555 (lowerBound + mlir::getAffineDimExpr(0, op.getContext())) *556 mlir::getAffineSymbolExpr(1, op.getContext()));557 558 auto affineFor = affine::AffineForOp::create(559 rewriter, op.getLoc(), ValueRange(),560 AffineMap::getConstantMap(0, op.getContext()),561 genericUpperBound.getResult(),562 mlir::AffineMap::get(0, 1,563 1 + mlir::getAffineSymbolExpr(0, op.getContext())),564 1, op.getIterOperands());565 rewriter.setInsertionPointToStart(affineFor.getBody());566 auto actualIndex = affine::AffineApplyOp::create(567 rewriter, op.getLoc(), actualIndexMap,568 ValueRange(569 {affineFor.getInductionVar(), op.getLowerBound(), op.getStep()}));570 return std::make_pair(affineFor, actualIndex.getResult());571 }572 573 AffineFunctionAnalysis &functionAnalysis;574};575 576/// Convert `fir.if` to `affine.if`.577class AffineIfConversion : public mlir::OpRewritePattern<fir::IfOp> {578public:579 using OpRewritePattern::OpRewritePattern;580 AffineIfConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)581 : OpRewritePattern(context), functionAnalysis(afa) {}582 llvm::LogicalResult583 matchAndRewrite(fir::IfOp op,584 mlir::PatternRewriter &rewriter) const override {585 LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: rewriting if:\n";586 op.dump(););587 if (!functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine())588 return rewriter.notifyMatchFailure(op, "cannot promote to affine");589 auto &ifOps = op.getThenRegion().front().getOperations();590 auto affineCondition = AffineIfCondition(op.getCondition());591 if (!affineCondition.hasIntegerSet()) {592 LLVM_DEBUG(593 llvm::dbgs()594 << "AffineIfConversion: couldn't calculate affine condition\n";);595 return failure();596 }597 auto affineIf = affine::AffineIfOp::create(598 rewriter, op.getLoc(), affineCondition.getIntegerSet(),599 affineCondition.getAffineArgs(), !op.getElseRegion().empty());600 rewriter.startOpModification(affineIf);601 affineIf.getThenBlock()->getOperations().splice(602 std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(),603 std::prev(ifOps.end()));604 if (!op.getElseRegion().empty()) {605 auto &otherOps = op.getElseRegion().front().getOperations();606 affineIf.getElseBlock()->getOperations().splice(607 std::prev(affineIf.getElseBlock()->end()), otherOps, otherOps.begin(),608 std::prev(otherOps.end()));609 }610 rewriter.finalizeOpModification(affineIf);611 rewriteMemoryOps(affineIf.getBody(), rewriter);612 613 LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: if converted to:\n";614 affineIf.dump(););615 rewriter.replaceOp(op, affineIf.getOperation()->getResults());616 return success();617 }618 619 AffineFunctionAnalysis &functionAnalysis;620};621 622/// Promote fir.do_loop and fir.if to affine.for and affine.if, in the cases623/// where such a promotion is possible.624class AffineDialectPromotion625 : public fir::impl::AffineDialectPromotionBase<AffineDialectPromotion> {626public:627 void runOnOperation() override {628 629 auto *context = &getContext();630 auto function = getOperation();631 markAllAnalysesPreserved();632 auto functionAnalysis = AffineFunctionAnalysis(function);633 mlir::RewritePatternSet patterns(context);634 patterns.insert<AffineIfConversion>(context, functionAnalysis);635 patterns.insert<AffineLoopConversion>(context, functionAnalysis);636 LLVM_DEBUG(llvm::dbgs()637 << "AffineDialectPromotion: running promotion on: \n";638 function.print(llvm::dbgs()););639 // apply the patterns640 walkAndApplyPatterns(function, std::move(patterns));641 }642};643} // namespace644 645/// Convert FIR loop constructs to the Affine dialect646std::unique_ptr<mlir::Pass> fir::createPromoteToAffinePass() {647 return std::make_unique<AffineDialectPromotion>();648}649