208 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements loop invariant code motion.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/Passes.h"14 15#include "mlir/Dialect/Affine/Analysis/Utils.h"16#include "mlir/Dialect/Func/IR/FuncOps.h"17#include "mlir/Interfaces/SideEffectInterfaces.h"18 19namespace mlir {20namespace affine {21#define GEN_PASS_DEF_AFFINELOOPINVARIANTCODEMOTION22#include "mlir/Dialect/Affine/Passes.h.inc"23} // namespace affine24} // namespace mlir25 26#define DEBUG_TYPE "affine-licm"27 28using namespace mlir;29using namespace mlir::affine;30 31namespace {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 pass36/// has some special handling for affine read/write ops but such handling37/// requires aliasing to be sound, and as such this pass is unsound. In38/// addition, this handling is nothing particular to affine memory ops but would39/// apply to any memory read/write effect ops. Either aliasing should be handled40/// or this pass can be removed and the standard LICM can be used.41struct LoopInvariantCodeMotion42 : public affine::impl::AffineLoopInvariantCodeMotionBase<43 LoopInvariantCodeMotion> {44 void runOnOperation() override;45 void runOnAffineForOp(AffineForOp forOp);46};47} // namespace48 49static bool50checkInvarianceOfNestedIfOps(AffineIfOp ifOp, AffineForOp loop,51 SmallPtrSetImpl<Operation *> &opsWithUsers,52 SmallPtrSetImpl<Operation *> &opsToHoist);53static bool isOpLoopInvariant(Operation &op, AffineForOp loop,54 SmallPtrSetImpl<Operation *> &opsWithUsers,55 SmallPtrSetImpl<Operation *> &opsToHoist);56 57static bool58areAllOpsInTheBlockListInvariant(Region &blockList, AffineForOp loop,59 SmallPtrSetImpl<Operation *> &opsWithUsers,60 SmallPtrSetImpl<Operation *> &opsToHoist);61 62/// Returns true if `op` is invariant on `loop`.63static 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 handled82 // 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);87 SmallVector<AffineForOp, 8> userIVs;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 in93 // the loop nest, we do not hoist. Similarly, if the memref used in a94 // 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 this126 // 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.138static bool139areAllOpsInTheBlockListInvariant(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.154static bool155checkInvarianceOfNestedIfOps(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 169void 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 used179 // to prevent hoisting ops that depend on these ops that are180 // 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 sequentially191 // right before the for loop.192 for (auto *op : opsToMove) {193 op->moveBefore(forOp);194 }195}196 197void LoopInvariantCodeMotion::runOnOperation() {198 // Walk through all loops in a function in innermost-loop-first order. This199 // way, we first LICM from the inner loop, and place the ops in200 // the outer loop, which in turn can be further LICM'ed.201 getOperation().walk([&](AffineForOp op) { runOnAffineForOp(op); });202}203 204std::unique_ptr<OperationPass<func::FuncOp>>205mlir::affine::createAffineLoopInvariantCodeMotionPass() {206 return std::make_unique<LoopInvariantCodeMotion>();207}208