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 13 #include "PassDetail.h" 14 #include "mlir/Analysis/SliceAnalysis.h" 15 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" 16 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h" 17 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" 18 #include "mlir/Dialect/Affine/Analysis/Utils.h" 19 #include "mlir/Dialect/Affine/IR/AffineOps.h" 20 #include "mlir/Dialect/Affine/LoopUtils.h" 21 #include "mlir/Dialect/Affine/Passes.h" 22 #include "mlir/Dialect/Affine/Utils.h" 23 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" 24 #include "mlir/IR/AffineExpr.h" 25 #include "mlir/IR/AffineMap.h" 26 #include "mlir/IR/Builders.h" 27 #include "mlir/IR/Matchers.h" 28 #include "llvm/ADT/DenseMap.h" 29 #include "llvm/ADT/DenseSet.h" 30 #include "llvm/ADT/SmallPtrSet.h" 31 #include "llvm/Support/CommandLine.h" 32 #include "llvm/Support/Debug.h" 33 #include "llvm/Support/raw_ostream.h" 34 35 #define DEBUG_TYPE "licm" 36 37 using namespace mlir; 38 39 namespace { 40 41 /// Loop invariant code motion (LICM) pass. 42 /// TODO: The pass is missing zero-trip tests. 43 /// TODO: Check for the presence of side effects before hoisting. 44 /// TODO: This code should be removed once the new LICM pass can handle its 45 /// uses. 46 struct LoopInvariantCodeMotion 47 : public AffineLoopInvariantCodeMotionBase<LoopInvariantCodeMotion> { 48 void runOnOperation() override; 49 void runOnAffineForOp(AffineForOp forOp); 50 }; 51 } // namespace 52 53 static bool 54 checkInvarianceOfNestedIfOps(Operation *op, Value indVar, ValueRange iterArgs, 55 SmallPtrSetImpl<Operation *> &opsWithUsers, 56 SmallPtrSetImpl<Operation *> &opsToHoist); 57 static bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs, 58 SmallPtrSetImpl<Operation *> &opsWithUsers, 59 SmallPtrSetImpl<Operation *> &opsToHoist); 60 61 static bool 62 areAllOpsInTheBlockListInvariant(Region &blockList, Value indVar, 63 ValueRange iterArgs, 64 SmallPtrSetImpl<Operation *> &opsWithUsers, 65 SmallPtrSetImpl<Operation *> &opsToHoist); 66 67 // Returns true if the individual op is loop invariant. 68 bool isOpLoopInvariant(Operation &op, Value indVar, ValueRange iterArgs, 69 SmallPtrSetImpl<Operation *> &opsWithUsers, 70 SmallPtrSetImpl<Operation *> &opsToHoist) { 71 LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;); 72 73 if (isa<AffineIfOp>(op)) { 74 if (!checkInvarianceOfNestedIfOps(&op, indVar, iterArgs, opsWithUsers, 75 opsToHoist)) { 76 return false; 77 } 78 } else if (auto forOp = dyn_cast<AffineForOp>(op)) { 79 if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar, iterArgs, 80 opsWithUsers, opsToHoist)) { 81 return false; 82 } 83 } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) { 84 // TODO: Support DMA ops. 85 return false; 86 } else if (!matchPattern(&op, m_Constant())) { 87 // Register op in the set of ops that have users. 88 opsWithUsers.insert(&op); 89 if (isa<AffineMapAccessInterface>(op)) { 90 Value memref = isa<AffineReadOpInterface>(op) 91 ? cast<AffineReadOpInterface>(op).getMemRef() 92 : cast<AffineWriteOpInterface>(op).getMemRef(); 93 for (auto *user : memref.getUsers()) { 94 // If this memref has a user that is a DMA, give up because these 95 // operations write to this memref. 96 if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) { 97 return false; 98 } 99 // If the memref used by the load/store is used in a store elsewhere in 100 // the loop nest, we do not hoist. Similarly, if the memref used in a 101 // load is also being stored too, we do not hoist the load. 102 if (isa<AffineWriteOpInterface>(user) || 103 (isa<AffineReadOpInterface>(user) && 104 isa<AffineWriteOpInterface>(op))) { 105 if (&op != user) { 106 SmallVector<AffineForOp, 8> userIVs; 107 getLoopIVs(*user, &userIVs); 108 // Check that userIVs don't contain the for loop around the op. 109 if (llvm::is_contained(userIVs, getForInductionVarOwner(indVar))) { 110 return false; 111 } 112 } 113 } 114 } 115 } 116 117 if (op.getNumOperands() == 0 && !isa<AffineYieldOp>(op)) { 118 LLVM_DEBUG(llvm::dbgs() << "\nNon-constant op with 0 operands\n"); 119 return false; 120 } 121 } 122 123 // Check operands. 124 for (unsigned int i = 0; i < op.getNumOperands(); ++i) { 125 auto *operandSrc = op.getOperand(i).getDefiningOp(); 126 127 LLVM_DEBUG( 128 op.getOperand(i).print(llvm::dbgs() << "\nIterating on operand\n")); 129 130 // If the loop IV is the operand, this op isn't loop invariant. 131 if (indVar == op.getOperand(i)) { 132 LLVM_DEBUG(llvm::dbgs() << "\nLoop IV is the operand\n"); 133 return false; 134 } 135 136 // If the one of the iter_args is the operand, this op isn't loop invariant. 137 if (llvm::is_contained(iterArgs, op.getOperand(i))) { 138 LLVM_DEBUG(llvm::dbgs() << "\nOne of the iter_args is the operand\n"); 139 return false; 140 } 141 142 if (operandSrc != nullptr) { 143 LLVM_DEBUG(llvm::dbgs() << *operandSrc << "\nIterating on operand src\n"); 144 145 // If the value was defined in the loop (outside of the 146 // if/else region), and that operation itself wasn't meant to 147 // be hoisted, then mark this operation loop dependent. 148 if (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0) { 149 return false; 150 } 151 } 152 } 153 154 // If no operand was loop variant, mark this op for motion. 155 opsToHoist.insert(&op); 156 return true; 157 } 158 159 // Checks if all ops in a region (i.e. list of blocks) are loop invariant. 160 bool areAllOpsInTheBlockListInvariant( 161 Region &blockList, Value indVar, ValueRange iterArgs, 162 SmallPtrSetImpl<Operation *> &opsWithUsers, 163 SmallPtrSetImpl<Operation *> &opsToHoist) { 164 165 for (auto &b : blockList) { 166 for (auto &op : b) { 167 if (!isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) { 168 return false; 169 } 170 } 171 } 172 173 return true; 174 } 175 176 // Returns true if the affine.if op can be hoisted. 177 bool checkInvarianceOfNestedIfOps(Operation *op, Value indVar, 178 ValueRange iterArgs, 179 SmallPtrSetImpl<Operation *> &opsWithUsers, 180 SmallPtrSetImpl<Operation *> &opsToHoist) { 181 assert(isa<AffineIfOp>(op)); 182 auto ifOp = cast<AffineIfOp>(op); 183 184 if (!areAllOpsInTheBlockListInvariant(ifOp.getThenRegion(), indVar, iterArgs, 185 opsWithUsers, opsToHoist)) { 186 return false; 187 } 188 189 if (!areAllOpsInTheBlockListInvariant(ifOp.getElseRegion(), indVar, iterArgs, 190 opsWithUsers, opsToHoist)) { 191 return false; 192 } 193 194 return true; 195 } 196 197 void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) { 198 auto *loopBody = forOp.getBody(); 199 auto indVar = forOp.getInductionVar(); 200 ValueRange iterArgs = forOp.getRegionIterArgs(); 201 202 // This is the place where hoisted instructions would reside. 203 OpBuilder b(forOp.getOperation()); 204 205 SmallPtrSet<Operation *, 8> opsToHoist; 206 SmallVector<Operation *, 8> opsToMove; 207 SmallPtrSet<Operation *, 8> opsWithUsers; 208 209 for (auto &op : *loopBody) { 210 // Register op in the set of ops that have users. This set is used 211 // to prevent hoisting ops that depend on these ops that are 212 // not being hoisted. 213 if (!op.use_empty()) 214 opsWithUsers.insert(&op); 215 if (!isa<AffineYieldOp>(op)) { 216 if (isOpLoopInvariant(op, indVar, iterArgs, opsWithUsers, opsToHoist)) { 217 opsToMove.push_back(&op); 218 } 219 } 220 } 221 222 // For all instructions that we found to be invariant, place sequentially 223 // right before the for loop. 224 for (auto *op : opsToMove) { 225 op->moveBefore(forOp); 226 } 227 228 LLVM_DEBUG(forOp->print(llvm::dbgs() << "Modified loop\n")); 229 } 230 231 void LoopInvariantCodeMotion::runOnOperation() { 232 // Walk through all loops in a function in innermost-loop-first order. This 233 // way, we first LICM from the inner loop, and place the ops in 234 // the outer loop, which in turn can be further LICM'ed. 235 getOperation().walk([&](AffineForOp op) { 236 LLVM_DEBUG(op->print(llvm::dbgs() << "\nOriginal loop\n")); 237 runOnAffineForOp(op); 238 }); 239 } 240 241 std::unique_ptr<OperationPass<func::FuncOp>> 242 mlir::createAffineLoopInvariantCodeMotionPass() { 243 return std::make_unique<LoopInvariantCodeMotion>(); 244 } 245