diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h --- a/flang/include/flang/Optimizer/Transforms/Passes.h +++ b/flang/include/flang/Optimizer/Transforms/Passes.h @@ -22,6 +22,7 @@ namespace fir { +std::unique_ptr createPromoteToAffinePass(); std::unique_ptr createExternalNameConversionPass(); /// Support for inlining on FIR. diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td --- a/flang/include/flang/Optimizer/Transforms/Passes.td +++ b/flang/include/flang/Optimizer/Transforms/Passes.td @@ -16,6 +16,31 @@ include "mlir/Pass/PassBase.td" +def AffineDialectPromotion : FunctionPass<"promote-to-affine"> { + let summary = "Promotes `fir.{do_loop,if}` to `affine.{for,if}`."; + let description = [{ + Convert fir operations which satisfy affine constraints to the affine + dialect. + + `fir.do_loop` will be converted to `affine.for` if the loops inside the body + can be converted and the indices for memory loads and stores satisfy + `affine.apply` criteria for symbols and dimensions. + + `fir.if` will be converted to `affine.if` where possible. `affine.if`'s + condition uses an integer set (==, >=) and an analysis is done to determine + the fir condition's parent operations to construct the integer set. + + `fir.load` (`fir.store`) will be converted to `affine.load` (`affine.store`) + where possible. This conversion includes adding a dummy `fir.convert` cast + to adapt values of type `!fir.ref` to `memref`. This is done + because the affine dialect presently only understands the `memref` type. + }]; + let constructor = "::fir::createPromoteToAffinePass()"; + let dependentDialects = [ + "fir::FIROpsDialect", "mlir::StandardOpsDialect", "mlir::AffineDialect" + ]; +} + def ExternalNameConversion : Pass<"external-name-interop", "mlir::ModuleOp"> { let summary = "Convert name for external interoperability"; let description = [{ diff --git a/flang/lib/Optimizer/Transforms/AffinePromotion.cpp b/flang/lib/Optimizer/Transforms/AffinePromotion.cpp new file mode 100644 --- /dev/null +++ b/flang/lib/Optimizer/Transforms/AffinePromotion.cpp @@ -0,0 +1,609 @@ +//===-- AffinePromotion.cpp -----------------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "PassDetail.h" +#include "flang/Optimizer/Dialect/FIRDialect.h" +#include "flang/Optimizer/Dialect/FIROps.h" +#include "flang/Optimizer/Dialect/FIRType.h" +#include "flang/Optimizer/Transforms/Passes.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/SCF/SCF.h" +#include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/IntegerSet.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Transforms/DialectConversion.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/Optional.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "flang-affine-promotion" + +using namespace fir; + +namespace { +struct AffineLoopAnalysis; +struct AffineIfAnalysis; + +/// Stores analysis objects for all loops and if operations inside a function +/// these analysis are used twice, first for marking operations for rewrite and +/// second when doing rewrite. +struct AffineFunctionAnalysis { + explicit AffineFunctionAnalysis(mlir::FuncOp funcOp) { + for (fir::DoLoopOp op : funcOp.getOps()) + loopAnalysisMap.try_emplace(op, op, *this); + } + + AffineLoopAnalysis getChildLoopAnalysis(fir::DoLoopOp op) const; + + AffineIfAnalysis getChildIfAnalysis(fir::IfOp op) const; + + llvm::DenseMap loopAnalysisMap; + llvm::DenseMap ifAnalysisMap; +}; +} // namespace + +static bool analyzeCoordinate(mlir::Value coordinate, mlir::Operation *op) { + if (auto blockArg = coordinate.dyn_cast()) { + if (isa(blockArg.getOwner()->getParentOp())) + return true; + LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a " + "loop induction variable (owner not loopOp)\n"; + op->dump()); + return false; + } + LLVM_DEBUG( + llvm::dbgs() << "AffineLoopAnalysis: array coordinate is not a loop " + "induction variable (not a block argument)\n"; + op->dump(); coordinate.getDefiningOp()->dump()); + return false; +} + +namespace { +struct AffineLoopAnalysis { + AffineLoopAnalysis() = default; + + explicit AffineLoopAnalysis(fir::DoLoopOp op, AffineFunctionAnalysis &afa) + : legality(analyzeLoop(op, afa)) {} + + bool canPromoteToAffine() { return legality; } + +private: + bool analyzeBody(fir::DoLoopOp loopOperation, + AffineFunctionAnalysis &functionAnalysis) { + for (auto loopOp : loopOperation.getOps()) { + auto analysis = functionAnalysis.loopAnalysisMap + .try_emplace(loopOp, loopOp, functionAnalysis) + .first->getSecond(); + if (!analysis.canPromoteToAffine()) + return false; + } + for (auto ifOp : loopOperation.getOps()) + functionAnalysis.ifAnalysisMap.try_emplace(ifOp, ifOp, functionAnalysis); + return true; + } + + bool analyzeLoop(fir::DoLoopOp loopOperation, + AffineFunctionAnalysis &functionAnalysis) { + LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: \n"; loopOperation.dump();); + return analyzeMemoryAccess(loopOperation) && + analyzeBody(loopOperation, functionAnalysis); + } + + bool analyzeReference(mlir::Value memref, mlir::Operation *op) { + if (auto acoOp = memref.getDefiningOp()) { + if (acoOp.memref().getType().isa()) { + // TODO: Look if and how fir.box can be promoted to affine. + LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: cannot promote loop, " + "array memory operation uses fir.box\n"; + op->dump(); acoOp.dump();); + return false; + } + bool canPromote = true; + for (auto coordinate : acoOp.indices()) + canPromote = canPromote && analyzeCoordinate(coordinate, op); + return canPromote; + } + if (auto coOp = memref.getDefiningOp()) { + LLVM_DEBUG(llvm::dbgs() + << "AffineLoopAnalysis: cannot promote loop, " + "array memory operation uses non ArrayCoorOp\n"; + op->dump(); coOp.dump();); + + return false; + } + LLVM_DEBUG(llvm::dbgs() << "AffineLoopAnalysis: unknown type of memory " + "reference for array load\n"; + op->dump();); + return false; + } + + bool analyzeMemoryAccess(fir::DoLoopOp loopOperation) { + for (auto loadOp : loopOperation.getOps()) + if (!analyzeReference(loadOp.memref(), loadOp)) + return false; + for (auto storeOp : loopOperation.getOps()) + if (!analyzeReference(storeOp.memref(), storeOp)) + return false; + return true; + } + + bool legality{}; +}; +} // namespace + +AffineLoopAnalysis +AffineFunctionAnalysis::getChildLoopAnalysis(fir::DoLoopOp op) const { + auto it = loopAnalysisMap.find_as(op); + if (it == loopAnalysisMap.end()) { + LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n"; + op.dump();); + op.emitError("error in fetching loop analysis in AffineFunctionAnalysis\n"); + return {}; + } + return it->getSecond(); +} + +namespace { +/// Calculates arguments for creating an IntegerSet. symCount, dimCount are the +/// final number of symbols and dimensions of the affine map. Integer set if +/// possible is in Optional IntegerSet. +struct AffineIfCondition { + using MaybeAffineExpr = llvm::Optional; + + explicit AffineIfCondition(mlir::Value fc) : firCondition(fc) { + if (auto condDef = firCondition.getDefiningOp()) + fromCmpIOp(condDef); + } + + bool hasIntegerSet() const { return integerSet.hasValue(); } + + mlir::IntegerSet getIntegerSet() const { + assert(hasIntegerSet() && "integer set is missing"); + return integerSet.getValue(); + } + + mlir::ValueRange getAffineArgs() const { return affineArgs; } + +private: + MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, mlir::Value lhs, + mlir::Value rhs) { + return affineBinaryOp(kind, toAffineExpr(lhs), toAffineExpr(rhs)); + } + + MaybeAffineExpr affineBinaryOp(mlir::AffineExprKind kind, MaybeAffineExpr lhs, + MaybeAffineExpr rhs) { + if (lhs.hasValue() && rhs.hasValue()) + return mlir::getAffineBinaryOpExpr(kind, lhs.getValue(), rhs.getValue()); + return {}; + } + + MaybeAffineExpr toAffineExpr(MaybeAffineExpr e) { return e; } + + MaybeAffineExpr toAffineExpr(int64_t value) { + return {mlir::getAffineConstantExpr(value, firCondition.getContext())}; + } + + /// Returns an AffineExpr if it is a result of operations that can be done + /// in an affine expression, this includes -, +, *, rem, constant. + /// block arguments of a loopOp or forOp are used as dimensions + MaybeAffineExpr toAffineExpr(mlir::Value value) { + if (auto op = value.getDefiningOp()) + return affineBinaryOp(mlir::AffineExprKind::Add, toAffineExpr(op.lhs()), + affineBinaryOp(mlir::AffineExprKind::Mul, + toAffineExpr(op.rhs()), + toAffineExpr(-1))); + if (auto op = value.getDefiningOp()) + return affineBinaryOp(mlir::AffineExprKind::Add, op.lhs(), op.rhs()); + if (auto op = value.getDefiningOp()) + return affineBinaryOp(mlir::AffineExprKind::Mul, op.lhs(), op.rhs()); + if (auto op = value.getDefiningOp()) + return affineBinaryOp(mlir::AffineExprKind::Mod, op.lhs(), op.rhs()); + if (auto op = value.getDefiningOp()) + if (auto intConstant = op.getValue().dyn_cast()) + return toAffineExpr(intConstant.getInt()); + if (auto blockArg = value.dyn_cast()) { + affineArgs.push_back(value); + if (isa(blockArg.getOwner()->getParentOp()) || + isa(blockArg.getOwner()->getParentOp())) + return {mlir::getAffineDimExpr(dimCount++, value.getContext())}; + return {mlir::getAffineSymbolExpr(symCount++, value.getContext())}; + } + return {}; + } + + void fromCmpIOp(mlir::CmpIOp cmpOp) { + auto lhsAffine = toAffineExpr(cmpOp.lhs()); + auto rhsAffine = toAffineExpr(cmpOp.rhs()); + if (!lhsAffine.hasValue() || !rhsAffine.hasValue()) + return; + auto constraintPair = constraint( + cmpOp.predicate(), rhsAffine.getValue() - lhsAffine.getValue()); + if (!constraintPair) + return; + integerSet = mlir::IntegerSet::get(dimCount, symCount, + {constraintPair.getValue().first}, + {constraintPair.getValue().second}); + return; + } + + llvm::Optional> + constraint(mlir::CmpIPredicate predicate, mlir::AffineExpr basic) { + switch (predicate) { + case mlir::CmpIPredicate::slt: + return {std::make_pair(basic - 1, false)}; + case mlir::CmpIPredicate::sle: + return {std::make_pair(basic, false)}; + case mlir::CmpIPredicate::sgt: + return {std::make_pair(1 - basic, false)}; + case mlir::CmpIPredicate::sge: + return {std::make_pair(0 - basic, false)}; + case mlir::CmpIPredicate::eq: + return {std::make_pair(basic, true)}; + default: + return {}; + } + } + + llvm::SmallVector affineArgs; + llvm::Optional integerSet; + mlir::Value firCondition; + unsigned symCount{0u}; + unsigned dimCount{0u}; +}; +} // namespace + +namespace { +/// Analysis for affine promotion of fir.if +struct AffineIfAnalysis { + AffineIfAnalysis() = default; + + explicit AffineIfAnalysis(fir::IfOp op, AffineFunctionAnalysis &afa) + : legality(analyzeIf(op, afa)) {} + + bool canPromoteToAffine() { return legality; } + +private: + bool analyzeIf(fir::IfOp op, AffineFunctionAnalysis &afa) { + if (op.getNumResults() == 0) + return true; + LLVM_DEBUG(llvm::dbgs() + << "AffineIfAnalysis: not promoting as op has results\n";); + return false; + } + + bool legality{}; +}; +} // namespace + +AffineIfAnalysis +AffineFunctionAnalysis::getChildIfAnalysis(fir::IfOp op) const { + auto it = ifAnalysisMap.find_as(op); + if (it == ifAnalysisMap.end()) { + LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n"; + op.dump();); + op.emitError("error in fetching if analysis in AffineFunctionAnalysis\n"); + return {}; + } + return it->getSecond(); +} + +/// AffineMap rewriting fir.array_coor operation to affine apply, +/// %dim = fir.gendim %lowerBound, %upperBound, %stride +/// %a = fir.array_coor %arr(%dim) %i +/// returning affineMap = affine_map<(i)[lb, ub, st] -> (i*st - lb)> +static mlir::AffineMap createArrayIndexAffineMap(unsigned dimensions, + MLIRContext *context) { + auto index = mlir::getAffineConstantExpr(0, context); + auto accuExtent = mlir::getAffineConstantExpr(1, context); + for (unsigned i = 0; i < dimensions; ++i) { + mlir::AffineExpr idx = mlir::getAffineDimExpr(i, context), + lowerBound = mlir::getAffineSymbolExpr(i * 3, context), + currentExtent = + mlir::getAffineSymbolExpr(i * 3 + 1, context), + stride = mlir::getAffineSymbolExpr(i * 3 + 2, context), + currentPart = (idx * stride - lowerBound) * accuExtent; + index = currentPart + index; + accuExtent = accuExtent * currentExtent; + } + return mlir::AffineMap::get(dimensions, dimensions * 3, index); +} + +static Optional constantIntegerLike(const mlir::Value value) { + if (auto definition = value.getDefiningOp()) + if (auto stepAttr = definition.getValue().dyn_cast()) + return stepAttr.getInt(); + return {}; +} + +static mlir::Type coordinateArrayElement(fir::ArrayCoorOp op) { + if (auto refType = op.memref().getType().dyn_cast_or_null()) { + if (auto seqType = refType.getEleTy().dyn_cast_or_null()) { + return seqType.getEleTy(); + } + } + op.emitError( + "AffineLoopConversion: array type in coordinate operation not valid\n"); + return mlir::Type(); +} + +static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeOp shape, + SmallVectorImpl &indexArgs, + mlir::PatternRewriter &rewriter) { + auto one = rewriter.create( + acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1)); + auto extents = shape.extents(); + for (auto i = extents.begin(); i < extents.end(); i++) { + indexArgs.push_back(one); + indexArgs.push_back(*i); + indexArgs.push_back(one); + } +} + +static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeShiftOp shape, + SmallVectorImpl &indexArgs, + mlir::PatternRewriter &rewriter) { + auto one = rewriter.create( + acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1)); + auto extents = shape.pairs(); + for (auto i = extents.begin(); i < extents.end();) { + indexArgs.push_back(*i++); + indexArgs.push_back(*i++); + indexArgs.push_back(one); + } +} + +static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::SliceOp slice, + SmallVectorImpl &indexArgs, + mlir::PatternRewriter &rewriter) { + auto extents = slice.triples(); + for (auto i = extents.begin(); i < extents.end();) { + indexArgs.push_back(*i++); + indexArgs.push_back(*i++); + indexArgs.push_back(*i++); + } +} + +static void populateIndexArgs(fir::ArrayCoorOp acoOp, + SmallVectorImpl &indexArgs, + mlir::PatternRewriter &rewriter) { + if (auto shape = acoOp.shape().getDefiningOp()) + return populateIndexArgs(acoOp, shape, indexArgs, rewriter); + if (auto shapeShift = acoOp.shape().getDefiningOp()) + return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter); + if (auto slice = acoOp.shape().getDefiningOp()) + return populateIndexArgs(acoOp, slice, indexArgs, rewriter); + return; +} + +/// Returns affine.apply and fir.convert from array_coor and gendims +static std::pair +createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) { + auto acoOp = arrayRef.getDefiningOp(); + auto affineMap = + createArrayIndexAffineMap(acoOp.indices().size(), acoOp.getContext()); + SmallVector indexArgs; + indexArgs.append(acoOp.indices().begin(), acoOp.indices().end()); + + populateIndexArgs(acoOp, indexArgs, rewriter); + + auto affineApply = rewriter.create(acoOp.getLoc(), + affineMap, indexArgs); + auto arrayElementType = coordinateArrayElement(acoOp); + auto newType = mlir::MemRefType::get({-1}, arrayElementType); + auto arrayConvert = + rewriter.create(acoOp.getLoc(), newType, acoOp.memref()); + return std::make_pair(affineApply, arrayConvert); +} + +static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) { + rewriter.setInsertionPoint(loadOp); + auto affineOps = createAffineOps(loadOp.memref(), rewriter); + rewriter.replaceOpWithNewOp( + loadOp, affineOps.second.getResult(), affineOps.first.getResult()); +} + +static void rewriteStore(fir::StoreOp storeOp, + mlir::PatternRewriter &rewriter) { + rewriter.setInsertionPoint(storeOp); + auto affineOps = createAffineOps(storeOp.memref(), rewriter); + rewriter.replaceOpWithNewOp(storeOp, storeOp.value(), + affineOps.second.getResult(), + affineOps.first.getResult()); +} + +static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) { + for (auto &bodyOp : block->getOperations()) { + if (isa(bodyOp)) + rewriteLoad(cast(bodyOp), rewriter); + if (isa(bodyOp)) + rewriteStore(cast(bodyOp), rewriter); + } +} + +namespace { +/// Convert `fir.do_loop` to `affine.for`, creates fir.convert for arrays to +/// memref, rewrites array_coor to affine.apply with affine_map. Rewrites fir +/// loads and stores to affine. +class AffineLoopConversion : public mlir::OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + AffineLoopConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa) + : OpRewritePattern(context), functionAnalysis(afa) {} + + mlir::LogicalResult + matchAndRewrite(fir::DoLoopOp loop, + mlir::PatternRewriter &rewriter) const override { + LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: rewriting loop:\n"; + loop.dump();); + LLVM_ATTRIBUTE_UNUSED auto loopAnalysis = + functionAnalysis.getChildLoopAnalysis(loop); + auto &loopOps = loop.getBody()->getOperations(); + auto loopAndIndex = createAffineFor(loop, rewriter); + auto affineFor = loopAndIndex.first; + auto inductionVar = loopAndIndex.second; + + rewriter.startRootUpdate(affineFor.getOperation()); + affineFor.getBody()->getOperations().splice( + std::prev(affineFor.getBody()->end()), loopOps, loopOps.begin(), + std::prev(loopOps.end())); + rewriter.finalizeRootUpdate(affineFor.getOperation()); + + rewriter.startRootUpdate(loop.getOperation()); + loop.getInductionVar().replaceAllUsesWith(inductionVar); + rewriter.finalizeRootUpdate(loop.getOperation()); + + rewriteMemoryOps(affineFor.getBody(), rewriter); + + LLVM_DEBUG(llvm::dbgs() << "AffineLoopConversion: loop rewriten to:\n"; + affineFor.dump();); + rewriter.replaceOp(loop, affineFor.getOperation()->getResults()); + return success(); + } + +private: + std::pair + createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const { + if (auto constantStep = constantIntegerLike(op.step())) + if (constantStep.getValue() > 0) + return positiveConstantStep(op, constantStep.getValue(), rewriter); + return genericBounds(op, rewriter); + } + + // when step for the loop is positive compile time constant + std::pair + positiveConstantStep(fir::DoLoopOp op, int64_t step, + mlir::PatternRewriter &rewriter) const { + auto affineFor = rewriter.create( + op.getLoc(), ValueRange(op.lowerBound()), + mlir::AffineMap::get(0, 1, + mlir::getAffineSymbolExpr(0, op.getContext())), + ValueRange(op.upperBound()), + mlir::AffineMap::get(0, 1, + 1 + mlir::getAffineSymbolExpr(0, op.getContext())), + step); + return std::make_pair(affineFor, affineFor.getInductionVar()); + } + + std::pair + genericBounds(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const { + auto lowerBound = mlir::getAffineSymbolExpr(0, op.getContext()); + auto upperBound = mlir::getAffineSymbolExpr(1, op.getContext()); + auto step = mlir::getAffineSymbolExpr(2, op.getContext()); + mlir::AffineMap upperBoundMap = mlir::AffineMap::get( + 0, 3, (upperBound - lowerBound + step).floorDiv(step)); + auto genericUpperBound = rewriter.create( + op.getLoc(), upperBoundMap, + ValueRange({op.lowerBound(), op.upperBound(), op.step()})); + auto actualIndexMap = mlir::AffineMap::get( + 1, 2, + (lowerBound + mlir::getAffineDimExpr(0, op.getContext())) * + mlir::getAffineSymbolExpr(1, op.getContext())); + + auto affineFor = rewriter.create( + op.getLoc(), ValueRange(), + AffineMap::getConstantMap(0, op.getContext()), + genericUpperBound.getResult(), + mlir::AffineMap::get(0, 1, + 1 + mlir::getAffineSymbolExpr(0, op.getContext())), + 1); + rewriter.setInsertionPointToStart(affineFor.getBody()); + auto actualIndex = rewriter.create( + op.getLoc(), actualIndexMap, + ValueRange({affineFor.getInductionVar(), op.lowerBound(), op.step()})); + return std::make_pair(affineFor, actualIndex.getResult()); + } + + AffineFunctionAnalysis &functionAnalysis; +}; + +/// Convert `fir.if` to `affine.if`. +class AffineIfConversion : public mlir::OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + AffineIfConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa) + : OpRewritePattern(context) {} + mlir::LogicalResult + matchAndRewrite(fir::IfOp op, + mlir::PatternRewriter &rewriter) const override { + LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: rewriting if:\n"; + op.dump();); + auto &ifOps = op.thenRegion().front().getOperations(); + auto affineCondition = AffineIfCondition(op.condition()); + if (!affineCondition.hasIntegerSet()) { + LLVM_DEBUG( + llvm::dbgs() + << "AffineIfConversion: couldn't calculate affine condition\n";); + return failure(); + } + auto affineIf = rewriter.create( + op.getLoc(), affineCondition.getIntegerSet(), + affineCondition.getAffineArgs(), !op.elseRegion().empty()); + rewriter.startRootUpdate(affineIf); + affineIf.getThenBlock()->getOperations().splice( + std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(), + std::prev(ifOps.end())); + if (!op.elseRegion().empty()) { + auto &otherOps = op.elseRegion().front().getOperations(); + affineIf.getElseBlock()->getOperations().splice( + std::prev(affineIf.getElseBlock()->end()), otherOps, otherOps.begin(), + std::prev(otherOps.end())); + } + rewriter.finalizeRootUpdate(affineIf); + rewriteMemoryOps(affineIf.getBody(), rewriter); + + LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: if converted to:\n"; + affineIf.dump();); + rewriter.replaceOp(op, affineIf.getOperation()->getResults()); + return success(); + } +}; + +/// Promote fir.do_loop and fir.if to affine.for and affine.if, in the cases +/// where such a promotion is possible. +class AffineDialectPromotion + : public AffineDialectPromotionBase { +public: + void runOnFunction() override { + + auto *context = &getContext(); + auto function = getFunction(); + markAllAnalysesPreserved(); + auto functionAnalysis = AffineFunctionAnalysis(function); + mlir::OwningRewritePatternList patterns(context); + patterns.insert(context, functionAnalysis); + patterns.insert(context, functionAnalysis); + mlir::ConversionTarget target = *context; + target.addLegalDialect(); + target.addDynamicallyLegalOp([&functionAnalysis](fir::IfOp op) { + return !(functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine()); + }); + target.addDynamicallyLegalOp([&functionAnalysis]( + fir::DoLoopOp op) { + return !(functionAnalysis.getChildLoopAnalysis(op).canPromoteToAffine()); + }); + + LLVM_DEBUG(llvm::dbgs() + << "AffineDialectPromotion: running promotion on: \n"; + function.print(llvm::dbgs());); + // apply the patterns + if (mlir::failed(mlir::applyPartialConversion(function, target, + std::move(patterns)))) { + mlir::emitError(mlir::UnknownLoc::get(context), + "error in converting to affine dialect\n"); + signalPassFailure(); + } + } +}; +} // namespace + +/// Convert FIR loop constructs to the Affine dialect +std::unique_ptr fir::createPromoteToAffinePass() { + return std::make_unique(); +} diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt --- a/flang/lib/Optimizer/Transforms/CMakeLists.txt +++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt @@ -1,4 +1,5 @@ add_flang_library(FIRTransforms + AffinePromotion.cpp Inliner.cpp ExternalNameConversion.cpp diff --git a/flang/lib/Optimizer/Transforms/PassDetail.h b/flang/lib/Optimizer/Transforms/PassDetail.h --- a/flang/lib/Optimizer/Transforms/PassDetail.h +++ b/flang/lib/Optimizer/Transforms/PassDetail.h @@ -9,6 +9,7 @@ #define FORTRAN_OPTMIZER_TRANSFORMS_PASSDETAIL_H #include "flang/Optimizer/Dialect/FIRDialect.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/OpenACC/OpenACC.h" #include "mlir/Dialect/OpenMP/OpenMPDialect.h" diff --git a/flang/test/Fir/affine-promotion.fir b/flang/test/Fir/affine-promotion.fir new file mode 100644 --- /dev/null +++ b/flang/test/Fir/affine-promotion.fir @@ -0,0 +1,133 @@ +// Test loop affine promotion + +// RUN: fir-opt --split-input-file --promote-to-affine --affine-loop-invariant-code-motion --cse %s | FileCheck %s + +!arr_d1 = type !fir.ref> +#arr_len = affine_map<()[j1,k1] -> (k1 - j1 + 1)> + +func @calc(%a1: !arr_d1, %a2: !arr_d1, %a3: !arr_d1) { + %c1 = constant 1 : index + %c0 = constant 0 : index + %len = constant 100 : index + %dims = fir.shape %len : (index) -> !fir.shape<1> + %siz = affine.apply #arr_len()[%c1,%len] + %t1 = fir.alloca !fir.array, %siz + + fir.do_loop %i = %c1 to %len step %c1 { + %a1_idx = fir.array_coor %a1(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + %a1_v = fir.load %a1_idx : !fir.ref + + %a2_idx = fir.array_coor %a2(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + %a2_v = fir.load %a2_idx : !fir.ref + + %v = addf %a1_v, %a2_v : f32 + %t1_idx = fir.array_coor %t1(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + + fir.store %v to %t1_idx : !fir.ref + } + fir.do_loop %i = %c1 to %len step %c1 { + %t1_idx = fir.array_coor %t1(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + %t1_v = fir.load %t1_idx : !fir.ref + + %a2_idx = fir.array_coor %a2(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + %a2_v = fir.load %a2_idx : !fir.ref + + %v = mulf %t1_v, %a2_v : f32 + %a3_idx = fir.array_coor %a3(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + + fir.store %v to %a3_idx : !fir.ref + } + return +} + +// CHECK: func @calc(%[[VAL_0:.*]]: !fir.ref>, %[[VAL_1:.*]]: !fir.ref>, %[[VAL_2:.*]]: !fir.ref>) { +// CHECK: %[[VAL_3:.*]] = constant 1 : index +// CHECK: %[[VAL_4:.*]] = constant 100 : index +// CHECK: %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_6:.*]] = affine.apply #map0(){{\[}}%[[VAL_3]], %[[VAL_4]]] +// CHECK: %[[VAL_7:.*]] = fir.alloca !fir.array, %[[VAL_6]] +// CHECK: %[[VAL_8:.*]] = fir.convert %[[VAL_0]] : (!fir.ref>) -> memref +// CHECK: %[[VAL_9:.*]] = fir.convert %[[VAL_1]] : (!fir.ref>) -> memref +// CHECK: %[[VAL_10:.*]] = fir.convert %[[VAL_7]] : (!fir.ref>) -> memref +// CHECK: affine.for %[[VAL_11:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] { +// CHECK: %[[VAL_12:.*]] = affine.apply #map2(%[[VAL_11]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]] +// CHECK: %[[VAL_13:.*]] = affine.load %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref +// CHECK: %[[VAL_14:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref +// CHECK: %[[VAL_15:.*]] = addf %[[VAL_13]], %[[VAL_14]] : f32 +// CHECK: affine.store %[[VAL_15]], %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref +// CHECK: } +// CHECK: %[[VAL_16:.*]] = fir.convert %[[VAL_2]] : (!fir.ref>) -> memref +// CHECK: affine.for %[[VAL_17:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] { +// CHECK: %[[VAL_18:.*]] = affine.apply #map2(%[[VAL_17]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]] +// CHECK: %[[VAL_19:.*]] = affine.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref +// CHECK: %[[VAL_20:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref +// CHECK: %[[VAL_21:.*]] = mulf %[[VAL_19]], %[[VAL_20]] : f32 +// CHECK: affine.store %[[VAL_21]], %[[VAL_16]]{{\[}}%[[VAL_18]]] : memref +// CHECK: } +// CHECK: return +// CHECK: } + +// ----- + +!arr_d1 = type !fir.ref> +#arr_len = affine_map<()[j1,k1] -> (k1 - j1 + 1)> + +func @calc(%a: !arr_d1, %v: f32) { + %c0 = constant 0 : index + %c1 = constant 1 : index + %c2 = constant 2 : index + %len = constant 100 : index + %dims = fir.shape %len : (index) -> !fir.shape<1> + + fir.do_loop %i = %c1 to %len step %c1 { + fir.do_loop %j = %c1 to %len step %c1 { + fir.do_loop %k = %c1 to %len step %c1 { + %im2 = subi %i, %c2 : index + %cond = cmpi "sgt", %im2, %c0 : index + fir.if %cond { + %a_idx = fir.array_coor %a(%dims) %i + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + fir.store %v to %a_idx : !fir.ref + } + %aj_idx = fir.array_coor %a(%dims) %j + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + fir.store %v to %aj_idx : !fir.ref + %ak_idx = fir.array_coor %a(%dims) %k + : (!arr_d1, !fir.shape<1>, index) -> !fir.ref + fir.store %v to %ak_idx : !fir.ref + } + } + } + return +} + +// CHECK: func @calc(%[[VAL_0:.*]]: !fir.ref>, %[[VAL_1:.*]]: f32) { +// CHECK: %[[VAL_2:.*]] = constant 0 : index +// CHECK: %[[VAL_3:.*]] = constant 1 : index +// CHECK: %[[VAL_4:.*]] = constant 2 : index +// CHECK: %[[VAL_5:.*]] = constant 100 : index +// CHECK: %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1> +// CHECK: %[[VAL_7:.*]] = fir.convert %[[VAL_0]] : (!fir.ref>) -> memref +// CHECK: affine.for %[[VAL_8:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] { +// CHECK: %[[VAL_9:.*]] = affine.apply #map1(%[[VAL_8]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]] +// CHECK: affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_9]]] : memref +// CHECK: } +// CHECK: affine.for %[[VAL_10:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] { +// CHECK: %[[VAL_11:.*]] = affine.apply #map1(%[[VAL_10]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]] +// CHECK: affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref +// CHECK: } +// CHECK: affine.for %[[VAL_12:.*]] = %[[VAL_3]] to #map0(){{\[}}%[[VAL_5]]] { +// CHECK: %[[VAL_13:.*]] = subi %[[VAL_12]], %[[VAL_4]] : index +// CHECK: affine.if #set(%[[VAL_12]]) { +// CHECK: %[[VAL_14:.*]] = affine.apply #map1(%[[VAL_12]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]] +// CHECK: affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref +// CHECK: } +// CHECK: } +// CHECK: return +// CHECK: }