//===-- 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
//
//===----------------------------------------------------------------------===//
//
// This transformation is a prototype that promote FIR loops operations
// to affine dialect operations.
// It is not part of the production pipeline and would need more work in order
// to be used in production.
// More information can be found in this presentation:
// https://slides.com/rajanwalia/deck
//
//===----------------------------------------------------------------------===//
#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/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.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/Support/Debug.h"
#include <optional>
namespace fir {
#define GEN_PASS_DEF_AFFINEDIALECTPROMOTION
#include "flang/Optimizer/Transforms/Passes.h.inc"
} // namespace fir
#define DEBUG_TYPE "flang-affine-promotion"
using namespace fir;
using namespace mlir;
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::func::FuncOp funcOp) {
for (fir::DoLoopOp op : funcOp.getOps<fir::DoLoopOp>())
loopAnalysisMap.try_emplace(op, op, *this);
}
AffineLoopAnalysis getChildLoopAnalysis(fir::DoLoopOp op) const;
AffineIfAnalysis getChildIfAnalysis(fir::IfOp op) const;
llvm::DenseMap<mlir::Operation *, AffineLoopAnalysis> loopAnalysisMap;
llvm::DenseMap<mlir::Operation *, AffineIfAnalysis> ifAnalysisMap;
};
} // namespace
static bool analyzeCoordinate(mlir::Value coordinate, mlir::Operation *op) {
if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(coordinate)) {
if (isa<fir::DoLoopOp>(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<fir::DoLoopOp>()) {
auto analysis = functionAnalysis.loopAnalysisMap
.try_emplace(loopOp, loopOp, functionAnalysis)
.first->getSecond();
if (!analysis.canPromoteToAffine())
return false;
}
for (auto ifOp : loopOperation.getOps<fir::IfOp>())
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<ArrayCoorOp>()) {
if (mlir::isa<fir::BoxType>(acoOp.getMemref().getType())) {
// 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.getIndices())
canPromote = canPromote && analyzeCoordinate(coordinate, op);
return canPromote;
}
if (auto coOp = memref.getDefiningOp<CoordinateOp>()) {
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<fir::LoadOp>())
if (!analyzeReference(loadOp.getMemref(), loadOp))
return false;
for (auto storeOp : loopOperation.getOps<fir::StoreOp>())
if (!analyzeReference(storeOp.getMemref(), 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 = std::optional<mlir::AffineExpr>;
explicit AffineIfCondition(mlir::Value fc) : firCondition(fc) {
if (auto condDef = firCondition.getDefiningOp<mlir::arith::CmpIOp>())
fromCmpIOp(condDef);
}
bool hasIntegerSet() const { return integerSet.has_value(); }
mlir::IntegerSet getIntegerSet() const {
assert(hasIntegerSet() && "integer set is missing");
return *integerSet;
}
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 && rhs)
return mlir::getAffineBinaryOpExpr(kind, *lhs, *rhs);
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<mlir::arith::SubIOp>())
return affineBinaryOp(
mlir::AffineExprKind::Add, toAffineExpr(op.getLhs()),
affineBinaryOp(mlir::AffineExprKind::Mul, toAffineExpr(op.getRhs()),
toAffineExpr(-1)));
if (auto op = value.getDefiningOp<mlir::arith::AddIOp>())
return affineBinaryOp(mlir::AffineExprKind::Add, op.getLhs(),
op.getRhs());
if (auto op = value.getDefiningOp<mlir::arith::MulIOp>())
return affineBinaryOp(mlir::AffineExprKind::Mul, op.getLhs(),
op.getRhs());
if (auto op = value.getDefiningOp<mlir::arith::RemUIOp>())
return affineBinaryOp(mlir::AffineExprKind::Mod, op.getLhs(),
op.getRhs());
if (auto op = value.getDefiningOp<mlir::arith::ConstantOp>())
if (auto intConstant = mlir::dyn_cast<IntegerAttr>(op.getValue()))
return toAffineExpr(intConstant.getInt());
if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(value)) {
affineArgs.push_back(value);
if (isa<fir::DoLoopOp>(blockArg.getOwner()->getParentOp()) ||
isa<mlir::affine::AffineForOp>(blockArg.getOwner()->getParentOp()))
return {mlir::getAffineDimExpr(dimCount++, value.getContext())};
return {mlir::getAffineSymbolExpr(symCount++, value.getContext())};
}
return {};
}
void fromCmpIOp(mlir::arith::CmpIOp cmpOp) {
auto lhsAffine = toAffineExpr(cmpOp.getLhs());
auto rhsAffine = toAffineExpr(cmpOp.getRhs());
if (!lhsAffine || !rhsAffine)
return;
auto constraintPair =
constraint(cmpOp.getPredicate(), *rhsAffine - *lhsAffine);
if (!constraintPair)
return;
integerSet = mlir::IntegerSet::get(
dimCount, symCount, {constraintPair->first}, {constraintPair->second});
}
std::optional<std::pair<AffineExpr, bool>>
constraint(mlir::arith::CmpIPredicate predicate, mlir::AffineExpr basic) {
switch (predicate) {
case mlir::arith::CmpIPredicate::slt:
return {std::make_pair(basic - 1, false)};
case mlir::arith::CmpIPredicate::sle:
return {std::make_pair(basic, false)};
case mlir::arith::CmpIPredicate::sgt:
return {std::make_pair(1 - basic, false)};
case mlir::arith::CmpIPredicate::sge:
return {std::make_pair(0 - basic, false)};
case mlir::arith::CmpIPredicate::eq:
return {std::make_pair(basic, true)};
default:
return {};
}
}
llvm::SmallVector<mlir::Value> affineArgs;
std::optional<mlir::IntegerSet> 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 std::optional<int64_t> constantIntegerLike(const mlir::Value value) {
if (auto definition = value.getDefiningOp<mlir::arith::ConstantOp>())
if (auto stepAttr = mlir::dyn_cast<IntegerAttr>(definition.getValue()))
return stepAttr.getInt();
return {};
}
static mlir::Type coordinateArrayElement(fir::ArrayCoorOp op) {
if (auto refType =
mlir::dyn_cast_or_null<ReferenceType>(op.getMemref().getType())) {
if (auto seqType =
mlir::dyn_cast_or_null<SequenceType>(refType.getEleTy())) {
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<mlir::Value> &indexArgs,
mlir::PatternRewriter &rewriter) {
auto one = rewriter.create<mlir::arith::ConstantOp>(
acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1));
auto extents = shape.getExtents();
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<mlir::Value> &indexArgs,
mlir::PatternRewriter &rewriter) {
auto one = rewriter.create<mlir::arith::ConstantOp>(
acoOp.getLoc(), rewriter.getIndexType(), rewriter.getIndexAttr(1));
auto extents = shape.getPairs();
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<mlir::Value> &indexArgs,
mlir::PatternRewriter &rewriter) {
auto extents = slice.getTriples();
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<mlir::Value> &indexArgs,
mlir::PatternRewriter &rewriter) {
if (auto shape = acoOp.getShape().getDefiningOp<ShapeOp>())
return populateIndexArgs(acoOp, shape, indexArgs, rewriter);
if (auto shapeShift = acoOp.getShape().getDefiningOp<ShapeShiftOp>())
return populateIndexArgs(acoOp, shapeShift, indexArgs, rewriter);
if (auto slice = acoOp.getShape().getDefiningOp<SliceOp>())
return populateIndexArgs(acoOp, slice, indexArgs, rewriter);
}
/// Returns affine.apply and fir.convert from array_coor and gendims
static std::pair<affine::AffineApplyOp, fir::ConvertOp>
createAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter) {
auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>();
auto affineMap =
createArrayIndexAffineMap(acoOp.getIndices().size(), acoOp.getContext());
SmallVector<mlir::Value> indexArgs;
indexArgs.append(acoOp.getIndices().begin(), acoOp.getIndices().end());
populateIndexArgs(acoOp, indexArgs, rewriter);
auto affineApply = rewriter.create<affine::AffineApplyOp>(
acoOp.getLoc(), affineMap, indexArgs);
auto arrayElementType = coordinateArrayElement(acoOp);
auto newType =
mlir::MemRefType::get({mlir::ShapedType::kDynamic}, arrayElementType);
auto arrayConvert = rewriter.create<fir::ConvertOp>(acoOp.getLoc(), newType,
acoOp.getMemref());
return std::make_pair(affineApply, arrayConvert);
}
static void rewriteLoad(fir::LoadOp loadOp, mlir::PatternRewriter &rewriter) {
rewriter.setInsertionPoint(loadOp);
auto affineOps = createAffineOps(loadOp.getMemref(), rewriter);
rewriter.replaceOpWithNewOp<affine::AffineLoadOp>(
loadOp, affineOps.second.getResult(), affineOps.first.getResult());
}
static void rewriteStore(fir::StoreOp storeOp,
mlir::PatternRewriter &rewriter) {
rewriter.setInsertionPoint(storeOp);
auto affineOps = createAffineOps(storeOp.getMemref(), rewriter);
rewriter.replaceOpWithNewOp<affine::AffineStoreOp>(
storeOp, storeOp.getValue(), affineOps.second.getResult(),
affineOps.first.getResult());
}
static void rewriteMemoryOps(Block *block, mlir::PatternRewriter &rewriter) {
for (auto &bodyOp : block->getOperations()) {
if (isa<fir::LoadOp>(bodyOp))
rewriteLoad(cast<fir::LoadOp>(bodyOp), rewriter);
if (isa<fir::StoreOp>(bodyOp))
rewriteStore(cast<fir::StoreOp>(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<fir::DoLoopOp> {
public:
using OpRewritePattern::OpRewritePattern;
AffineLoopConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)
: OpRewritePattern(context), functionAnalysis(afa) {}
llvm::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.startOpModification(affineFor.getOperation());
affineFor.getBody()->getOperations().splice(
std::prev(affineFor.getBody()->end()), loopOps, loopOps.begin(),
std::prev(loopOps.end()));
rewriter.finalizeOpModification(affineFor.getOperation());
rewriter.startOpModification(loop.getOperation());
loop.getInductionVar().replaceAllUsesWith(inductionVar);
rewriter.finalizeOpModification(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<affine::AffineForOp, mlir::Value>
createAffineFor(fir::DoLoopOp op, mlir::PatternRewriter &rewriter) const {
if (auto constantStep = constantIntegerLike(op.getStep()))
if (*constantStep > 0)
return positiveConstantStep(op, *constantStep, rewriter);
return genericBounds(op, rewriter);
}
// when step for the loop is positive compile time constant
std::pair<affine::AffineForOp, mlir::Value>
positiveConstantStep(fir::DoLoopOp op, int64_t step,
mlir::PatternRewriter &rewriter) const {
auto affineFor = rewriter.create<affine::AffineForOp>(
op.getLoc(), ValueRange(op.getLowerBound()),
mlir::AffineMap::get(0, 1,
mlir::getAffineSymbolExpr(0, op.getContext())),
ValueRange(op.getUpperBound()),
mlir::AffineMap::get(0, 1,
1 + mlir::getAffineSymbolExpr(0, op.getContext())),
step);
return std::make_pair(affineFor, affineFor.getInductionVar());
}
std::pair<affine::AffineForOp, mlir::Value>
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<affine::AffineApplyOp>(
op.getLoc(), upperBoundMap,
ValueRange({op.getLowerBound(), op.getUpperBound(), op.getStep()}));
auto actualIndexMap = mlir::AffineMap::get(
1, 2,
(lowerBound + mlir::getAffineDimExpr(0, op.getContext())) *
mlir::getAffineSymbolExpr(1, op.getContext()));
auto affineFor = rewriter.create<affine::AffineForOp>(
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<affine::AffineApplyOp>(
op.getLoc(), actualIndexMap,
ValueRange(
{affineFor.getInductionVar(), op.getLowerBound(), op.getStep()}));
return std::make_pair(affineFor, actualIndex.getResult());
}
AffineFunctionAnalysis &functionAnalysis;
};
/// Convert `fir.if` to `affine.if`.
class AffineIfConversion : public mlir::OpRewritePattern<fir::IfOp> {
public:
using OpRewritePattern::OpRewritePattern;
AffineIfConversion(mlir::MLIRContext *context, AffineFunctionAnalysis &afa)
: OpRewritePattern(context) {}
llvm::LogicalResult
matchAndRewrite(fir::IfOp op,
mlir::PatternRewriter &rewriter) const override {
LLVM_DEBUG(llvm::dbgs() << "AffineIfConversion: rewriting if:\n";
op.dump(););
auto &ifOps = op.getThenRegion().front().getOperations();
auto affineCondition = AffineIfCondition(op.getCondition());
if (!affineCondition.hasIntegerSet()) {
LLVM_DEBUG(
llvm::dbgs()
<< "AffineIfConversion: couldn't calculate affine condition\n";);
return failure();
}
auto affineIf = rewriter.create<affine::AffineIfOp>(
op.getLoc(), affineCondition.getIntegerSet(),
affineCondition.getAffineArgs(), !op.getElseRegion().empty());
rewriter.startOpModification(affineIf);
affineIf.getThenBlock()->getOperations().splice(
std::prev(affineIf.getThenBlock()->end()), ifOps, ifOps.begin(),
std::prev(ifOps.end()));
if (!op.getElseRegion().empty()) {
auto &otherOps = op.getElseRegion().front().getOperations();
affineIf.getElseBlock()->getOperations().splice(
std::prev(affineIf.getElseBlock()->end()), otherOps, otherOps.begin(),
std::prev(otherOps.end()));
}
rewriter.finalizeOpModification(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 fir::impl::AffineDialectPromotionBase<AffineDialectPromotion> {
public:
void runOnOperation() override {
auto *context = &getContext();
auto function = getOperation();
markAllAnalysesPreserved();
auto functionAnalysis = AffineFunctionAnalysis(function);
mlir::RewritePatternSet patterns(context);
patterns.insert<AffineIfConversion>(context, functionAnalysis);
patterns.insert<AffineLoopConversion>(context, functionAnalysis);
mlir::ConversionTarget target = *context;
target.addLegalDialect<mlir::affine::AffineDialect, FIROpsDialect,
mlir::scf::SCFDialect, mlir::arith::ArithDialect,
mlir::func::FuncDialect>();
target.addDynamicallyLegalOp<IfOp>([&functionAnalysis](fir::IfOp op) {
return !(functionAnalysis.getChildIfAnalysis(op).canPromoteToAffine());
});
target.addDynamicallyLegalOp<DoLoopOp>([&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<mlir::Pass> fir::createPromoteToAffinePass() {
return std::make_unique<AffineDialectPromotion>();
}
Codebase Browser
llvm