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.
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).
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`.
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.
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.
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   RewritePatternSet patterns(ifOp.getContext());
192   AffineIfOp::getCanonicalizationPatterns(patterns, ifOp.getContext());
193   bool erased;
194   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
195   (void)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   (void)applyPatternsAndFoldGreedily(
224       hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
225       frozenPatterns);
226 
227   return success();
228 }
229 
230 // Return the min expr after replacing the given dim.
231 AffineExpr mlir::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
232                               AffineExpr max, bool positivePath) {
233   if (e == dim)
234     return positivePath ? min : max;
235   if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
236     AffineExpr lhs = bin.getLHS();
237     AffineExpr rhs = bin.getRHS();
238     if (bin.getKind() == mlir::AffineExprKind::Add)
239       return substWithMin(lhs, dim, min, max, positivePath) +
240              substWithMin(rhs, dim, min, max, positivePath);
241 
242     auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
243     auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
244     if (c1 && c1.getValue() < 0)
245       return getAffineBinaryOpExpr(
246           bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
247     if (c2 && c2.getValue() < 0)
248       return getAffineBinaryOpExpr(
249           bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
250     return getAffineBinaryOpExpr(
251         bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
252         substWithMin(rhs, dim, min, max, positivePath));
253   }
254   return e;
255 }
256