371 lines · cpp
1//===-- ControlFlowConverter.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/Dialect/FIROps.h"11#include "flang/Optimizer/Dialect/FIROpsSupport.h"12#include "flang/Optimizer/Dialect/Support/FIRContext.h"13#include "flang/Optimizer/Dialect/Support/KindMapping.h"14#include "flang/Optimizer/Support/InternalNames.h"15#include "flang/Optimizer/Support/TypeCode.h"16#include "flang/Optimizer/Transforms/Passes.h"17#include "flang/Runtime/derived-api.h"18#include "mlir/Dialect/Affine/IR/AffineOps.h"19#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"20#include "mlir/Dialect/Func/IR/FuncOps.h"21#include "mlir/Pass/Pass.h"22#include "mlir/Transforms/DialectConversion.h"23#include "llvm/ADT/SmallSet.h"24#include "llvm/Support/CommandLine.h"25 26namespace fir {27#define GEN_PASS_DEF_CFGCONVERSION28#include "flang/Optimizer/Transforms/Passes.h.inc"29} // namespace fir30 31using namespace fir;32using namespace mlir;33 34namespace {35 36// Conversion of fir control ops to more primitive control-flow.37//38// FIR loops that cannot be converted to the affine dialect will remain as39// `fir.do_loop` operations. These can be converted to control-flow operations.40 41/// Convert `fir.do_loop` to CFG42class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {43public:44 using OpRewritePattern::OpRewritePattern;45 46 CfgLoopConv(mlir::MLIRContext *ctx, bool forceLoopToExecuteOnce, bool setNSW)47 : mlir::OpRewritePattern<fir::DoLoopOp>(ctx),48 forceLoopToExecuteOnce(forceLoopToExecuteOnce), setNSW(setNSW) {}49 50 llvm::LogicalResult51 matchAndRewrite(DoLoopOp loop,52 mlir::PatternRewriter &rewriter) const override {53 auto loc = loop.getLoc();54 mlir::arith::IntegerOverflowFlags flags{};55 if (setNSW)56 flags = bitEnumSet(flags, mlir::arith::IntegerOverflowFlags::nsw);57 auto iofAttr = mlir::arith::IntegerOverflowFlagsAttr::get(58 rewriter.getContext(), flags);59 60 // Create the start and end blocks that will wrap the DoLoopOp with an61 // initalizer and an end point62 auto *initBlock = rewriter.getInsertionBlock();63 auto initPos = rewriter.getInsertionPoint();64 auto *endBlock = rewriter.splitBlock(initBlock, initPos);65 66 // Split the first DoLoopOp block in two parts. The part before will be the67 // conditional block since it already has the induction variable and68 // loop-carried values as arguments.69 auto *conditionalBlock = &loop.getRegion().front();70 conditionalBlock->addArgument(rewriter.getIndexType(), loc);71 auto *firstBlock =72 rewriter.splitBlock(conditionalBlock, conditionalBlock->begin());73 auto *lastBlock = &loop.getRegion().back();74 75 // Move the blocks from the DoLoopOp between initBlock and endBlock76 rewriter.inlineRegionBefore(loop.getRegion(), endBlock);77 78 // Get loop values from the DoLoopOp79 auto low = loop.getLowerBound();80 auto high = loop.getUpperBound();81 assert(low && high && "must be a Value");82 auto step = loop.getStep();83 84 // Initalization block85 rewriter.setInsertionPointToEnd(initBlock);86 auto diff = mlir::arith::SubIOp::create(rewriter, loc, high, low);87 auto distance = mlir::arith::AddIOp::create(rewriter, loc, diff, step);88 mlir::Value iters =89 mlir::arith::DivSIOp::create(rewriter, loc, distance, step);90 91 if (forceLoopToExecuteOnce) {92 auto zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);93 auto cond = mlir::arith::CmpIOp::create(94 rewriter, loc, arith::CmpIPredicate::sle, iters, zero);95 auto one = mlir::arith::ConstantIndexOp::create(rewriter, loc, 1);96 iters = mlir::arith::SelectOp::create(rewriter, loc, cond, one, iters);97 }98 99 llvm::SmallVector<mlir::Value> loopOperands;100 loopOperands.push_back(low);101 auto operands = loop.getIterOperands();102 loopOperands.append(operands.begin(), operands.end());103 loopOperands.push_back(iters);104 105 mlir::cf::BranchOp::create(rewriter, loc, conditionalBlock, loopOperands);106 107 // Last loop block108 auto *terminator = lastBlock->getTerminator();109 rewriter.setInsertionPointToEnd(lastBlock);110 auto iv = conditionalBlock->getArgument(0);111 mlir::Value steppedIndex =112 mlir::arith::AddIOp::create(rewriter, loc, iv, step, iofAttr);113 assert(steppedIndex && "must be a Value");114 auto lastArg = conditionalBlock->getNumArguments() - 1;115 auto itersLeft = conditionalBlock->getArgument(lastArg);116 auto one = mlir::arith::ConstantIndexOp::create(rewriter, loc, 1);117 mlir::Value itersMinusOne =118 mlir::arith::SubIOp::create(rewriter, loc, itersLeft, one);119 120 llvm::SmallVector<mlir::Value> loopCarried;121 loopCarried.push_back(steppedIndex);122 auto begin = loop.getFinalValue() ? std::next(terminator->operand_begin())123 : terminator->operand_begin();124 loopCarried.append(begin, terminator->operand_end());125 loopCarried.push_back(itersMinusOne);126 auto backEdge = mlir::cf::BranchOp::create(rewriter, loc, conditionalBlock,127 loopCarried);128 rewriter.eraseOp(terminator);129 130 // Copy loop annotations from the do loop to the loop back edge.131 if (auto ann = loop.getLoopAnnotation())132 backEdge->setAttr("loop_annotation", *ann);133 134 // Conditional block135 rewriter.setInsertionPointToEnd(conditionalBlock);136 auto zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);137 auto comparison = mlir::arith::CmpIOp::create(138 rewriter, loc, arith::CmpIPredicate::sgt, itersLeft, zero);139 140 mlir::cf::CondBranchOp::create(rewriter, loc, comparison, firstBlock,141 llvm::ArrayRef<mlir::Value>(), endBlock,142 llvm::ArrayRef<mlir::Value>());143 144 // The result of the loop operation is the values of the condition block145 // arguments except the induction variable on the last iteration.146 auto args = loop.getFinalValue()147 ? conditionalBlock->getArguments()148 : conditionalBlock->getArguments().drop_front();149 rewriter.replaceOp(loop, args.drop_back());150 return success();151 }152 153private:154 bool forceLoopToExecuteOnce;155 bool setNSW;156};157 158/// Convert `fir.if` to control-flow159class CfgIfConv : public mlir::OpRewritePattern<fir::IfOp> {160public:161 using OpRewritePattern::OpRewritePattern;162 163 CfgIfConv(mlir::MLIRContext *ctx, bool forceLoopToExecuteOnce, bool setNSW)164 : mlir::OpRewritePattern<fir::IfOp>(ctx) {}165 166 llvm::LogicalResult167 matchAndRewrite(IfOp ifOp, mlir::PatternRewriter &rewriter) const override {168 auto loc = ifOp.getLoc();169 170 // Split the block containing the 'fir.if' into two parts. The part before171 // will contain the condition, the part after will be the continuation172 // point.173 auto *condBlock = rewriter.getInsertionBlock();174 auto opPosition = rewriter.getInsertionPoint();175 auto *remainingOpsBlock = rewriter.splitBlock(condBlock, opPosition);176 mlir::Block *continueBlock;177 if (ifOp.getNumResults() == 0) {178 continueBlock = remainingOpsBlock;179 } else {180 continueBlock = rewriter.createBlock(181 remainingOpsBlock, ifOp.getResultTypes(),182 llvm::SmallVector<mlir::Location>(ifOp.getNumResults(), loc));183 mlir::cf::BranchOp::create(rewriter, loc, remainingOpsBlock);184 }185 186 // Move blocks from the "then" region to the region containing 'fir.if',187 // place it before the continuation block, and branch to it.188 auto &ifOpRegion = ifOp.getThenRegion();189 auto *ifOpBlock = &ifOpRegion.front();190 auto *ifOpTerminator = ifOpRegion.back().getTerminator();191 auto ifOpTerminatorOperands = ifOpTerminator->getOperands();192 rewriter.setInsertionPointToEnd(&ifOpRegion.back());193 mlir::cf::BranchOp::create(rewriter, loc, continueBlock,194 ifOpTerminatorOperands);195 rewriter.eraseOp(ifOpTerminator);196 rewriter.inlineRegionBefore(ifOpRegion, continueBlock);197 198 // Move blocks from the "else" region (if present) to the region containing199 // 'fir.if', place it before the continuation block and branch to it. It200 // will be placed after the "then" regions.201 auto *otherwiseBlock = continueBlock;202 auto &otherwiseRegion = ifOp.getElseRegion();203 if (!otherwiseRegion.empty()) {204 otherwiseBlock = &otherwiseRegion.front();205 auto *otherwiseTerm = otherwiseRegion.back().getTerminator();206 auto otherwiseTermOperands = otherwiseTerm->getOperands();207 rewriter.setInsertionPointToEnd(&otherwiseRegion.back());208 mlir::cf::BranchOp::create(rewriter, loc, continueBlock,209 otherwiseTermOperands);210 rewriter.eraseOp(otherwiseTerm);211 rewriter.inlineRegionBefore(otherwiseRegion, continueBlock);212 }213 214 rewriter.setInsertionPointToEnd(condBlock);215 auto branchOp = mlir::cf::CondBranchOp::create(216 rewriter, loc, ifOp.getCondition(), ifOpBlock,217 llvm::ArrayRef<mlir::Value>(), otherwiseBlock,218 llvm::ArrayRef<mlir::Value>());219 llvm::ArrayRef<int32_t> weights = ifOp.getWeights();220 if (!weights.empty())221 branchOp.setWeights(weights);222 rewriter.replaceOp(ifOp, continueBlock->getArguments());223 return success();224 }225};226 227/// Convert `fir.iter_while` to control-flow.228class CfgIterWhileConv : public mlir::OpRewritePattern<fir::IterWhileOp> {229public:230 using OpRewritePattern::OpRewritePattern;231 232 CfgIterWhileConv(mlir::MLIRContext *ctx, bool forceLoopToExecuteOnce,233 bool setNSW)234 : mlir::OpRewritePattern<fir::IterWhileOp>(ctx), setNSW(setNSW) {}235 236 llvm::LogicalResult237 matchAndRewrite(fir::IterWhileOp whileOp,238 mlir::PatternRewriter &rewriter) const override {239 auto loc = whileOp.getLoc();240 mlir::arith::IntegerOverflowFlags flags{};241 if (setNSW)242 flags = bitEnumSet(flags, mlir::arith::IntegerOverflowFlags::nsw);243 auto iofAttr = mlir::arith::IntegerOverflowFlagsAttr::get(244 rewriter.getContext(), flags);245 246 // Start by splitting the block containing the 'fir.do_loop' into two parts.247 // The part before will get the init code, the part after will be the end248 // point.249 auto *initBlock = rewriter.getInsertionBlock();250 auto initPosition = rewriter.getInsertionPoint();251 auto *endBlock = rewriter.splitBlock(initBlock, initPosition);252 253 // Use the first block of the loop body as the condition block since it is254 // the block that has the induction variable and loop-carried values as255 // arguments. Split out all operations from the first block into a new256 // block. Move all body blocks from the loop body region to the region257 // containing the loop.258 auto *conditionBlock = &whileOp.getRegion().front();259 auto *firstBodyBlock =260 rewriter.splitBlock(conditionBlock, conditionBlock->begin());261 auto *lastBodyBlock = &whileOp.getRegion().back();262 rewriter.inlineRegionBefore(whileOp.getRegion(), endBlock);263 auto iv = conditionBlock->getArgument(0);264 auto iterateVar = conditionBlock->getArgument(1);265 266 // Append the induction variable stepping logic to the last body block and267 // branch back to the condition block. Loop-carried values are taken from268 // operands of the loop terminator.269 auto *terminator = lastBodyBlock->getTerminator();270 rewriter.setInsertionPointToEnd(lastBodyBlock);271 auto step = whileOp.getStep();272 mlir::Value stepped =273 mlir::arith::AddIOp::create(rewriter, loc, iv, step, iofAttr);274 assert(stepped && "must be a Value");275 276 llvm::SmallVector<mlir::Value> loopCarried;277 loopCarried.push_back(stepped);278 auto begin = whileOp.getFinalValue()279 ? std::next(terminator->operand_begin())280 : terminator->operand_begin();281 loopCarried.append(begin, terminator->operand_end());282 mlir::cf::BranchOp::create(rewriter, loc, conditionBlock, loopCarried);283 rewriter.eraseOp(terminator);284 285 // Compute loop bounds before branching to the condition.286 rewriter.setInsertionPointToEnd(initBlock);287 auto lowerBound = whileOp.getLowerBound();288 auto upperBound = whileOp.getUpperBound();289 assert(lowerBound && upperBound && "must be a Value");290 291 // The initial values of loop-carried values is obtained from the operands292 // of the loop operation.293 llvm::SmallVector<mlir::Value> destOperands;294 destOperands.push_back(lowerBound);295 auto iterOperands = whileOp.getIterOperands();296 destOperands.append(iterOperands.begin(), iterOperands.end());297 mlir::cf::BranchOp::create(rewriter, loc, conditionBlock, destOperands);298 299 // With the body block done, we can fill in the condition block.300 rewriter.setInsertionPointToEnd(conditionBlock);301 // The comparison depends on the sign of the step value. We fully expect302 // this expression to be folded by the optimizer or LLVM. This expression303 // is written this way so that `step == 0` always returns `false`.304 auto zero = mlir::arith::ConstantIndexOp::create(rewriter, loc, 0);305 auto compl0 = mlir::arith::CmpIOp::create(306 rewriter, loc, arith::CmpIPredicate::slt, zero, step);307 auto compl1 = mlir::arith::CmpIOp::create(308 rewriter, loc, arith::CmpIPredicate::sle, iv, upperBound);309 auto compl2 = mlir::arith::CmpIOp::create(310 rewriter, loc, arith::CmpIPredicate::slt, step, zero);311 auto compl3 = mlir::arith::CmpIOp::create(312 rewriter, loc, arith::CmpIPredicate::sle, upperBound, iv);313 auto cmp0 = mlir::arith::AndIOp::create(rewriter, loc, compl0, compl1);314 auto cmp1 = mlir::arith::AndIOp::create(rewriter, loc, compl2, compl3);315 auto cmp2 = mlir::arith::OrIOp::create(rewriter, loc, cmp0, cmp1);316 // Remember to AND in the early-exit bool.317 auto comparison =318 mlir::arith::AndIOp::create(rewriter, loc, iterateVar, cmp2);319 mlir::cf::CondBranchOp::create(rewriter, loc, comparison, firstBodyBlock,320 llvm::ArrayRef<mlir::Value>(), endBlock,321 llvm::ArrayRef<mlir::Value>());322 // The result of the loop operation is the values of the condition block323 // arguments except the induction variable on the last iteration.324 auto args = whileOp.getFinalValue()325 ? conditionBlock->getArguments()326 : conditionBlock->getArguments().drop_front();327 rewriter.replaceOp(whileOp, args);328 return success();329 }330 331private:332 bool setNSW;333};334 335/// Convert FIR structured control flow ops to CFG ops.336class CfgConversion : public fir::impl::CFGConversionBase<CfgConversion> {337public:338 using CFGConversionBase<CfgConversion>::CFGConversionBase;339 340 void runOnOperation() override {341 auto *context = &this->getContext();342 mlir::RewritePatternSet patterns(context);343 fir::populateCfgConversionRewrites(patterns, this->forceLoopToExecuteOnce,344 this->setNSW);345 mlir::ConversionTarget target(*context);346 target.addLegalDialect<mlir::affine::AffineDialect,347 mlir::cf::ControlFlowDialect, FIROpsDialect,348 mlir::func::FuncDialect>();349 350 // apply the patterns351 target.addIllegalOp<ResultOp, DoLoopOp, IfOp, IterWhileOp>();352 target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });353 if (mlir::failed(mlir::applyPartialConversion(this->getOperation(), target,354 std::move(patterns)))) {355 mlir::emitError(mlir::UnknownLoc::get(context),356 "error in converting to CFG\n");357 this->signalPassFailure();358 }359 }360};361 362} // namespace363 364/// Expose conversion rewriters to other passes365void fir::populateCfgConversionRewrites(mlir::RewritePatternSet &patterns,366 bool forceLoopToExecuteOnce,367 bool setNSW) {368 patterns.insert<CfgLoopConv, CfgIfConv, CfgIterWhileConv>(369 patterns.getContext(), forceLoopToExecuteOnce, setNSW);370}371