235 lines · cpp
1//===-- FIRToSCF.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#include "flang/Optimizer/Dialect/FIRDialect.h"10#include "flang/Optimizer/Transforms/Passes.h"11#include "mlir/Dialect/SCF/IR/SCF.h"12#include "mlir/Transforms/WalkPatternRewriteDriver.h"13 14namespace fir {15#define GEN_PASS_DEF_FIRTOSCFPASS16#include "flang/Optimizer/Transforms/Passes.h.inc"17} // namespace fir18 19namespace {20class FIRToSCFPass : public fir::impl::FIRToSCFPassBase<FIRToSCFPass> {21 using FIRToSCFPassBase::FIRToSCFPassBase;22 23public:24 void runOnOperation() override;25};26 27struct DoLoopConversion : public mlir::OpRewritePattern<fir::DoLoopOp> {28 using OpRewritePattern<fir::DoLoopOp>::OpRewritePattern;29 30 DoLoopConversion(mlir::MLIRContext *context,31 bool parallelUnorderedLoop = false,32 mlir::PatternBenefit benefit = 1)33 : OpRewritePattern<fir::DoLoopOp>(context, benefit),34 parallelUnorderedLoop(parallelUnorderedLoop) {}35 36 mlir::LogicalResult37 matchAndRewrite(fir::DoLoopOp doLoopOp,38 mlir::PatternRewriter &rewriter) const override {39 mlir::Location loc = doLoopOp.getLoc();40 bool hasFinalValue = doLoopOp.getFinalValue().has_value();41 bool isUnordered = doLoopOp.getUnordered().has_value();42 43 // Get loop values from the DoLoopOp44 mlir::Value low = doLoopOp.getLowerBound();45 mlir::Value high = doLoopOp.getUpperBound();46 assert(low && high && "must be a Value");47 mlir::Value step = doLoopOp.getStep();48 mlir::SmallVector<mlir::Value> iterArgs;49 if (hasFinalValue)50 iterArgs.push_back(low);51 iterArgs.append(doLoopOp.getIterOperands().begin(),52 doLoopOp.getIterOperands().end());53 54 // fir.do_loop iterates over the interval [%l, %u], and the step may be55 // negative. But scf.for iterates over the interval [%l, %u), and the step56 // must be a positive value.57 // For easier conversion, we calculate the trip count and use a canonical58 // induction variable.59 auto diff = mlir::arith::SubIOp::create(rewriter, loc, high, low);60 auto distance = mlir::arith::AddIOp::create(rewriter, loc, diff, step);61 auto tripCount =62 mlir::arith::DivSIOp::create(rewriter, loc, distance, step);63 auto zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);64 auto one = mlir::arith::ConstantIndexOp::create(rewriter, loc, 1);65 66 // Create the scf.for or scf.parallel operation67 mlir::Operation *scfLoopOp = nullptr;68 if (isUnordered && parallelUnorderedLoop) {69 scfLoopOp = mlir::scf::ParallelOp::create(rewriter, loc, {zero},70 {tripCount}, {one}, iterArgs);71 } else {72 scfLoopOp = mlir::scf::ForOp::create(rewriter, loc, zero, tripCount, one,73 iterArgs);74 }75 76 // Move the body of the fir.do_loop to the scf.for or scf.parallel77 auto &loopOps = doLoopOp.getBody()->getOperations();78 auto resultOp =79 mlir::cast<fir::ResultOp>(doLoopOp.getBody()->getTerminator());80 auto results = resultOp.getOperands();81 auto scfLoopLikeOp = mlir::cast<mlir::LoopLikeOpInterface>(scfLoopOp);82 mlir::Block &scfLoopBody = scfLoopLikeOp.getLoopRegions().front()->front();83 84 scfLoopBody.getOperations().splice(scfLoopBody.begin(), loopOps,85 loopOps.begin(),86 std::prev(loopOps.end()));87 88 rewriter.setInsertionPointToStart(&scfLoopBody);89 mlir::Value iv = mlir::arith::MulIOp::create(90 rewriter, loc, scfLoopLikeOp.getSingleInductionVar().value(), step);91 iv = mlir::arith::AddIOp::create(rewriter, loc, low, iv);92 93 if (!results.empty()) {94 rewriter.setInsertionPointToEnd(&scfLoopBody);95 mlir::scf::YieldOp::create(rewriter, resultOp->getLoc(), results);96 }97 doLoopOp.getInductionVar().replaceAllUsesWith(iv);98 rewriter.replaceAllUsesWith(99 doLoopOp.getRegionIterArgs(),100 hasFinalValue ? scfLoopLikeOp.getRegionIterArgs().drop_front()101 : scfLoopLikeOp.getRegionIterArgs());102 103 // Copy loop annotations from the fir.do_loop to scf loop op.104 if (auto ann = doLoopOp.getLoopAnnotation())105 scfLoopOp->setAttr("loop_annotation", *ann);106 107 rewriter.replaceOp(doLoopOp, scfLoopOp);108 return mlir::success();109 }110 111private:112 bool parallelUnorderedLoop;113};114 115struct IterWhileConversion : public mlir::OpRewritePattern<fir::IterWhileOp> {116 using OpRewritePattern<fir::IterWhileOp>::OpRewritePattern;117 118 mlir::LogicalResult119 matchAndRewrite(fir::IterWhileOp iterWhileOp,120 mlir::PatternRewriter &rewriter) const override {121 122 mlir::Location loc = iterWhileOp.getLoc();123 mlir::Value lowerBound = iterWhileOp.getLowerBound();124 mlir::Value upperBound = iterWhileOp.getUpperBound();125 mlir::Value step = iterWhileOp.getStep();126 127 mlir::Value okInit = iterWhileOp.getIterateIn();128 mlir::ValueRange iterArgs = iterWhileOp.getInitArgs();129 130 mlir::SmallVector<mlir::Value> initVals;131 initVals.push_back(lowerBound);132 initVals.push_back(okInit);133 initVals.append(iterArgs.begin(), iterArgs.end());134 135 mlir::SmallVector<mlir::Type> loopTypes;136 loopTypes.push_back(lowerBound.getType());137 loopTypes.push_back(okInit.getType());138 for (auto val : iterArgs)139 loopTypes.push_back(val.getType());140 141 auto scfWhileOp =142 mlir::scf::WhileOp::create(rewriter, loc, loopTypes, initVals);143 144 auto &beforeBlock = *rewriter.createBlock(145 &scfWhileOp.getBefore(), scfWhileOp.getBefore().end(), loopTypes,146 mlir::SmallVector<mlir::Location>(loopTypes.size(), loc));147 148 mlir::Region::BlockArgListType argsInBefore =149 scfWhileOp.getBefore().getArguments();150 auto ivInBefore = argsInBefore[0];151 auto earlyExitInBefore = argsInBefore[1];152 153 rewriter.setInsertionPointToStart(&beforeBlock);154 155 mlir::Value inductionCmp = mlir::arith::CmpIOp::create(156 rewriter, loc, mlir::arith::CmpIPredicate::sle, ivInBefore, upperBound);157 mlir::Value cond = mlir::arith::AndIOp::create(rewriter, loc, inductionCmp,158 earlyExitInBefore);159 160 mlir::scf::ConditionOp::create(rewriter, loc, cond, argsInBefore);161 162 rewriter.moveBlockBefore(iterWhileOp.getBody(), &scfWhileOp.getAfter(),163 scfWhileOp.getAfter().begin());164 165 auto *afterBody = scfWhileOp.getAfterBody();166 auto resultOp = mlir::cast<fir::ResultOp>(afterBody->getTerminator());167 mlir::SmallVector<mlir::Value> results(resultOp->getOperands());168 mlir::Value ivInAfter = scfWhileOp.getAfterArguments()[0];169 170 rewriter.setInsertionPointToStart(afterBody);171 results[0] = mlir::arith::AddIOp::create(rewriter, loc, ivInAfter, step);172 173 rewriter.setInsertionPointToEnd(afterBody);174 rewriter.replaceOpWithNewOp<mlir::scf::YieldOp>(resultOp, results);175 176 scfWhileOp->setAttrs(iterWhileOp->getAttrs());177 rewriter.replaceOp(iterWhileOp, scfWhileOp);178 return mlir::success();179 }180};181 182void copyBlockAndTransformResult(mlir::PatternRewriter &rewriter,183 mlir::Block &srcBlock, mlir::Block &dstBlock) {184 mlir::Operation *srcTerminator = srcBlock.getTerminator();185 auto resultOp = mlir::cast<fir::ResultOp>(srcTerminator);186 187 dstBlock.getOperations().splice(dstBlock.begin(), srcBlock.getOperations(),188 srcBlock.begin(), std::prev(srcBlock.end()));189 190 if (!resultOp->getOperands().empty()) {191 rewriter.setInsertionPointToEnd(&dstBlock);192 mlir::scf::YieldOp::create(rewriter, resultOp->getLoc(),193 resultOp->getOperands());194 }195 196 rewriter.eraseOp(srcTerminator);197}198 199struct IfConversion : public mlir::OpRewritePattern<fir::IfOp> {200 using OpRewritePattern<fir::IfOp>::OpRewritePattern;201 mlir::LogicalResult202 matchAndRewrite(fir::IfOp ifOp,203 mlir::PatternRewriter &rewriter) const override {204 bool hasElse = !ifOp.getElseRegion().empty();205 auto scfIfOp =206 mlir::scf::IfOp::create(rewriter, ifOp.getLoc(), ifOp.getResultTypes(),207 ifOp.getCondition(), hasElse);208 209 copyBlockAndTransformResult(rewriter, ifOp.getThenRegion().front(),210 scfIfOp.getThenRegion().front());211 212 if (hasElse) {213 copyBlockAndTransformResult(rewriter, ifOp.getElseRegion().front(),214 scfIfOp.getElseRegion().front());215 }216 217 scfIfOp->setAttrs(ifOp->getAttrs());218 rewriter.replaceOp(ifOp, scfIfOp);219 return mlir::success();220 }221};222} // namespace223 224void fir::populateFIRToSCFRewrites(mlir::RewritePatternSet &patterns,225 bool parallelUnordered) {226 patterns.add<IterWhileConversion, IfConversion>(patterns.getContext());227 patterns.add<DoLoopConversion>(patterns.getContext(), parallelUnordered);228}229 230void FIRToSCFPass::runOnOperation() {231 mlir::RewritePatternSet patterns(&getContext());232 fir::populateFIRToSCFRewrites(patterns, parallelUnordered);233 walkAndApplyPatterns(getOperation(), std::move(patterns));234}235