182 lines · cpp
1//===-- AffineDemotion.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 demote affine dialects operations10// after optimizations to FIR loops operations.11// It is used after the AffinePromotion pass.12// It is not part of the production pipeline and would need more work in order13// to be used in production.14// More information can be found in this presentation:15// https://slides.com/rajanwalia/deck16//17//===----------------------------------------------------------------------===//18 19#include "flang/Optimizer/Dialect/FIRDialect.h"20#include "flang/Optimizer/Dialect/FIROps.h"21#include "flang/Optimizer/Dialect/FIRType.h"22#include "flang/Optimizer/Transforms/Passes.h"23#include "mlir/Dialect/Affine/IR/AffineOps.h"24#include "mlir/Dialect/Affine/Utils.h"25#include "mlir/Dialect/Func/IR/FuncOps.h"26#include "mlir/Dialect/MemRef/IR/MemRef.h"27#include "mlir/Dialect/SCF/IR/SCF.h"28#include "mlir/IR/BuiltinAttributes.h"29#include "mlir/IR/IntegerSet.h"30#include "mlir/IR/Visitors.h"31#include "mlir/Pass/Pass.h"32#include "mlir/Transforms/DialectConversion.h"33#include "llvm/ADT/DenseMap.h"34#include "llvm/Support/CommandLine.h"35#include "llvm/Support/Debug.h"36 37namespace fir {38#define GEN_PASS_DEF_AFFINEDIALECTDEMOTION39#include "flang/Optimizer/Transforms/Passes.h.inc"40} // namespace fir41 42#define DEBUG_TYPE "flang-affine-demotion"43 44using namespace fir;45using namespace mlir;46 47namespace {48 49class AffineLoadConversion50 : public OpConversionPattern<mlir::affine::AffineLoadOp> {51public:52 using OpConversionPattern<mlir::affine::AffineLoadOp>::OpConversionPattern;53 54 LogicalResult55 matchAndRewrite(mlir::affine::AffineLoadOp op, OpAdaptor adaptor,56 ConversionPatternRewriter &rewriter) const override {57 SmallVector<Value> indices(adaptor.getIndices());58 auto maybeExpandedMap = affine::expandAffineMap(rewriter, op.getLoc(),59 op.getAffineMap(), indices);60 if (!maybeExpandedMap)61 return failure();62 63 auto coorOp = fir::CoordinateOp::create(64 rewriter, op.getLoc(),65 fir::ReferenceType::get(op.getResult().getType()), adaptor.getMemref(),66 *maybeExpandedMap);67 68 rewriter.replaceOpWithNewOp<fir::LoadOp>(op, coorOp.getResult());69 return success();70 }71};72 73class AffineStoreConversion74 : public OpConversionPattern<mlir::affine::AffineStoreOp> {75public:76 using OpConversionPattern<mlir::affine::AffineStoreOp>::OpConversionPattern;77 78 LogicalResult79 matchAndRewrite(mlir::affine::AffineStoreOp op, OpAdaptor adaptor,80 ConversionPatternRewriter &rewriter) const override {81 SmallVector<Value> indices(op.getIndices());82 auto maybeExpandedMap = affine::expandAffineMap(rewriter, op.getLoc(),83 op.getAffineMap(), indices);84 if (!maybeExpandedMap)85 return failure();86 87 auto coorOp = fir::CoordinateOp::create(88 rewriter, op.getLoc(),89 fir::ReferenceType::get(op.getValueToStore().getType()),90 adaptor.getMemref(), *maybeExpandedMap);91 rewriter.replaceOpWithNewOp<fir::StoreOp>(op, adaptor.getValue(),92 coorOp.getResult());93 return success();94 }95};96 97class ConvertConversion : public mlir::OpRewritePattern<fir::ConvertOp> {98public:99 using OpRewritePattern::OpRewritePattern;100 llvm::LogicalResult101 matchAndRewrite(fir::ConvertOp op,102 mlir::PatternRewriter &rewriter) const override {103 if (mlir::isa<mlir::MemRefType>(op.getRes().getType())) {104 // due to index calculation moving to affine maps we still need to105 // add converts for sequence types this has a side effect of losing106 // some information about arrays with known dimensions by creating:107 // fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) ->108 // !fir.ref<!fir.array<?xi32>>109 if (auto refTy =110 mlir::dyn_cast<fir::ReferenceType>(op.getValue().getType()))111 if (auto arrTy = mlir::dyn_cast<fir::SequenceType>(refTy.getEleTy())) {112 fir::SequenceType::Shape flatShape = {113 fir::SequenceType::getUnknownExtent()};114 auto flatArrTy = fir::SequenceType::get(flatShape, arrTy.getEleTy());115 auto flatTy = fir::ReferenceType::get(flatArrTy);116 rewriter.replaceOpWithNewOp<fir::ConvertOp>(op, flatTy,117 op.getValue());118 return success();119 }120 rewriter.replaceOp(op, op.getValue());121 }122 return success();123 }124};125 126mlir::Type convertMemRef(mlir::MemRefType type) {127 return fir::SequenceType::get(SmallVector<int64_t>(type.getShape()),128 type.getElementType());129}130 131class StdAllocConversion : public mlir::OpRewritePattern<memref::AllocOp> {132public:133 using OpRewritePattern::OpRewritePattern;134 llvm::LogicalResult135 matchAndRewrite(memref::AllocOp op,136 mlir::PatternRewriter &rewriter) const override {137 rewriter.replaceOpWithNewOp<fir::AllocaOp>(op, convertMemRef(op.getType()),138 op.getMemref());139 return success();140 }141};142 143class AffineDialectDemotion144 : public fir::impl::AffineDialectDemotionBase<AffineDialectDemotion> {145public:146 void runOnOperation() override {147 auto *context = &getContext();148 auto function = getOperation();149 LLVM_DEBUG(llvm::dbgs() << "AffineDemotion: running on function:\n";150 function.print(llvm::dbgs()););151 152 mlir::RewritePatternSet patterns(context);153 patterns.insert<ConvertConversion>(context);154 patterns.insert<AffineLoadConversion>(context);155 patterns.insert<AffineStoreConversion>(context);156 patterns.insert<StdAllocConversion>(context);157 mlir::ConversionTarget target(*context);158 target.addIllegalOp<memref::AllocOp>();159 target.addDynamicallyLegalOp<fir::ConvertOp>([](fir::ConvertOp op) {160 if (mlir::isa<mlir::MemRefType>(op.getRes().getType()))161 return false;162 return true;163 });164 target165 .addLegalDialect<FIROpsDialect, mlir::scf::SCFDialect,166 mlir::arith::ArithDialect, mlir::func::FuncDialect>();167 168 if (mlir::failed(mlir::applyPartialConversion(function, target,169 std::move(patterns)))) {170 mlir::emitError(mlir::UnknownLoc::get(context),171 "error in converting affine dialect\n");172 signalPassFailure();173 }174 }175};176 177} // namespace178 179std::unique_ptr<mlir::Pass> fir::createAffineDemotionPass() {180 return std::make_unique<AffineDialectDemotion>();181}182