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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "PassDetail.h"
10 #include "flang/Optimizer/Dialect/FIRDialect.h"
11 #include "flang/Optimizer/Dialect/FIROps.h"
12 #include "flang/Optimizer/Dialect/FIRType.h"
13 #include "flang/Optimizer/Transforms/Passes.h"
14 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/Dialect/MemRef/IR/MemRef.h"
17 #include "mlir/Dialect/SCF/SCF.h"
18 #include "mlir/Dialect/StandardOps/IR/Ops.h"
19 #include "mlir/IR/BuiltinAttributes.h"
20 #include "mlir/IR/IntegerSet.h"
21 #include "mlir/IR/Visitors.h"
22 #include "mlir/Pass/Pass.h"
23 #include "mlir/Transforms/DialectConversion.h"
24 #include "llvm/ADT/DenseMap.h"
25 #include "llvm/ADT/Optional.h"
26 #include "llvm/Support/CommandLine.h"
27 #include "llvm/Support/Debug.h"
28 
29 #define DEBUG_TYPE "flang-affine-demotion"
30 
31 using namespace fir;
32 
33 namespace {
34 
35 class AffineLoadConversion : public OpRewritePattern<mlir::AffineLoadOp> {
36 public:
37   using OpRewritePattern<mlir::AffineLoadOp>::OpRewritePattern;
38 
39   LogicalResult matchAndRewrite(mlir::AffineLoadOp op,
40                                 PatternRewriter &rewriter) const override {
41     SmallVector<Value> indices(op.getMapOperands());
42     auto maybeExpandedMap =
43         expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
44     if (!maybeExpandedMap)
45       return failure();
46 
47     auto coorOp = rewriter.create<fir::CoordinateOp>(
48         op.getLoc(), fir::ReferenceType::get(op.getResult().getType()),
49         op.getMemRef(), *maybeExpandedMap);
50 
51     rewriter.replaceOpWithNewOp<fir::LoadOp>(op, coorOp.getResult());
52     return success();
53   }
54 };
55 
56 class AffineStoreConversion : public OpRewritePattern<mlir::AffineStoreOp> {
57 public:
58   using OpRewritePattern<mlir::AffineStoreOp>::OpRewritePattern;
59 
60   LogicalResult matchAndRewrite(mlir::AffineStoreOp op,
61                                 PatternRewriter &rewriter) const override {
62     SmallVector<Value> indices(op.getMapOperands());
63     auto maybeExpandedMap =
64         expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
65     if (!maybeExpandedMap)
66       return failure();
67 
68     auto coorOp = rewriter.create<fir::CoordinateOp>(
69         op.getLoc(), fir::ReferenceType::get(op.getValueToStore().getType()),
70         op.getMemRef(), *maybeExpandedMap);
71     rewriter.replaceOpWithNewOp<fir::StoreOp>(op, op.getValueToStore(),
72                                               coorOp.getResult());
73     return success();
74   }
75 };
76 
77 class ConvertConversion : public mlir::OpRewritePattern<fir::ConvertOp> {
78 public:
79   using OpRewritePattern::OpRewritePattern;
80   mlir::LogicalResult
81   matchAndRewrite(fir::ConvertOp op,
82                   mlir::PatternRewriter &rewriter) const override {
83     if (op.res().getType().isa<mlir::MemRefType>()) {
84       // due to index calculation moving to affine maps we still need to
85       // add converts for sequence types this has a side effect of losing
86       // some information about arrays with known dimensions by creating:
87       // fir.convert %arg0 : (!fir.ref<!fir.array<5xi32>>) ->
88       // !fir.ref<!fir.array<?xi32>>
89       if (auto refTy = op.value().getType().dyn_cast<fir::ReferenceType>())
90         if (auto arrTy = refTy.getEleTy().dyn_cast<fir::SequenceType>()) {
91           fir::SequenceType::Shape flatShape = {
92               fir::SequenceType::getUnknownExtent()};
93           auto flatArrTy = fir::SequenceType::get(flatShape, arrTy.getEleTy());
94           auto flatTy = fir::ReferenceType::get(flatArrTy);
95           rewriter.replaceOpWithNewOp<fir::ConvertOp>(op, flatTy, op.value());
96           return success();
97         }
98       rewriter.startRootUpdate(op->getParentOp());
99       op.getResult().replaceAllUsesWith(op.value());
100       rewriter.finalizeRootUpdate(op->getParentOp());
101       rewriter.eraseOp(op);
102     }
103     return success();
104   }
105 };
106 
107 mlir::Type convertMemRef(mlir::MemRefType type) {
108   return fir::SequenceType::get(
109       SmallVector<int64_t>(type.getShape().begin(), type.getShape().end()),
110       type.getElementType());
111 }
112 
113 class StdAllocConversion : public mlir::OpRewritePattern<memref::AllocOp> {
114 public:
115   using OpRewritePattern::OpRewritePattern;
116   mlir::LogicalResult
117   matchAndRewrite(memref::AllocOp op,
118                   mlir::PatternRewriter &rewriter) const override {
119     rewriter.replaceOpWithNewOp<fir::AllocaOp>(op, convertMemRef(op.getType()),
120                                                op.memref());
121     return success();
122   }
123 };
124 
125 class AffineDialectDemotion
126     : public AffineDialectDemotionBase<AffineDialectDemotion> {
127 public:
128   void runOnFunction() override {
129     auto *context = &getContext();
130     auto function = getFunction();
131     LLVM_DEBUG(llvm::dbgs() << "AffineDemotion: running on function:\n";
132                function.print(llvm::dbgs()););
133 
134     mlir::OwningRewritePatternList patterns(context);
135     patterns.insert<ConvertConversion>(context);
136     patterns.insert<AffineLoadConversion>(context);
137     patterns.insert<AffineStoreConversion>(context);
138     patterns.insert<StdAllocConversion>(context);
139     mlir::ConversionTarget target(*context);
140     target.addIllegalOp<memref::AllocOp>();
141     target.addDynamicallyLegalOp<fir::ConvertOp>([](fir::ConvertOp op) {
142       if (op.res().getType().isa<mlir::MemRefType>())
143         return false;
144       return true;
145     });
146     target.addLegalDialect<FIROpsDialect, mlir::scf::SCFDialect,
147                            mlir::StandardOpsDialect>();
148 
149     if (mlir::failed(mlir::applyPartialConversion(function, target,
150                                                   std::move(patterns)))) {
151       mlir::emitError(mlir::UnknownLoc::get(context),
152                       "error in converting affine dialect\n");
153       signalPassFailure();
154     }
155   }
156 };
157 
158 } // namespace
159 
160 std::unique_ptr<mlir::Pass> fir::createAffineDemotionPass() {
161   return std::make_unique<AffineDialectDemotion>();
162 }
163