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/Analysis/AffineAnalysis.h"
16 #include "mlir/Analysis/Utils.h"
17 #include "mlir/Dialect/Affine/IR/AffineOps.h"
18 #include "mlir/IR/BlockAndValueMapping.h"
19 #include "mlir/IR/BuiltinOps.h"
20 #include "mlir/IR/IntegerSet.h"
21 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
22 #include "llvm/Support/Debug.h"
23 #include "llvm/Support/raw_ostream.h"
24 
25 using namespace mlir;
26 
27 /// Promotes the `then` or the `else` block of `ifOp` (depending on whether
28 /// `elseBlock` is false or true) into `ifOp`'s containing block, and discards
29 /// the rest of the op.
30 static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) {
31   if (elseBlock)
32     assert(ifOp.hasElse() && "else block expected");
33 
34   Block *destBlock = ifOp->getBlock();
35   Block *srcBlock = elseBlock ? ifOp.getElseBlock() : ifOp.getThenBlock();
36   destBlock->getOperations().splice(
37       Block::iterator(ifOp), srcBlock->getOperations(), srcBlock->begin(),
38       std::prev(srcBlock->end()));
39   ifOp.erase();
40 }
41 
42 /// Returns the outermost affine.for/parallel op that the `ifOp` is invariant
43 /// on. The `ifOp` could be hoisted and placed right before such an operation.
44 /// This method assumes that the ifOp has been canonicalized (to be correct and
45 /// effective).
46 static Operation *getOutermostInvariantForOp(AffineIfOp ifOp) {
47   // Walk up the parents past all for op that this conditional is invariant on.
48   auto ifOperands = ifOp.getOperands();
49   auto *res = ifOp.getOperation();
50   while (!isa<FuncOp>(res->getParentOp())) {
51     auto *parentOp = res->getParentOp();
52     if (auto forOp = dyn_cast<AffineForOp>(parentOp)) {
53       if (llvm::is_contained(ifOperands, forOp.getInductionVar()))
54         break;
55     } else if (auto parallelOp = dyn_cast<AffineParallelOp>(parentOp)) {
56       for (auto iv : parallelOp.getIVs())
57         if (llvm::is_contained(ifOperands, iv))
58           break;
59     } else if (!isa<AffineIfOp>(parentOp)) {
60       // Won't walk up past anything other than affine.for/if ops.
61       break;
62     }
63     // You can always hoist up past any affine.if ops.
64     res = parentOp;
65   }
66   return res;
67 }
68 
69 /// A helper for the mechanics of mlir::hoistAffineIfOp. Hoists `ifOp` just over
70 /// `hoistOverOp`. Returns the new hoisted op if any hoisting happened,
71 /// otherwise the same `ifOp`.
72 static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
73   // No hoisting to do.
74   if (hoistOverOp == ifOp)
75     return ifOp;
76 
77   // Create the hoisted 'if' first. Then, clone the op we are hoisting over for
78   // the else block. Then drop the else block of the original 'if' in the 'then'
79   // branch while promoting its then block, and analogously drop the 'then'
80   // block of the original 'if' from the 'else' branch while promoting its else
81   // block.
82   BlockAndValueMapping operandMap;
83   OpBuilder b(hoistOverOp);
84   auto hoistedIfOp = b.create<AffineIfOp>(ifOp.getLoc(), ifOp.getIntegerSet(),
85                                           ifOp.getOperands(),
86                                           /*elseBlock=*/true);
87 
88   // Create a clone of hoistOverOp to use for the else branch of the hoisted
89   // conditional. The else block may get optimized away if empty.
90   Operation *hoistOverOpClone = nullptr;
91   // We use this unique name to identify/find  `ifOp`'s clone in the else
92   // version.
93   Identifier idForIfOp = b.getIdentifier("__mlir_if_hoisting");
94   operandMap.clear();
95   b.setInsertionPointAfter(hoistOverOp);
96   // We'll set an attribute to identify this op in a clone of this sub-tree.
97   ifOp->setAttr(idForIfOp, b.getBoolAttr(true));
98   hoistOverOpClone = b.clone(*hoistOverOp, operandMap);
99 
100   // Promote the 'then' block of the original affine.if in the then version.
101   promoteIfBlock(ifOp, /*elseBlock=*/false);
102 
103   // Move the then version to the hoisted if op's 'then' block.
104   auto *thenBlock = hoistedIfOp.getThenBlock();
105   thenBlock->getOperations().splice(thenBlock->begin(),
106                                     hoistOverOp->getBlock()->getOperations(),
107                                     Block::iterator(hoistOverOp));
108 
109   // Find the clone of the original affine.if op in the else version.
110   AffineIfOp ifCloneInElse;
111   hoistOverOpClone->walk([&](AffineIfOp ifClone) {
112     if (!ifClone->getAttr(idForIfOp))
113       return WalkResult::advance();
114     ifCloneInElse = ifClone;
115     return WalkResult::interrupt();
116   });
117   assert(ifCloneInElse && "if op clone should exist");
118   // For the else block, promote the else block of the original 'if' if it had
119   // one; otherwise, the op itself is to be erased.
120   if (!ifCloneInElse.hasElse())
121     ifCloneInElse.erase();
122   else
123     promoteIfBlock(ifCloneInElse, /*elseBlock=*/true);
124 
125   // Move the else version into the else block of the hoisted if op.
126   auto *elseBlock = hoistedIfOp.getElseBlock();
127   elseBlock->getOperations().splice(
128       elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(),
129       Block::iterator(hoistOverOpClone));
130 
131   return hoistedIfOp;
132 }
133 
134 /// Replace affine.for with a 1-d affine.parallel and clone the former's body
135 /// into the latter while remapping values. Parallelizes the specified
136 /// reductions. Parallelization will fail in presence of loop iteration
137 /// arguments that are not listed in `parallelReductions`.
138 LogicalResult
139 mlir::affineParallelize(AffineForOp forOp,
140                         ArrayRef<LoopReduction> parallelReductions) {
141   // Fail early if there are iter arguments that are not reductions.
142   unsigned numReductions = parallelReductions.size();
143   if (numReductions != forOp.getNumIterOperands())
144     return failure();
145 
146   Location loc = forOp.getLoc();
147   OpBuilder outsideBuilder(forOp);
148   AffineMap lowerBoundMap = forOp.getLowerBoundMap();
149   ValueRange lowerBoundOperands = forOp.getLowerBoundOperands();
150   AffineMap upperBoundMap = forOp.getUpperBoundMap();
151   ValueRange upperBoundOperands = forOp.getUpperBoundOperands();
152 
153   // Creating empty 1-D affine.parallel op.
154   auto reducedValues = llvm::to_vector<4>(llvm::map_range(
155       parallelReductions, [](const LoopReduction &red) { return red.value; }));
156   auto reductionKinds = llvm::to_vector<4>(llvm::map_range(
157       parallelReductions, [](const LoopReduction &red) { return red.kind; }));
158   AffineParallelOp newPloop = outsideBuilder.create<AffineParallelOp>(
159       loc, ValueRange(reducedValues).getTypes(), reductionKinds,
160       llvm::makeArrayRef(lowerBoundMap), lowerBoundOperands,
161       llvm::makeArrayRef(upperBoundMap), upperBoundOperands,
162       llvm::makeArrayRef(forOp.getStep()));
163   // Steal the body of the old affine for op.
164   newPloop.region().takeBody(forOp.region());
165   Operation *yieldOp = &newPloop.getBody()->back();
166 
167   // Handle the initial values of reductions because the parallel loop always
168   // starts from the neutral value.
169   SmallVector<Value> newResults;
170   newResults.reserve(numReductions);
171   for (unsigned i = 0; i < numReductions; ++i) {
172     Value init = forOp.getIterOperands()[i];
173     // This works because we are only handling single-op reductions at the
174     // moment. A switch on reduction kind or a mechanism to collect operations
175     // participating in the reduction will be necessary for multi-op reductions.
176     Operation *reductionOp = yieldOp->getOperand(i).getDefiningOp();
177     assert(reductionOp && "yielded value is expected to be produced by an op");
178     outsideBuilder.getInsertionBlock()->getOperations().splice(
179         outsideBuilder.getInsertionPoint(), newPloop.getBody()->getOperations(),
180         reductionOp);
181     reductionOp->setOperands({init, newPloop->getResult(i)});
182     forOp->getResult(i).replaceAllUsesWith(reductionOp->getResult(0));
183   }
184 
185   // Update the loop terminator to yield reduced values bypassing the reduction
186   // operation itself (now moved outside of the loop) and erase the block
187   // arguments that correspond to reductions. Note that the loop always has one
188   // "main" induction variable whenc coming from a non-parallel for.
189   unsigned numIVs = 1;
190   yieldOp->setOperands(reducedValues);
191   newPloop.getBody()->eraseArguments(
192       llvm::to_vector<4>(llvm::seq<unsigned>(numIVs, numReductions + numIVs)));
193 
194   forOp.erase();
195   return success();
196 }
197 
198 // Returns success if any hoisting happened.
199 LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
200   // Bail out early if the ifOp returns a result.  TODO: Consider how to
201   // properly support this case.
202   if (ifOp.getNumResults() != 0)
203     return failure();
204 
205   // Apply canonicalization patterns and folding - this is necessary for the
206   // hoisting check to be correct (operands should be composed), and to be more
207   // effective (no unused operands). Since the pattern rewriter's folding is
208   // entangled with application of patterns, we may fold/end up erasing the op,
209   // in which case we return with `folded` being set.
210   RewritePatternSet patterns(ifOp.getContext());
211   AffineIfOp::getCanonicalizationPatterns(patterns, ifOp.getContext());
212   bool erased;
213   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
214   (void)applyOpPatternsAndFold(ifOp, frozenPatterns, &erased);
215   if (erased) {
216     if (folded)
217       *folded = true;
218     return failure();
219   }
220   if (folded)
221     *folded = false;
222 
223   // The folding above should have ensured this, but the affine.if's
224   // canonicalization is missing composition of affine.applys into it.
225   assert(llvm::all_of(ifOp.getOperands(),
226                       [](Value v) {
227                         return isTopLevelValue(v) || isForInductionVar(v);
228                       }) &&
229          "operands not composed");
230 
231   // We are going hoist as high as possible.
232   // TODO: this could be customized in the future.
233   auto *hoistOverOp = getOutermostInvariantForOp(ifOp);
234 
235   AffineIfOp hoistedIfOp = ::hoistAffineIfOp(ifOp, hoistOverOp);
236   // Nothing to hoist over.
237   if (hoistedIfOp == ifOp)
238     return failure();
239 
240   // Canonicalize to remove dead else blocks (happens whenever an 'if' moves up
241   // a sequence of affine.fors that are all perfectly nested).
242   (void)applyPatternsAndFoldGreedily(
243       hoistedIfOp->getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
244       frozenPatterns);
245 
246   return success();
247 }
248 
249 // Return the min expr after replacing the given dim.
250 AffineExpr mlir::substWithMin(AffineExpr e, AffineExpr dim, AffineExpr min,
251                               AffineExpr max, bool positivePath) {
252   if (e == dim)
253     return positivePath ? min : max;
254   if (auto bin = e.dyn_cast<AffineBinaryOpExpr>()) {
255     AffineExpr lhs = bin.getLHS();
256     AffineExpr rhs = bin.getRHS();
257     if (bin.getKind() == mlir::AffineExprKind::Add)
258       return substWithMin(lhs, dim, min, max, positivePath) +
259              substWithMin(rhs, dim, min, max, positivePath);
260 
261     auto c1 = bin.getLHS().dyn_cast<AffineConstantExpr>();
262     auto c2 = bin.getRHS().dyn_cast<AffineConstantExpr>();
263     if (c1 && c1.getValue() < 0)
264       return getAffineBinaryOpExpr(
265           bin.getKind(), c1, substWithMin(rhs, dim, min, max, !positivePath));
266     if (c2 && c2.getValue() < 0)
267       return getAffineBinaryOpExpr(
268           bin.getKind(), substWithMin(lhs, dim, min, max, !positivePath), c2);
269     return getAffineBinaryOpExpr(
270         bin.getKind(), substWithMin(lhs, dim, min, max, positivePath),
271         substWithMin(rhs, dim, min, max, positivePath));
272   }
273   return e;
274 }
275