408 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements functions concerned with hoisting invariant operations10// 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/AffineOps.h"18#include "mlir/Dialect/Affine/Utils.h"19#include "mlir/Dialect/Linalg/Transforms/Transforms.h"20#include "mlir/Dialect/SCF/IR/SCF.h"21#include "mlir/Dialect/SCF/Utils/Utils.h"22#include "mlir/Dialect/Vector/IR/VectorOps.h"23#include "mlir/Dialect/Vector/Utils/VectorUtils.h"24#include "mlir/IR/Dominance.h"25#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"26#include "llvm/Support/Debug.h"27 28using llvm::dbgs;29 30#define DEBUG_TYPE "linalg-hoisting"31 32#define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")33 34using namespace mlir;35using namespace mlir::linalg;36 37/// Replace `loop` with a new loop that has a different init operand at38/// position `index`. The body of this loop is moved over to the new loop.39///40/// `newInitOperands` specifies the replacement "init" operands.41/// `newYieldValue` is the replacement yield value of the loop at position42/// `index`.43static scf::ForOp replaceWithDifferentYield(RewriterBase &rewriter,44 scf::ForOp loop,45 Value newInitOperand,46 unsigned index,47 Value newYieldValue) {48 OpBuilder::InsertionGuard g(rewriter);49 rewriter.setInsertionPoint(loop.getOperation());50 auto inits = llvm::to_vector(loop.getInits());51 52 // Replace the init value with the new operand.53 assert(index < inits.size());54 inits[index] = newInitOperand;55 56 scf::ForOp newLoop = scf::ForOp::create(57 rewriter, loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(),58 loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},59 loop.getUnsignedCmp());60 61 // Generate the new yield with the replaced operand.62 auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());63 yieldOp.setOperand(index, newYieldValue);64 65 // Move the loop body to the new op.66 rewriter.mergeBlocks(loop.getBody(), newLoop.getBody(),67 newLoop.getBody()->getArguments());68 69 // Replace the old loop.70 rewriter.replaceOp(loop.getOperation(), newLoop->getResults());71 return newLoop;72}73 74// Hoist out a pair of corresponding vector.extract+vector.broadcast75// operations. This function transforms a loop like this:76// %res = scf.for _ = _ to _ step _ iter_args(%iarg = %v) -> (t1) {77// %e = vector.extract %iarg : t1 to t278// %u = "some_use"(%e) : (t2) -> t279// %b = vector.broadcast %u : t2 to t180// scf.yield %b : t181// }82// into the following:83// %e = vector.extract %v: t1 to t284// %res' = scf.for _ = _ to _ step _ iter_args(%iarg = %e) -> (t2) {85// %u' = "some_use"(%iarg) : (t2) -> t286// scf.yield %u' : t287// }88// %res = vector.broadcast %res' : t2 to t189void mlir::linalg::hoistRedundantVectorBroadcasts(RewriterBase &rewriter,90 Operation *root) {91 bool changed = true;92 while (changed) {93 changed = false;94 // First move loop invariant ops outside of their loop. This needs to be95 // done before as we cannot move ops without interrupting the function walk.96 root->walk(97 [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });98 99 root->walk([&](vector::ExtractOp extractOp) {100 LLVM_DEBUG(DBGS() << "Candidate for hoisting: "101 << *extractOp.getOperation() << "\n");102 103 auto loop = dyn_cast<scf::ForOp>(extractOp->getParentOp());104 if (!loop)105 return WalkResult::advance();106 107 // Check that the vector to extract from is a BlockArgument.108 auto blockArg = dyn_cast<BlockArgument>(extractOp.getSource());109 if (!blockArg)110 return WalkResult::advance();111 112 // Check that the blockArg is an iter_arg of the loop.113 OpOperand *initArg = loop.getTiedLoopInit(blockArg);114 if (!initArg)115 return WalkResult::advance();116 117 // If the iter_arg does not have only one use, it won't be possible to118 // hoist the extractOp out.119 if (!blockArg.hasOneUse())120 return WalkResult::advance();121 122 unsigned index = blockArg.getArgNumber() - loop.getNumInductionVars();123 124 // Check that the loop yields a broadcast that has just one use.125 Operation *yieldedVal =126 loop.getTiedLoopYieldedValue(blockArg)->get().getDefiningOp();127 auto broadcast = dyn_cast<vector::BroadcastOp>(yieldedVal);128 if (!broadcast || !broadcast.getResult().hasOneUse())129 return WalkResult::advance();130 131 LLVM_DEBUG(DBGS() << "Candidate broadcast: " << broadcast << "\n");132 133 Type broadcastInputType = broadcast.getSourceType();134 if (broadcastInputType != extractOp.getType())135 return WalkResult::advance();136 137 // The position of the extract must be defined outside of the loop if138 // it is dynamic.139 for (auto operand : extractOp.getDynamicPosition())140 if (!loop.isDefinedOutsideOfLoop(operand))141 return WalkResult::advance();142 143 rewriter.modifyOpInPlace(broadcast, [&] {144 extractOp.getSourceMutable().assign(initArg->get());145 });146 loop.moveOutOfLoop(extractOp);147 rewriter.moveOpAfter(broadcast, loop);148 149 scf::ForOp newLoop = replaceWithDifferentYield(150 rewriter, loop, extractOp.getResult(), index, broadcast.getSource());151 152 LLVM_DEBUG(DBGS() << "New loop: " << newLoop << "\n");153 154 rewriter.replaceAllUsesWith(newLoop.getResult(index), broadcast);155 rewriter.modifyOpInPlace(156 broadcast, [&] { broadcast.setOperand(newLoop.getResult(index)); });157 158 changed = true;159 return WalkResult::interrupt();160 });161 }162}163 164static bool noAliasingUseInLoop(vector::TransferReadOp transferRead,165 LoopLikeOpInterface loop) {166 Value source = transferRead.getBase();167 168 // Skip view-like Ops and retrive the actual soruce Operation169 while (auto viewLike = source.getDefiningOp<ViewLikeOpInterface>()) {170 if (viewLike.getViewDest() != source) {171 break;172 }173 source = viewLike.getViewSource();174 }175 176 llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),177 source.getUsers().end());178 llvm::SmallDenseSet<Operation *, 32> processed;179 while (!users.empty()) {180 Operation *user = users.pop_back_val();181 // If the user has already been processed skip.182 if (!processed.insert(user).second)183 continue;184 if (auto viewLike = dyn_cast<ViewLikeOpInterface>(user)) {185 Value viewDest = viewLike.getViewDest();186 users.append(viewDest.getUsers().begin(), viewDest.getUsers().end());187 continue;188 }189 if (isMemoryEffectFree(user) || isa<vector::TransferReadOp>(user))190 continue;191 if (!loop->isAncestor(user))192 continue;193 return false;194 }195 return true;196}197 198void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,199 bool verifyNonZeroTrip) {200 bool changed = true;201 while (changed) {202 changed = false;203 // First move loop invariant ops outside of their loop. This needs to be204 // done before as we cannot move ops without interrupting the function walk.205 root->walk(206 [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });207 208 // Find all loops that are certain to have non zero trip count. Any loops209 // that are not part of this set cannot be hoisted from, since hoisting from210 // a potentially zero trip count loop may cause a vector transfer to be211 // executed when it shouldn't be.212 llvm::DenseSet<LoopLikeOpInterface> definiteNonZeroTripCountLoops;213 if (verifyNonZeroTrip) {214 root->walk([&](LoopLikeOpInterface loopLike) {215 std::optional<SmallVector<OpFoldResult>> lbs =216 loopLike.getLoopLowerBounds();217 std::optional<SmallVector<OpFoldResult>> ubs =218 loopLike.getLoopUpperBounds();219 // If loop bounds cannot be found, assume possibly zero trip count.220 if (!lbs || !ubs)221 return;222 223 // Otherwise, use ValueBounds to find the maximum lower bound and224 // minimum upper bound. If the bounds are found, and maxLb is less225 // than the minUb, then the loop will not have zero trip count.226 for (auto [lb, ub] : llvm::zip_equal(lbs.value(), ubs.value())) {227 FailureOr<int64_t> maxLb =228 ValueBoundsConstraintSet::computeConstantBound(229 presburger::BoundType::UB, lb,230 /*stopCondition=*/nullptr, /*closedUB=*/true);231 if (failed(maxLb))232 return;233 FailureOr<int64_t> minUb =234 ValueBoundsConstraintSet::computeConstantBound(235 presburger::BoundType::LB, ub);236 if (failed(minUb))237 return;238 if (minUb.value() <= maxLb.value())239 return;240 definiteNonZeroTripCountLoops.insert(loopLike);241 }242 });243 }244 245 root->walk([&](vector::TransferReadOp transferRead) {246 if (!isa<MemRefType>(transferRead.getShapedType()))247 return WalkResult::advance();248 249 LLVM_DEBUG(DBGS() << "Candidate for hoisting: "250 << *transferRead.getOperation() << "\n");251 auto loop = dyn_cast<LoopLikeOpInterface>(transferRead->getParentOp());252 LLVM_DEBUG(DBGS() << "Parent op: " << *transferRead->getParentOp()253 << "\n");254 if (!isa_and_nonnull<scf::ForOp, affine::AffineForOp>(loop))255 return WalkResult::advance();256 257 if (verifyNonZeroTrip && !definiteNonZeroTripCountLoops.contains(loop)) {258 LLVM_DEBUG(DBGS() << "Loop may have zero trip count: " << *loop259 << "\n");260 return WalkResult::advance();261 }262 263 LLVM_DEBUG(DBGS() << "Candidate read: " << *transferRead.getOperation()264 << "\n");265 266 SetVector<Operation *> forwardSlice;267 getForwardSlice(transferRead.getOperation(), &forwardSlice);268 269 // Look for the last TransferWriteOp in the forwardSlice of270 // `transferRead` that operates on the same memref.271 vector::TransferWriteOp transferWrite;272 for (auto *sliceOp : llvm::reverse(forwardSlice)) {273 auto candidateWrite = dyn_cast<vector::TransferWriteOp>(sliceOp);274 if (!candidateWrite ||275 candidateWrite.getBase() != transferRead.getBase())276 continue;277 transferWrite = candidateWrite;278 }279 280 // All operands of the TransferRead must be defined outside of the loop.281 for (auto operand : transferRead.getOperands())282 if (!loop.isDefinedOutsideOfLoop(operand))283 return WalkResult::advance();284 285 // Only hoist transfer_read / transfer_write pairs and singleton286 // transfer_reads for now.287 if (!transferWrite) {288 // Make sure there are no other accesses to the memref before289 // hoisting transfer_read.290 if (noAliasingUseInLoop(transferRead, loop))291 loop.moveOutOfLoop(transferRead);292 return WalkResult::advance();293 }294 295 LLVM_DEBUG(DBGS() << "Candidate: " << *transferWrite.getOperation()296 << "\n");297 298 // Approximate aliasing by checking that:299 // 1. indices, vector type and permutation map are the same (i.e., the300 // transfer_read/transfer_write ops are matching),301 // 2. source operands for transfer.{read|write} do not originate from302 // nor have users that are Ops implementing ViewLikeOpInterface.303 // 3. no other operations in the loop access the same memref except304 // for transfer_read/transfer_write accessing statically disjoint305 // slices.306 307 // Check 1.308 if (transferRead.getIndices() != transferWrite.getIndices() ||309 transferRead.getVectorType() != transferWrite.getVectorType() ||310 transferRead.getPermutationMap() != transferWrite.getPermutationMap())311 return WalkResult::advance();312 313 // Check 2. Note, since both xfer Ops share the source, we only need to314 // look at one of them.315 auto base = transferRead.getBase();316 auto *source = base.getDefiningOp();317 if (source) {318 // NOTE: We treat `memref.assume_alignment` as a special case.319 //320 // The idea is that it is safe to look past AssumeAlignmemtOp (i.e.321 // MemRef _before_ alignment) iff:322 // 1. It has exactly two uses (these have to be the xfer Ops323 // being looked at).324 // 2. The original MemRef has only one use (i.e.325 // AssumeAlignmentOp).326 //327 // Relaxing these conditions will most likely require proper alias328 // analysis.329 if (auto assume = dyn_cast<memref::AssumeAlignmentOp>(source)) {330 Value memPreAlignment = assume.getMemref();331 auto numInLoopUses =332 llvm::count_if(base.getUses(), [&loop](OpOperand &use) {333 return loop->isAncestor(use.getOwner());334 });335 336 if (numInLoopUses && memPreAlignment.hasOneUse())337 source = memPreAlignment.getDefiningOp();338 }339 if (isa_and_nonnull<ViewLikeOpInterface>(source))340 return WalkResult::advance();341 }342 343 if (llvm::any_of(base.getUsers(), llvm::IsaPred<ViewLikeOpInterface>))344 return WalkResult::advance();345 346 // Check 3.347 // TODO: may want to memoize this information for performance but it348 // likely gets invalidated often.349 DominanceInfo dom(loop);350 if (!dom.properlyDominates(transferRead.getOperation(), transferWrite))351 return WalkResult::advance();352 for (auto &use : transferRead.getBase().getUses()) {353 if (!loop->isAncestor(use.getOwner()))354 continue;355 if (use.getOwner() == transferRead.getOperation() ||356 use.getOwner() == transferWrite.getOperation())357 continue;358 if (auto transferWriteUse =359 dyn_cast<vector::TransferWriteOp>(use.getOwner())) {360 if (!vector::isDisjointTransferSet(361 cast<VectorTransferOpInterface>(*transferWrite),362 cast<VectorTransferOpInterface>(*transferWriteUse),363 /*testDynamicValueUsingBounds=*/true))364 return WalkResult::advance();365 } else if (auto transferReadUse =366 dyn_cast<vector::TransferReadOp>(use.getOwner())) {367 if (!vector::isDisjointTransferSet(368 cast<VectorTransferOpInterface>(*transferWrite),369 cast<VectorTransferOpInterface>(*transferReadUse),370 /*testDynamicValueUsingBounds=*/true))371 return WalkResult::advance();372 } else {373 // Unknown use, we cannot prove that it doesn't alias with the374 // transferRead/transferWrite operations.375 return WalkResult::advance();376 }377 }378 379 // Hoist read before.380 loop.moveOutOfLoop(transferRead);381 382 // Hoist write after.383 transferWrite->moveAfter(loop);384 385 // Rewrite `loop` with new yields by cloning and erase the original386 // loop.387 IRRewriter rewriter(transferRead.getContext());388 NewYieldValuesFn yieldFn = [&](OpBuilder &b, Location loc,389 ArrayRef<BlockArgument> newBBArgs) {390 return SmallVector<Value>{transferWrite.getVector()};391 };392 393 auto maybeNewLoop = loop.replaceWithAdditionalYields(394 rewriter, transferRead.getVector(),395 /*replaceInitOperandUsesInLoop=*/true, yieldFn);396 if (failed(maybeNewLoop))397 return WalkResult::interrupt();398 399 transferWrite.getValueToStoreMutable().assign(400 maybeNewLoop->getOperation()->getResults().back());401 changed = true;402 // Need to interrupt and restart because erasing the loop messes up403 // the walk.404 return WalkResult::interrupt();405 });406 }407}408