1 //===- Hoisting.cpp - Linalg hoisting transformations ---------------------===//
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 functions concerned with hoisting invariant operations
10 // in the context of Linalg transformations.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
15 #include "mlir/Analysis/SliceAnalysis.h"
16 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
17 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
18 #include "mlir/Dialect/Func/IR/FuncOps.h"
19 #include "mlir/Dialect/Linalg/IR/Linalg.h"
20 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
21 #include "mlir/Dialect/SCF/SCF.h"
22 #include "mlir/Dialect/SCF/Utils/Utils.h"
23 #include "mlir/Dialect/Tensor/IR/Tensor.h"
24 #include "mlir/Dialect/Vector/IR/VectorOps.h"
25 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
26 #include "mlir/IR/BuiltinOps.h"
27 #include "mlir/IR/Dominance.h"
28 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
29 #include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
30 #include "llvm/ADT/StringRef.h"
31 #include "llvm/Support/Debug.h"
32 
33 using llvm::dbgs;
34 
35 #define DEBUG_TYPE "linalg-hoisting"
36 
37 #define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")
38 
39 using namespace mlir;
40 using namespace mlir::linalg;
41 
42 namespace {
43 /// Represents a unit of hoistable TransferWriteOp. This may comprise other
44 /// instructions that need to be hoisted too.
45 struct HoistableWrite {
46   vector::TransferWriteOp transferWriteOp;
47   tensor::InsertSliceOp insertSliceOp;
48 };
49 /// Represents a unit of hoistable TransferReadOp. This may comprise other
50 /// instructions that need to be hoisted too.
51 struct HoistableRead {
52   vector::TransferReadOp transferReadOp;
53   tensor::ExtractSliceOp extractSliceOp;
54 };
55 } // namespace
56 
57 /// Return true if op1 and op2 are the same constant or the same SSA value.
58 static bool isEqualOffsetSizeOrStride(OpFoldResult op1, OpFoldResult op2) {
59   auto getConstantIntValue = [](OpFoldResult ofr) -> llvm::Optional<int64_t> {
60     Attribute attr = ofr.dyn_cast<Attribute>();
61     // Note: isa+cast-like pattern allows writing the condition below as 1 line.
62     if (!attr && ofr.get<Value>().getDefiningOp<arith::ConstantOp>())
63       attr = ofr.get<Value>().getDefiningOp<arith::ConstantOp>().getValue();
64     if (auto intAttr = attr.dyn_cast_or_null<IntegerAttr>())
65       return intAttr.getValue().getSExtValue();
66     return llvm::None;
67   };
68   auto cst1 = getConstantIntValue(op1), cst2 = getConstantIntValue(op2);
69   if (cst1 && cst2 && *cst1 == *cst2)
70     return true;
71   auto v1 = op1.dyn_cast<Value>(), v2 = op2.dyn_cast<Value>();
72   return v1 && v2 && v1 == v2;
73 }
74 
75 /// Return true is all offsets, sizes and strides are equal.
76 static bool sameOffsetsSizesAndStrides(tensor::ExtractSliceOp s,
77                                        tensor::InsertSliceOp si) {
78   if (s.static_offsets().size() != si.static_offsets().size())
79     return false;
80   if (s.static_sizes().size() != si.static_sizes().size())
81     return false;
82   if (s.static_strides().size() != si.static_strides().size())
83     return false;
84   for (auto it : llvm::zip(s.getMixedOffsets(), si.getMixedOffsets()))
85     if (!isEqualOffsetSizeOrStride(std::get<0>(it), std::get<1>(it)))
86       return false;
87   for (auto it : llvm::zip(s.getMixedSizes(), si.getMixedSizes()))
88     if (!isEqualOffsetSizeOrStride(std::get<0>(it), std::get<1>(it)))
89       return false;
90   for (auto it : llvm::zip(s.getMixedStrides(), si.getMixedStrides()))
91     if (!isEqualOffsetSizeOrStride(std::get<0>(it), std::get<1>(it)))
92       return false;
93   return true;
94 }
95 
96 /// Look for a HoistableRead, in the given tensor uses, accessing the same
97 /// offset as the HoistableWrite.
98 static HoistableRead findMatchingTransferRead(HoistableWrite write,
99                                               Value srcTensor) {
100   assert(write.transferWriteOp &&
101          "expected hoistable write to have a .transfer_write");
102 
103   LLVM_DEBUG(DBGS() << "findMatchingTransferRead for: "
104                     << *write.transferWriteOp.getOperation() << "\n");
105   if (write.insertSliceOp)
106     LLVM_DEBUG(DBGS() << "findMatchingTransferRead inserSliceOp: "
107                       << *write.insertSliceOp.getOperation() << "\n");
108 
109   for (Operation *user : srcTensor.getUsers()) {
110     LLVM_DEBUG(DBGS() << "findMatchingTransferRead inspect user: " << *user
111                       << "\n");
112 
113     // If HoistableWrite involves a InsertSliceOp, we need to find a
114     // matching ExtractSliceOp.
115     tensor::ExtractSliceOp sliceOp;
116     Operation *maybeTransferReadUser = user;
117     if (write.insertSliceOp) {
118       sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
119       if (!sliceOp || sliceOp.getResult().getType() !=
120                           write.insertSliceOp.source().getType())
121         continue;
122 
123       LLVM_DEBUG(DBGS() << "check whether sameOffsetsSizesAndStrides: "
124                         << *sliceOp << " vs " << *write.insertSliceOp << "\n");
125       if (!sameOffsetsSizesAndStrides(sliceOp, write.insertSliceOp))
126         continue;
127 
128       LLVM_DEBUG(DBGS() << "sameOffsetsSizesAndStrides: SUCCESS\n");
129       // If we got here, sliceOp is hoistable iff it has exactly 2 uses:
130       //   1. the transfer_write we want to hoist.
131       //   2. a matching transfer_read.
132       // Anything else, we skip.
133       bool skip = false;
134       Operation *otherUser = nullptr;
135       for (Operation *u : sliceOp->getUsers()) {
136         if (u == write.transferWriteOp)
137           continue;
138         if (otherUser) {
139           skip = true;
140           break;
141         }
142         otherUser = u;
143       }
144       if (skip || !otherUser)
145         continue;
146       maybeTransferReadUser = otherUser;
147     }
148 
149     LLVM_DEBUG(DBGS() << "maybeTransferReadUser: " << *maybeTransferReadUser
150                       << "\n");
151     auto read = dyn_cast<vector::TransferReadOp>(maybeTransferReadUser);
152     if (read && read.getIndices() == write.transferWriteOp.getIndices() &&
153         read.getVectorType() == write.transferWriteOp.getVectorType())
154       return HoistableRead{read, sliceOp};
155   }
156   return HoistableRead();
157 }
158 
159 /// Check if the chunk of data inserted by the HoistableWrite are read by any
160 /// other op than the HoistableRead candidate.
161 static bool tensorChunkAccessedByUnknownOp(HoistableWrite write,
162                                            HoistableRead candidateRead,
163                                            BlockArgument tensorArg) {
164   // Make sure none of the other uses read the part of the tensor modified
165   // by the transfer_write.
166   llvm::SmallVector<Value::use_range, 1> uses;
167   uses.push_back(tensorArg.getUses());
168   while (!uses.empty()) {
169     for (OpOperand &use : uses.pop_back_val()) {
170       Operation *user = use.getOwner();
171       // Skip the candidate use, only inspect the "other" uses.
172       if (user == candidateRead.transferReadOp ||
173           user == candidateRead.extractSliceOp ||
174           user == write.transferWriteOp || user == write.insertSliceOp)
175         continue;
176       // Consider all transitive uses through a extract_slice / insert_slice.
177       // TODO: atm we just bail because a stronger analysis is needed for these
178       // cases.
179       if (isa<tensor::ExtractSliceOp, tensor::InsertSliceOp>(user))
180         return true;
181       // Consider all transitive uses through a vector.transfer_write.
182       if (auto writeUser = dyn_cast<vector::TransferWriteOp>(user)) {
183         uses.push_back(writeUser->getResult(0).getUses());
184         continue;
185       }
186       // Consider all nested uses through an scf::ForOp. We may have
187       // pass-through tensor arguments left from previous level of
188       // hoisting.
189       if (auto forUser = dyn_cast<scf::ForOp>(user)) {
190         Value arg = forUser.getLoopBody().getArgument(
191             use.getOperandNumber() - forUser.getNumControlOperands() +
192             /*iv value*/ 1);
193         uses.push_back(arg.getUses());
194         continue;
195       }
196       // Follow the use yield as long as it doesn't escape the original
197       // region.
198       scf::YieldOp yieldUser = dyn_cast<scf::YieldOp>(user);
199       if (yieldUser && write.transferWriteOp->getParentOp()->isAncestor(
200                            yieldUser->getParentOp())) {
201         Value ret = yieldUser->getParentOp()->getResult(use.getOperandNumber());
202         uses.push_back(ret.getUses());
203         continue;
204       }
205       auto read = dyn_cast<vector::TransferReadOp>(user);
206       if (!read || !vector::isDisjointTransferIndices(
207                        cast<VectorTransferOpInterface>(read.getOperation()),
208                        cast<VectorTransferOpInterface>(
209                            write.transferWriteOp.getOperation()))) {
210         return true;
211       }
212     }
213   }
214   return false;
215 }
216 
217 /// Return the `forOp`-invariant HoistableWrite that produces `yieldOperand`.
218 /// Return the null HoistableWrite() if it is not comprised of a
219 /// vector.transfer_write + optional insert_slice or if any of the indexings
220 /// is `forOp`-dependent.
221 static HoistableWrite
222 getLoopInvariantTransferWriteOpDefining(scf::ForOp forOp,
223                                         OpOperand &yieldOperand) {
224   Value v = yieldOperand.get();
225   if (auto write = v.getDefiningOp<vector::TransferWriteOp>()) {
226     // Indexing must not depend on `forOp`.
227     for (Value operand : write.getIndices())
228       if (!forOp.isDefinedOutsideOfLoop(operand))
229         return HoistableWrite();
230 
231     return HoistableWrite{write, nullptr};
232   }
233 
234   if (auto insertSliceOp = v.getDefiningOp<tensor::InsertSliceOp>()) {
235     // Inserted slice must come from vector.transfer_write.
236     auto write =
237         insertSliceOp.source().getDefiningOp<vector::TransferWriteOp>();
238     if (!write)
239       return HoistableWrite();
240 
241     // Tensor inserted into must be a BBArg at position matching yieldOperand's.
242     auto bbArg = insertSliceOp.dest().dyn_cast<BlockArgument>();
243     if (!bbArg || bbArg.getOwner()->getParentOp() != forOp ||
244         bbArg.getArgNumber() != /*num iv=*/1 + yieldOperand.getOperandNumber())
245       return HoistableWrite();
246 
247     // Indexing inserted into must not depend on `forOp`.
248     for (Value operand : insertSliceOp->getOperands().drop_front(
249              tensor::InsertSliceOp::getOffsetSizeAndStrideStartOperandIndex()))
250       if (!forOp.isDefinedOutsideOfLoop(operand))
251         return HoistableWrite();
252 
253     return HoistableWrite{write, insertSliceOp};
254   }
255 
256   return HoistableWrite();
257 }
258 
259 /// Mechanical hoisting of a matching HoistableRead / HoistableWrite pair.
260 static void hoistReadWrite(HoistableRead read, HoistableWrite write,
261                            BlockArgument tensorBBArg) {
262   scf::ForOp forOp = cast<scf::ForOp>(tensorBBArg.getOwner()->getParentOp());
263   assert(read.transferReadOp && write.transferWriteOp &&
264          "expected transfer_read and transfer_write ops to be set");
265   assert(((read.extractSliceOp && write.insertSliceOp) ||
266           (!read.extractSliceOp && !write.insertSliceOp)) &&
267          "expected matching extract_slice / insert_slice");
268   LLVM_DEBUG(DBGS() << "In forOp:\n"
269                     << *forOp.getOperation()
270                     << "\nHoist: " << *read.transferReadOp.getOperation()
271                     << "\nHoist: " << *write.transferWriteOp.getOperation()
272                     << "\nInvolving: " << tensorBBArg << "\n");
273 
274   // If a read slice is present, hoist it.
275   if (read.extractSliceOp)
276     forOp.moveOutOfLoop(read.extractSliceOp);
277 
278   // Hoist the transfer_read op.
279   forOp.moveOutOfLoop(read.transferReadOp);
280 
281   // TODO: don't hardcode /*numIvs=*/1.
282   assert(tensorBBArg.getArgNumber() >= /*numIvs=*/1);
283   unsigned initArgNumber = tensorBBArg.getArgNumber() - /*numIvs=*/1;
284 
285   // Update the source tensor.
286   if (read.extractSliceOp)
287     read.extractSliceOp.sourceMutable().assign(
288         forOp.getInitArgs()[initArgNumber]);
289   else
290     read.transferReadOp.getSourceMutable().assign(
291         forOp.getInitArgs()[initArgNumber]);
292 
293   // Hoist write after.
294   if (write.insertSliceOp)
295     write.insertSliceOp->moveAfter(forOp);
296   write.transferWriteOp->moveAfter(forOp);
297 
298   // Update the yield.
299   auto yieldOp = cast<scf::YieldOp>(forOp.getRegion().front().getTerminator());
300   if (write.insertSliceOp)
301     yieldOp->setOperand(initArgNumber, write.insertSliceOp.dest());
302   else
303     yieldOp->setOperand(initArgNumber, write.transferWriteOp.getSource());
304 
305   // Rewrite `loop` with additional new yields.
306   OpBuilder b(read.transferReadOp);
307   auto newForOp = cloneWithNewYields(b, forOp, read.transferReadOp.getVector(),
308                                      write.transferWriteOp.getVector());
309   // Transfer write has been hoisted, need to update the vector and tensor
310   // source. Replace the result of the loop to use the new tensor created
311   // outside the loop.
312   // Depending on whether a insert_slice is present or not, it carries the
313   // update on the tensor operands.
314   if (write.insertSliceOp) {
315     newForOp.getResult(initArgNumber)
316         .replaceAllUsesWith(write.insertSliceOp.getResult());
317     write.transferWriteOp.getSourceMutable().assign(
318         read.extractSliceOp.result());
319     write.insertSliceOp.destMutable().assign(read.extractSliceOp.source());
320   } else {
321     newForOp.getResult(initArgNumber)
322         .replaceAllUsesWith(write.transferWriteOp.getResult());
323     write.transferWriteOp.getSourceMutable().assign(
324         newForOp.getResult(initArgNumber));
325   }
326 
327   // Always update with the newly yield tensor and vector.
328   write.transferWriteOp.getVectorMutable().assign(newForOp.getResults().back());
329 }
330 
331 // To hoist transfer op on tensor the logic can be significantly simplified
332 // compared to the case on buffer. The transformation follows this logic:
333 // 1. Look for transfer_write with a single use from ForOp yield
334 // 2. Check the uses of the matching block argument and look for a transfer_read
335 // with the same indices.
336 // 3. Check that all the other uses of the tensor argument are either disjoint
337 // tensor_read or transfer_write. For transfer_write uses recurse to make sure
338 // the new tensor has the same restrictions on its uses.
339 // 4. Hoist the tensor_read/tensor_write and update the tensor SSA links.
340 // After this transformation the scf.forOp may have unused arguments that can be
341 // remove by the canonicalization pass.
342 void mlir::linalg::hoistRedundantVectorTransfersOnTensor(func::FuncOp func) {
343   bool changed = true;
344   while (changed) {
345     changed = false;
346     func.walk([&](scf::ForOp forOp) {
347       Operation *yield = forOp.getBody()->getTerminator();
348       for (const auto &it : llvm::enumerate(forOp.getRegionIterArgs())) {
349         OpOperand &ret = yield->getOpOperand(it.index());
350         HoistableWrite write =
351             getLoopInvariantTransferWriteOpDefining(forOp, ret);
352         if (!write.transferWriteOp || !write.transferWriteOp->hasOneUse())
353           continue;
354         LLVM_DEBUG(dbgs() << "\n";
355                    DBGS() << "Candidate write for hoisting: "
356                           << *write.transferWriteOp.getOperation() << "\n");
357         if (write.insertSliceOp)
358           LLVM_DEBUG(DBGS() << "Candidate insert_slice for hoisting: "
359                             << *write.insertSliceOp.getOperation() << "\n");
360         if (llvm::any_of(write.transferWriteOp.getIndices(),
361                          [&forOp](Value index) {
362                            return !forOp.isDefinedOutsideOfLoop(index);
363                          }))
364           continue;
365         // Find a read with the same type and indices.
366         HoistableRead matchingRead =
367             findMatchingTransferRead(write, it.value());
368         // Make sure none of the other uses read the part of the tensor modified
369         // by the transfer_write.
370         if (!matchingRead.transferReadOp ||
371             tensorChunkAccessedByUnknownOp(write, matchingRead, it.value()))
372           continue;
373 
374         LLVM_DEBUG(DBGS() << "Start hoisting\n");
375         hoistReadWrite(matchingRead, write, it.value());
376         changed = true;
377         forOp.erase();
378 
379         // Need to interrupt and restart: erasing the loop messes up the walk.
380         return WalkResult::interrupt();
381       }
382       return WalkResult::advance();
383     });
384     // Apply canonicalization so the newForOp + yield folds immediately, thus
385     // cleaning up the IR and potentially enabling more hoisting.
386     if (changed) {
387       RewritePatternSet patterns(func->getContext());
388       scf::ForOp::getCanonicalizationPatterns(patterns, func->getContext());
389       (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
390     }
391   }
392 }
393 
394 void mlir::linalg::hoistRedundantVectorTransfers(func::FuncOp func) {
395   bool changed = true;
396   while (changed) {
397     changed = false;
398     // First move loop invariant ops outside of their loop. This needs to be
399     // done before as we cannot move ops without interputing the function walk.
400     func.walk([&](LoopLikeOpInterface loopLike) {
401       moveLoopInvariantCode(loopLike);
402     });
403 
404     func.walk([&](vector::TransferReadOp transferRead) {
405       if (!transferRead.getShapedType().isa<MemRefType>())
406         return WalkResult::advance();
407 
408       LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
409                         << *transferRead.getOperation() << "\n");
410       auto loop = dyn_cast<scf::ForOp>(transferRead->getParentOp());
411       LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()
412                         << "\n");
413       if (!loop)
414         return WalkResult::advance();
415 
416       LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()
417                         << "\n");
418 
419       SetVector<Operation *> forwardSlice;
420       getForwardSlice(transferRead.getOperation(), &forwardSlice);
421 
422       // Look for the last TransferWriteOp in the forwardSlice of
423       // `transferRead` that operates on the same memref.
424       vector::TransferWriteOp transferWrite;
425       for (auto *sliceOp : llvm::reverse(forwardSlice)) {
426         auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);
427         if (!candidateWrite ||
428             candidateWrite.getSource() != transferRead.getSource())
429           continue;
430         transferWrite = candidateWrite;
431       }
432 
433       // All operands of the TransferRead must be defined outside of the loop.
434       for (auto operand : transferRead.getOperands())
435         if (!loop.isDefinedOutsideOfLoop(operand))
436           return WalkResult::advance();
437 
438       // Only hoist transfer_read / transfer_write pairs for now.
439       if (!transferWrite)
440         return WalkResult::advance();
441 
442       LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()
443                         << "\n");
444 
445       // Approximate aliasing by checking that:
446       //   1. indices are the same,
447       //   2. no other operations in the loop access the same memref except
448       //      for transfer_read/transfer_write accessing statically disjoint
449       //      slices.
450       if (transferRead.getIndices() != transferWrite.getIndices() &&
451           transferRead.getVectorType() == transferWrite.getVectorType())
452         return WalkResult::advance();
453 
454       // TODO: may want to memoize this information for performance but it
455       // likely gets invalidated often.
456       DominanceInfo dom(loop);
457       if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))
458         return WalkResult::advance();
459       for (auto &use : transferRead.getSource().getUses()) {
460         if (!loop->isAncestor(use.getOwner()))
461           continue;
462         if (use.getOwner() == transferRead.getOperation() ||
463             use.getOwner() == transferWrite.getOperation())
464           continue;
465         if (auto transferWriteUse =
466                 dyn_cast<vector::TransferWriteOp>(use.getOwner())) {
467           if (!vector::isDisjointTransferSet(
468                   cast<VectorTransferOpInterface>(transferWrite.getOperation()),
469                   cast<VectorTransferOpInterface>(
470                       transferWriteUse.getOperation())))
471             return WalkResult::advance();
472         } else if (auto transferReadUse =
473                        dyn_cast<vector::TransferReadOp>(use.getOwner())) {
474           if (!vector::isDisjointTransferSet(
475                   cast<VectorTransferOpInterface>(transferWrite.getOperation()),
476                   cast<VectorTransferOpInterface>(
477                       transferReadUse.getOperation())))
478             return WalkResult::advance();
479         } else {
480           // Unknown use, we cannot prove that it doesn't alias with the
481           // transferRead/transferWrite operations.
482           return WalkResult::advance();
483         }
484       }
485 
486       // Hoist read before.
487       loop.moveOutOfLoop(transferRead);
488 
489       // Hoist write after.
490       transferWrite->moveAfter(loop);
491 
492       // Rewrite `loop` with new yields by cloning and erase the original loop.
493       OpBuilder b(transferRead);
494       auto newForOp = cloneWithNewYields(b, loop, transferRead.getVector(),
495                                          transferWrite.getVector());
496 
497       // Transfer write has been hoisted, need to update the written value to
498       // the value yielded by the newForOp.
499       transferWrite.getVector().replaceAllUsesWith(
500           newForOp.getResults().take_back()[0]);
501 
502       changed = true;
503       loop.erase();
504       // Need to interrupt and restart because erasing the loop messes up the
505       // walk.
506       return WalkResult::interrupt();
507     });
508   }
509 }
510