1 //===- Utils.cpp ---- Utilities for affine dialect transformation ---------===//
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 implements miscellaneous transformation utilities for the Affine
10 // dialect.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "mlir/Dialect/Affine/Utils.h"
15 #include "mlir/Dialect/Affine/IR/AffineOps.h"
16 #include "mlir/IR/BlockAndValueMapping.h"
17 #include "mlir/IR/BuiltinOps.h"
18 #include "mlir/IR/IntegerSet.h"
19 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
20 #include "llvm/Support/Debug.h"
21 #include "llvm/Support/raw_ostream.h"
22
23 using namespace mlir;
24
25 /// Promotes the `then` or the `else` block of `ifOp` (depending on whether
26 /// `elseBlock` is false or true) into `ifOp`'s containing block, and discards
27 /// the rest of the op.
promoteIfBlock(AffineIfOp ifOp,bool elseBlock)28 static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) {
29 if (elseBlock)
30 assert(ifOp.hasElse() && "else block expected");
31
32 Block *destBlock = ifOp->getBlock();
33 Block *srcBlock = elseBlock ? ifOp.getElseBlock() : ifOp.getThenBlock();
34 destBlock->getOperations().splice(
35 Block::iterator(ifOp), srcBlock->getOperations(), srcBlock->begin(),
36 std::prev(srcBlock->end()));
37 ifOp.erase();
38 }
39
40 /// Returns the outermost affine.for/parallel op that the `ifOp` is invariant
41 /// on. The `ifOp` could be hoisted and placed right before such an operation.
42 /// This method assumes that the ifOp has been canonicalized (to be correct and
43 /// effective).
getOutermostInvariantForOp(AffineIfOp ifOp)44 static Operation *getOutermostInvariantForOp(AffineIfOp ifOp) {
45 // Walk up the parents past all for op that this conditional is invariant on.
46 auto ifOperands = ifOp.getOperands();
47 auto *res = ifOp.getOperation();
48 while (!isa<FuncOp>(res->getParentOp())) {
49 auto *parentOp = res->getParentOp();
50 if (auto forOp = dyn_cast<AffineForOp>(parentOp)) {
51 if (llvm::is_contained(ifOperands, forOp.getInductionVar()))
52 break;
53 } else if (auto parallelOp = dyn_cast<AffineParallelOp>(parentOp)) {
54 for (auto iv : parallelOp.getIVs())
55 if (llvm::is_contained(ifOperands, iv))
56 break;
57 } else if (!isa<AffineIfOp>(parentOp)) {
58 // Won't walk up past anything other than affine.for/if ops.
59 break;
60 }
61 // You can always hoist up past any affine.if ops.
62 res = parentOp;
63 }
64 return res;
65 }
66
67 /// A helper for the mechanics of mlir::hoistAffineIfOp. Hoists `ifOp` just over
68 /// `hoistOverOp`. Returns the new hoisted op if any hoisting happened,
69 /// otherwise the same `ifOp`.
hoistAffineIfOp(AffineIfOp ifOp,Operation * hoistOverOp)70 static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
71 // No hoisting to do.
72 if (hoistOverOp == ifOp)
73 return ifOp;
74
75 // Create the hoisted 'if' first. Then, clone the op we are hoisting over for
76 // the else block. Then drop the else block of the original 'if' in the 'then'
77 // branch while promoting its then block, and analogously drop the 'then'
78 // block of the original 'if' from the 'else' branch while promoting its else
79 // block.
80 BlockAndValueMapping operandMap;
81 OpBuilder b(hoistOverOp);
82 auto hoistedIfOp = b.create<AffineIfOp>(ifOp.getLoc(), ifOp.getIntegerSet(),
83 ifOp.getOperands(),
84 /*elseBlock=*/true);
85
86 // Create a clone of hoistOverOp to use for the else branch of the hoisted
87 // conditional. The else block may get optimized away if empty.
88 Operation *hoistOverOpClone = nullptr;
89 // We use this unique name to identify/find `ifOp`'s clone in the else
90 // version.
91 Identifier idForIfOp = b.getIdentifier("__mlir_if_hoisting");
92 operandMap.clear();
93 b.setInsertionPointAfter(hoistOverOp);
94 // We'll set an attribute to identify this op in a clone of this sub-tree.
95 ifOp.setAttr(idForIfOp, b.getBoolAttr(true));
96 hoistOverOpClone = b.clone(*hoistOverOp, operandMap);
97
98 // Promote the 'then' block of the original affine.if in the then version.
99 promoteIfBlock(ifOp, /*elseBlock=*/false);
100
101 // Move the then version to the hoisted if op's 'then' block.
102 auto *thenBlock = hoistedIfOp.getThenBlock();
103 thenBlock->getOperations().splice(thenBlock->begin(),
104 hoistOverOp->getBlock()->getOperations(),
105 Block::iterator(hoistOverOp));
106
107 // Find the clone of the original affine.if op in the else version.
108 AffineIfOp ifCloneInElse;
109 hoistOverOpClone->walk([&](AffineIfOp ifClone) {
110 if (!ifClone.getAttr(idForIfOp))
111 return WalkResult::advance();
112 ifCloneInElse = ifClone;
113 return WalkResult::interrupt();
114 });
115 assert(ifCloneInElse && "if op clone should exist");
116 // For the else block, promote the else block of the original 'if' if it had
117 // one; otherwise, the op itself is to be erased.
118 if (!ifCloneInElse.hasElse())
119 ifCloneInElse.erase();
120 else
121 promoteIfBlock(ifCloneInElse, /*elseBlock=*/true);
122
123 // Move the else version into the else block of the hoisted if op.
124 auto *elseBlock = hoistedIfOp.getElseBlock();
125 elseBlock->getOperations().splice(
126 elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(),
127 Block::iterator(hoistOverOpClone));
128
129 return hoistedIfOp;
130 }
131
132 /// Replace affine.for with a 1-d affine.parallel and clone the former's body
133 /// into the latter while remapping values.
affineParallelize(AffineForOp forOp)134 void mlir::affineParallelize(AffineForOp forOp) {
135 Location loc = forOp.getLoc();
136 OpBuilder outsideBuilder(forOp);
137
138 // If a loop has a 'max' in the lower bound, emit it outside the parallel loop
139 // as it does not have implicit 'max' behavior.
140 AffineMap lowerBoundMap = forOp.getLowerBoundMap();
141 ValueRange lowerBoundOperands = forOp.getLowerBoundOperands();
142 AffineMap upperBoundMap = forOp.getUpperBoundMap();
143 ValueRange upperBoundOperands = forOp.getUpperBoundOperands();
144
145 bool needsMax = lowerBoundMap.getNumResults() > 1;
146 bool needsMin = upperBoundMap.getNumResults() > 1;
147 AffineMap identityMap;
148 if (needsMax || needsMin) {
149 if (forOp->getParentOp() &&
150 !forOp->getParentOp()->hasTrait<OpTrait::AffineScope>())
151 return;
152
153 identityMap = AffineMap::getMultiDimIdentityMap(1, loc->getContext());
154 }
155 if (needsMax) {
156 auto maxOp = outsideBuilder.create<AffineMaxOp>(loc, lowerBoundMap,
157 lowerBoundOperands);
158 lowerBoundMap = identityMap;
159 lowerBoundOperands = maxOp->getResults();
160 }
161
162 // Same for the upper bound.
163 if (needsMin) {
164 auto minOp = outsideBuilder.create<AffineMinOp>(loc, upperBoundMap,
165 upperBoundOperands);
166 upperBoundMap = identityMap;
167 upperBoundOperands = minOp->getResults();
168 }
169
170 // Creating empty 1-D affine.parallel op.
171 AffineParallelOp newPloop = outsideBuilder.create<AffineParallelOp>(
172 loc, llvm::None, llvm::None, lowerBoundMap, lowerBoundOperands,
173 upperBoundMap, upperBoundOperands);
174 // Steal the body of the old affine for op and erase it.
175 newPloop.region().takeBody(forOp.region());
176 forOp.erase();
177 }
178
179 // Returns success if any hoisting happened.
hoistAffineIfOp(AffineIfOp ifOp,bool * folded)180 LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
181 // Bail out early if the ifOp returns a result. TODO: Consider how to
182 // properly support this case.
183 if (ifOp.getNumResults() != 0)
184 return failure();
185
186 // Apply canonicalization patterns and folding - this is necessary for the
187 // hoisting check to be correct (operands should be composed), and to be more
188 // effective (no unused operands). Since the pattern rewriter's folding is
189 // entangled with application of patterns, we may fold/end up erasing the op,
190 // in which case we return with `folded` being set.
191 OwningRewritePatternList patterns;
192 AffineIfOp::getCanonicalizationPatterns(patterns, ifOp.getContext());
193 bool erased;
194 FrozenRewritePatternList frozenPatterns(std::move(patterns));
195 applyOpPatternsAndFold(ifOp, frozenPatterns, &erased);
196 if (erased) {
197 if (folded)
198 *folded = true;
199 return failure();
200 }
201 if (folded)
202 *folded = false;
203
204 // The folding above should have ensured this, but the affine.if's
205 // canonicalization is missing composition of affine.applys into it.
206 assert(llvm::all_of(ifOp.getOperands(),
207 [](Value v) {
208 return isTopLevelValue(v) || isForInductionVar(v);
209 }) &&
210 "operands not composed");
211
212 // We are going hoist as high as possible.
213 // TODO: this could be customized in the future.
214 auto *hoistOverOp = getOutermostInvariantForOp(ifOp);
215
216 AffineIfOp hoistedIfOp = ::hoistAffineIfOp(ifOp, hoistOverOp);
217 // Nothing to hoist over.
218 if (hoistedIfOp == ifOp)
219 return failure();
220
221 // Canonicalize to remove dead else blocks (happens whenever an 'if' moves up
222 // a sequence of affine.fors that are all perfectly nested).
223 applyPatternsAndFoldGreedily(
224 hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
225 frozenPatterns);
226
227 return success();
228 }
229