1 | //===- LoopCanonicalization.cpp - Cross-dialect canonicalization patterns -===// |
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 | // This file contains cross-dialect canonicalization patterns that cannot be |
10 | // actual canonicalization patterns due to undesired additional dependencies. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #include "mlir/Dialect/SCF/Transforms/Passes.h" |
15 | |
16 | #include "mlir/Dialect/Affine/IR/AffineOps.h" |
17 | #include "mlir/Dialect/MemRef/IR/MemRef.h" |
18 | #include "mlir/Dialect/SCF/IR/SCF.h" |
19 | #include "mlir/Dialect/SCF/Transforms/Patterns.h" |
20 | #include "mlir/Dialect/SCF/Utils/AffineCanonicalizationUtils.h" |
21 | #include "mlir/Dialect/Tensor/IR/Tensor.h" |
22 | #include "mlir/IR/PatternMatch.h" |
23 | #include "mlir/Transforms/GreedyPatternRewriteDriver.h" |
24 | #include "llvm/ADT/TypeSwitch.h" |
25 | |
26 | namespace mlir { |
27 | #define GEN_PASS_DEF_SCFFORLOOPCANONICALIZATION |
28 | #include "mlir/Dialect/SCF/Transforms/Passes.h.inc" |
29 | } // namespace mlir |
30 | |
31 | using namespace mlir; |
32 | using namespace mlir::scf; |
33 | |
34 | /// A simple, conservative analysis to determine if the loop is shape |
35 | /// conserving. I.e., the type of the arg-th yielded value is the same as the |
36 | /// type of the corresponding basic block argument of the loop. |
37 | /// Note: This function handles only simple cases. Expand as needed. |
38 | static bool isShapePreserving(ForOp forOp, int64_t arg) { |
39 | assert(arg < static_cast<int64_t>(forOp.getNumResults()) && |
40 | "arg is out of bounds" ); |
41 | Value value = forOp.getYieldedValues()[arg]; |
42 | while (value) { |
43 | if (value == forOp.getRegionIterArgs()[arg]) |
44 | return true; |
45 | OpResult opResult = dyn_cast<OpResult>(value); |
46 | if (!opResult) |
47 | return false; |
48 | |
49 | using tensor::InsertSliceOp; |
50 | value = llvm::TypeSwitch<Operation *, Value>(opResult.getOwner()) |
51 | .template Case<InsertSliceOp>( |
52 | [&](InsertSliceOp op) { return op.getDest(); }) |
53 | .template Case<ForOp>([&](ForOp forOp) { |
54 | return isShapePreserving(forOp, opResult.getResultNumber()) |
55 | ? forOp.getInitArgs()[opResult.getResultNumber()] |
56 | : Value(); |
57 | }) |
58 | .Default([&](auto op) { return Value(); }); |
59 | } |
60 | return false; |
61 | } |
62 | |
63 | namespace { |
64 | /// Fold dim ops of iter_args to dim ops of their respective init args. E.g.: |
65 | /// |
66 | /// ``` |
67 | /// %0 = ... : tensor<?x?xf32> |
68 | /// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) { |
69 | /// %1 = tensor.dim %arg0, %c0 : tensor<?x?xf32> |
70 | /// ... |
71 | /// } |
72 | /// ``` |
73 | /// |
74 | /// is folded to: |
75 | /// |
76 | /// ``` |
77 | /// %0 = ... : tensor<?x?xf32> |
78 | /// scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) { |
79 | /// %1 = tensor.dim %0, %c0 : tensor<?x?xf32> |
80 | /// ... |
81 | /// } |
82 | /// ``` |
83 | /// |
84 | /// Note: Dim ops are folded only if it can be proven that the runtime type of |
85 | /// the iter arg does not change with loop iterations. |
86 | template <typename OpTy> |
87 | struct DimOfIterArgFolder : public OpRewritePattern<OpTy> { |
88 | using OpRewritePattern<OpTy>::OpRewritePattern; |
89 | |
90 | LogicalResult matchAndRewrite(OpTy dimOp, |
91 | PatternRewriter &rewriter) const override { |
92 | auto blockArg = dyn_cast<BlockArgument>(dimOp.getSource()); |
93 | if (!blockArg) |
94 | return failure(); |
95 | auto forOp = dyn_cast<ForOp>(blockArg.getParentBlock()->getParentOp()); |
96 | if (!forOp) |
97 | return failure(); |
98 | if (!isShapePreserving(forOp, blockArg.getArgNumber() - 1)) |
99 | return failure(); |
100 | |
101 | Value initArg = forOp.getTiedLoopInit(blockArg)->get(); |
102 | rewriter.modifyOpInPlace( |
103 | dimOp, [&]() { dimOp.getSourceMutable().assign(initArg); }); |
104 | |
105 | return success(); |
106 | }; |
107 | }; |
108 | |
109 | /// Fold dim ops of loop results to dim ops of their respective init args. E.g.: |
110 | /// |
111 | /// ``` |
112 | /// %0 = ... : tensor<?x?xf32> |
113 | /// %r = scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) { |
114 | /// ... |
115 | /// } |
116 | /// %1 = tensor.dim %r, %c0 : tensor<?x?xf32> |
117 | /// ``` |
118 | /// |
119 | /// is folded to: |
120 | /// |
121 | /// ``` |
122 | /// %0 = ... : tensor<?x?xf32> |
123 | /// %r = scf.for ... iter_args(%arg0 = %0) -> (tensor<?x?xf32>) { |
124 | /// ... |
125 | /// } |
126 | /// %1 = tensor.dim %0, %c0 : tensor<?x?xf32> |
127 | /// ``` |
128 | /// |
129 | /// Note: Dim ops are folded only if it can be proven that the runtime type of |
130 | /// the iter arg does not change with loop iterations. |
131 | template <typename OpTy> |
132 | struct DimOfLoopResultFolder : public OpRewritePattern<OpTy> { |
133 | using OpRewritePattern<OpTy>::OpRewritePattern; |
134 | |
135 | LogicalResult matchAndRewrite(OpTy dimOp, |
136 | PatternRewriter &rewriter) const override { |
137 | auto forOp = dimOp.getSource().template getDefiningOp<scf::ForOp>(); |
138 | if (!forOp) |
139 | return failure(); |
140 | auto opResult = cast<OpResult>(dimOp.getSource()); |
141 | unsigned resultNumber = opResult.getResultNumber(); |
142 | if (!isShapePreserving(forOp, resultNumber)) |
143 | return failure(); |
144 | rewriter.modifyOpInPlace(dimOp, [&]() { |
145 | dimOp.getSourceMutable().assign(forOp.getInitArgs()[resultNumber]); |
146 | }); |
147 | return success(); |
148 | } |
149 | }; |
150 | |
151 | /// Canonicalize AffineMinOp/AffineMaxOp operations in the context of scf.for |
152 | /// and scf.parallel loops with a known range. |
153 | template <typename OpTy> |
154 | struct AffineOpSCFCanonicalizationPattern : public OpRewritePattern<OpTy> { |
155 | using OpRewritePattern<OpTy>::OpRewritePattern; |
156 | |
157 | LogicalResult matchAndRewrite(OpTy op, |
158 | PatternRewriter &rewriter) const override { |
159 | return scf::canonicalizeMinMaxOpInLoop(rewriter, op, loopMatcher: scf::matchForLikeLoop); |
160 | } |
161 | }; |
162 | |
163 | struct SCFForLoopCanonicalization |
164 | : public impl::SCFForLoopCanonicalizationBase<SCFForLoopCanonicalization> { |
165 | void runOnOperation() override { |
166 | auto *parentOp = getOperation(); |
167 | MLIRContext *ctx = parentOp->getContext(); |
168 | RewritePatternSet patterns(ctx); |
169 | scf::populateSCFForLoopCanonicalizationPatterns(patterns); |
170 | if (failed(applyPatternsAndFoldGreedily(parentOp, std::move(patterns)))) |
171 | signalPassFailure(); |
172 | } |
173 | }; |
174 | } // namespace |
175 | |
176 | void mlir::scf::populateSCFForLoopCanonicalizationPatterns( |
177 | RewritePatternSet &patterns) { |
178 | MLIRContext *ctx = patterns.getContext(); |
179 | patterns |
180 | .add<AffineOpSCFCanonicalizationPattern<affine::AffineMinOp>, |
181 | AffineOpSCFCanonicalizationPattern<affine::AffineMaxOp>, |
182 | DimOfIterArgFolder<tensor::DimOp>, DimOfIterArgFolder<memref::DimOp>, |
183 | DimOfLoopResultFolder<tensor::DimOp>, |
184 | DimOfLoopResultFolder<memref::DimOp>>(ctx); |
185 | } |
186 | |
187 | std::unique_ptr<Pass> mlir::createSCFForLoopCanonicalizationPass() { |
188 | return std::make_unique<SCFForLoopCanonicalization>(); |
189 | } |
190 | |