MLIR  22.0.0git
AffineLoopInvariantCodeMotion.cpp
Go to the documentation of this file.
1 //===- AffineLoopInvariantCodeMotion.cpp - Code to perform loop fusion-----===//
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 loop invariant code motion.
10 //
11 //===----------------------------------------------------------------------===//
12 
14 
18 
19 namespace mlir {
20 namespace affine {
21 #define GEN_PASS_DEF_AFFINELOOPINVARIANTCODEMOTION
22 #include "mlir/Dialect/Affine/Passes.h.inc"
23 } // namespace affine
24 } // namespace mlir
25 
26 #define DEBUG_TYPE "affine-licm"
27 
28 using namespace mlir;
29 using namespace mlir::affine;
30 
31 namespace {
32 
33 /// Affine loop invariant code motion (LICM) pass.
34 /// TODO: The pass is missing zero tripcount tests.
35 /// TODO: When compared to the other standard LICM pass, this pass
36 /// has some special handling for affine read/write ops but such handling
37 /// requires aliasing to be sound, and as such this pass is unsound. In
38 /// addition, this handling is nothing particular to affine memory ops but would
39 /// apply to any memory read/write effect ops. Either aliasing should be handled
40 /// or this pass can be removed and the standard LICM can be used.
41 struct LoopInvariantCodeMotion
42  : public affine::impl::AffineLoopInvariantCodeMotionBase<
43  LoopInvariantCodeMotion> {
44  void runOnOperation() override;
45  void runOnAffineForOp(AffineForOp forOp);
46 };
47 } // namespace
48 
49 static bool
50 checkInvarianceOfNestedIfOps(AffineIfOp ifOp, AffineForOp loop,
51  SmallPtrSetImpl<Operation *> &opsWithUsers,
52  SmallPtrSetImpl<Operation *> &opsToHoist);
53 static bool isOpLoopInvariant(Operation &op, AffineForOp loop,
54  SmallPtrSetImpl<Operation *> &opsWithUsers,
55  SmallPtrSetImpl<Operation *> &opsToHoist);
56 
57 static bool
58 areAllOpsInTheBlockListInvariant(Region &blockList, AffineForOp loop,
59  SmallPtrSetImpl<Operation *> &opsWithUsers,
60  SmallPtrSetImpl<Operation *> &opsToHoist);
61 
62 /// Returns true if `op` is invariant on `loop`.
63 static bool isOpLoopInvariant(Operation &op, AffineForOp loop,
64  SmallPtrSetImpl<Operation *> &opsWithUsers,
65  SmallPtrSetImpl<Operation *> &opsToHoist) {
66  Value iv = loop.getInductionVar();
67 
68  if (auto ifOp = dyn_cast<AffineIfOp>(op)) {
69  if (!checkInvarianceOfNestedIfOps(ifOp, loop, opsWithUsers, opsToHoist))
70  return false;
71  } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
72  if (!areAllOpsInTheBlockListInvariant(forOp.getRegion(), loop, opsWithUsers,
73  opsToHoist))
74  return false;
75  } else if (auto parOp = dyn_cast<AffineParallelOp>(op)) {
76  if (!areAllOpsInTheBlockListInvariant(parOp.getRegion(), loop, opsWithUsers,
77  opsToHoist))
78  return false;
79  } else if (!isMemoryEffectFree(&op) &&
80  !isa<AffineReadOpInterface, AffineWriteOpInterface>(&op)) {
81  // Check for side-effecting ops. Affine read/write ops are handled
82  // separately below.
83  return false;
84  } else if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) {
85  // Register op in the set of ops that have users.
86  opsWithUsers.insert(&op);
88  auto read = dyn_cast<AffineReadOpInterface>(op);
89  Value memref =
90  read ? read.getMemRef() : cast<AffineWriteOpInterface>(op).getMemRef();
91  for (auto *user : memref.getUsers()) {
92  // If the memref used by the load/store is used in a store elsewhere in
93  // the loop nest, we do not hoist. Similarly, if the memref used in a
94  // load is also being stored too, we do not hoist the load.
95  // FIXME: This is missing checking aliases.
96  if (&op == user)
97  continue;
98  if (hasEffect<MemoryEffects::Write>(user, memref) ||
99  (hasEffect<MemoryEffects::Read>(user, memref) &&
100  isa<AffineWriteOpInterface>(op))) {
101  userIVs.clear();
102  getAffineForIVs(*user, &userIVs);
103  // Check that userIVs don't contain the for loop around the op.
104  if (llvm::is_contained(userIVs, loop))
105  return false;
106  }
107  }
108  }
109 
110  // Check operands.
111  ValueRange iterArgs = loop.getRegionIterArgs();
112  for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
113  auto *operandSrc = op.getOperand(i).getDefiningOp();
114 
115  // If the loop IV is the operand, this op isn't loop invariant.
116  if (iv == op.getOperand(i))
117  return false;
118 
119  // If the one of the iter_args is the operand, this op isn't loop invariant.
120  if (llvm::is_contained(iterArgs, op.getOperand(i)))
121  return false;
122 
123  if (operandSrc) {
124  // If the value was defined in the loop (outside of the if/else region),
125  // and that operation itself wasn't meant to be hoisted, then mark this
126  // operation loop dependent.
127  if (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0)
128  return false;
129  }
130  }
131 
132  // If no operand was loop variant, mark this op for motion.
133  opsToHoist.insert(&op);
134  return true;
135 }
136 
137 // Checks if all ops in a region (i.e. list of blocks) are loop invariant.
138 static bool
139 areAllOpsInTheBlockListInvariant(Region &blockList, AffineForOp loop,
140  SmallPtrSetImpl<Operation *> &opsWithUsers,
141  SmallPtrSetImpl<Operation *> &opsToHoist) {
142 
143  for (auto &b : blockList) {
144  for (auto &op : b) {
145  if (!isOpLoopInvariant(op, loop, opsWithUsers, opsToHoist))
146  return false;
147  }
148  }
149 
150  return true;
151 }
152 
153 // Returns true if the affine.if op can be hoisted.
154 static bool
155 checkInvarianceOfNestedIfOps(AffineIfOp ifOp, AffineForOp loop,
156  SmallPtrSetImpl<Operation *> &opsWithUsers,
157  SmallPtrSetImpl<Operation *> &opsToHoist) {
158  if (!areAllOpsInTheBlockListInvariant(ifOp.getThenRegion(), loop,
159  opsWithUsers, opsToHoist))
160  return false;
161 
162  if (!areAllOpsInTheBlockListInvariant(ifOp.getElseRegion(), loop,
163  opsWithUsers, opsToHoist))
164  return false;
165 
166  return true;
167 }
168 
169 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
170  // This is the place where hoisted instructions would reside.
171  OpBuilder b(forOp.getOperation());
172 
173  SmallPtrSet<Operation *, 8> opsToHoist;
174  SmallVector<Operation *, 8> opsToMove;
175  SmallPtrSet<Operation *, 8> opsWithUsers;
176 
177  for (Operation &op : *forOp.getBody()) {
178  // Register op in the set of ops that have users. This set is used
179  // to prevent hoisting ops that depend on these ops that are
180  // not being hoisted.
181  if (!op.use_empty())
182  opsWithUsers.insert(&op);
183  if (!isa<AffineYieldOp>(op)) {
184  if (isOpLoopInvariant(op, forOp, opsWithUsers, opsToHoist)) {
185  opsToMove.push_back(&op);
186  }
187  }
188  }
189 
190  // For all instructions that we found to be invariant, place sequentially
191  // right before the for loop.
192  for (auto *op : opsToMove) {
193  op->moveBefore(forOp);
194  }
195 }
196 
197 void LoopInvariantCodeMotion::runOnOperation() {
198  // Walk through all loops in a function in innermost-loop-first order. This
199  // way, we first LICM from the inner loop, and place the ops in
200  // the outer loop, which in turn can be further LICM'ed.
201  getOperation().walk([&](AffineForOp op) { runOnAffineForOp(op); });
202 }
203 
204 std::unique_ptr<OperationPass<func::FuncOp>>
206  return std::make_unique<LoopInvariantCodeMotion>();
207 }
static bool isOpLoopInvariant(Operation &op, AffineForOp loop, SmallPtrSetImpl< Operation * > &opsWithUsers, SmallPtrSetImpl< Operation * > &opsToHoist)
Returns true if op is invariant on loop.
static bool areAllOpsInTheBlockListInvariant(Region &blockList, AffineForOp loop, SmallPtrSetImpl< Operation * > &opsWithUsers, SmallPtrSetImpl< Operation * > &opsToHoist)
static bool checkInvarianceOfNestedIfOps(AffineIfOp ifOp, AffineForOp loop, SmallPtrSetImpl< Operation * > &opsWithUsers, SmallPtrSetImpl< Operation * > &opsToHoist)
This class helps build Operations.
Definition: Builders.h:205
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
Value getOperand(unsigned idx)
Definition: Operation.h:350
unsigned getNumOperands()
Definition: Operation.h:346
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
user_range getUsers() const
Definition: Value.h:218
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:18
std::unique_ptr< OperationPass< func::FuncOp > > createAffineLoopInvariantCodeMotionPass()
Creates a loop invariant code motion pass that hoists loop invariant operations out of affine loops.
void getAffineForIVs(Operation &op, SmallVectorImpl< AffineForOp > *loops)
Populates 'loops' with IVs of the affine.for ops surrounding 'op' ordered from the outermost 'affine....
Definition: Utils.cpp:756
Include the generated interface declarations.
bool isMemoryEffectFree(Operation *op)
Returns true if the given operation is free of memory effects.