289 lines · cpp
1//===- ParallelLoopFusion.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 fusion on parallel loops.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SCF/Transforms/Passes.h"14 15#include "mlir/Analysis/AliasAnalysis.h"16#include "mlir/Dialect/MemRef/IR/MemRef.h"17#include "mlir/Dialect/SCF/IR/SCF.h"18#include "mlir/Dialect/SCF/Transforms/Transforms.h"19#include "mlir/IR/Builders.h"20#include "mlir/IR/IRMapping.h"21#include "mlir/IR/OpDefinition.h"22#include "mlir/IR/OperationSupport.h"23#include "mlir/Interfaces/SideEffectInterfaces.h"24 25namespace mlir {26#define GEN_PASS_DEF_SCFPARALLELLOOPFUSION27#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"28} // namespace mlir29 30using namespace mlir;31using namespace mlir::scf;32 33/// Verify there are no nested ParallelOps.34static bool hasNestedParallelOp(ParallelOp ploop) {35 auto walkResult =36 ploop.getBody()->walk([](ParallelOp) { return WalkResult::interrupt(); });37 return walkResult.wasInterrupted();38}39 40/// Verify equal iteration spaces.41static bool equalIterationSpaces(ParallelOp firstPloop,42 ParallelOp secondPloop) {43 if (firstPloop.getNumLoops() != secondPloop.getNumLoops())44 return false;45 46 auto matchOperands = [&](const OperandRange &lhs,47 const OperandRange &rhs) -> bool {48 // TODO: Extend this to support aliases and equal constants.49 return std::equal(lhs.begin(), lhs.end(), rhs.begin());50 };51 return matchOperands(firstPloop.getLowerBound(),52 secondPloop.getLowerBound()) &&53 matchOperands(firstPloop.getUpperBound(),54 secondPloop.getUpperBound()) &&55 matchOperands(firstPloop.getStep(), secondPloop.getStep());56}57 58/// Checks if the parallel loops have mixed access to the same buffers. Returns59/// `true` if the first parallel loop writes to the same indices that the second60/// loop reads.61static bool haveNoReadsAfterWriteExceptSameIndex(62 ParallelOp firstPloop, ParallelOp secondPloop,63 const IRMapping &firstToSecondPloopIndices,64 llvm::function_ref<bool(Value, Value)> mayAlias) {65 DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;66 SmallVector<Value> bufferStoresVec;67 firstPloop.getBody()->walk([&](memref::StoreOp store) {68 bufferStores[store.getMemRef()].push_back(store.getIndices());69 bufferStoresVec.emplace_back(store.getMemRef());70 });71 auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {72 Value loadMem = load.getMemRef();73 // Stop if the memref is defined in secondPloop body. Careful alias analysis74 // is needed.75 auto *memrefDef = loadMem.getDefiningOp();76 if (memrefDef && memrefDef->getBlock() == load->getBlock())77 return WalkResult::interrupt();78 79 for (Value store : bufferStoresVec)80 if (store != loadMem && mayAlias(store, loadMem))81 return WalkResult::interrupt();82 83 auto write = bufferStores.find(loadMem);84 if (write == bufferStores.end())85 return WalkResult::advance();86 87 // Check that at last one store was retrieved88 if (write->second.empty())89 return WalkResult::interrupt();90 91 auto storeIndices = write->second.front();92 93 // Multiple writes to the same memref are allowed only on the same indices94 for (const auto &othStoreIndices : write->second) {95 if (othStoreIndices != storeIndices)96 return WalkResult::interrupt();97 }98 99 // Check that the load indices of secondPloop coincide with store indices of100 // firstPloop for the same memrefs.101 auto loadIndices = load.getIndices();102 if (storeIndices.size() != loadIndices.size())103 return WalkResult::interrupt();104 for (int i = 0, e = storeIndices.size(); i < e; ++i) {105 if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=106 loadIndices[i]) {107 auto *storeIndexDefOp = storeIndices[i].getDefiningOp();108 auto *loadIndexDefOp = loadIndices[i].getDefiningOp();109 if (storeIndexDefOp && loadIndexDefOp) {110 if (!isMemoryEffectFree(storeIndexDefOp))111 return WalkResult::interrupt();112 if (!isMemoryEffectFree(loadIndexDefOp))113 return WalkResult::interrupt();114 if (!OperationEquivalence::isEquivalentTo(115 storeIndexDefOp, loadIndexDefOp,116 [&](Value storeIndex, Value loadIndex) {117 if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=118 firstToSecondPloopIndices.lookupOrDefault(loadIndex))119 return failure();120 else121 return success();122 },123 /*markEquivalent=*/nullptr,124 OperationEquivalence::Flags::IgnoreLocations)) {125 return WalkResult::interrupt();126 }127 } else {128 return WalkResult::interrupt();129 }130 }131 }132 return WalkResult::advance();133 });134 return !walkResult.wasInterrupted();135}136 137/// Analyzes dependencies in the most primitive way by checking simple read and138/// write patterns.139static LogicalResult140verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,141 const IRMapping &firstToSecondPloopIndices,142 llvm::function_ref<bool(Value, Value)> mayAlias) {143 if (!haveNoReadsAfterWriteExceptSameIndex(144 firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))145 return failure();146 147 IRMapping secondToFirstPloopIndices;148 secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),149 firstPloop.getBody()->getArguments());150 return success(haveNoReadsAfterWriteExceptSameIndex(151 secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));152}153 154static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,155 const IRMapping &firstToSecondPloopIndices,156 llvm::function_ref<bool(Value, Value)> mayAlias) {157 return !hasNestedParallelOp(firstPloop) &&158 !hasNestedParallelOp(secondPloop) &&159 equalIterationSpaces(firstPloop, secondPloop) &&160 succeeded(verifyDependencies(firstPloop, secondPloop,161 firstToSecondPloopIndices, mayAlias));162}163 164/// Prepends operations of firstPloop's body into secondPloop's body.165/// Updates secondPloop with new loop.166static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,167 OpBuilder builder,168 llvm::function_ref<bool(Value, Value)> mayAlias) {169 Block *block1 = firstPloop.getBody();170 Block *block2 = secondPloop.getBody();171 IRMapping firstToSecondPloopIndices;172 firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());173 174 if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,175 mayAlias))176 return;177 178 DominanceInfo dom;179 // We are fusing first loop into second, make sure there are no users of the180 // first loop results between loops.181 for (Operation *user : firstPloop->getUsers())182 if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))183 return;184 185 ValueRange inits1 = firstPloop.getInitVals();186 ValueRange inits2 = secondPloop.getInitVals();187 188 SmallVector<Value> newInitVars(inits1.begin(), inits1.end());189 newInitVars.append(inits2.begin(), inits2.end());190 191 IRRewriter b(builder);192 b.setInsertionPoint(secondPloop);193 auto newSecondPloop = ParallelOp::create(194 b, secondPloop.getLoc(), secondPloop.getLowerBound(),195 secondPloop.getUpperBound(), secondPloop.getStep(), newInitVars);196 197 Block *newBlock = newSecondPloop.getBody();198 auto term1 = cast<ReduceOp>(block1->getTerminator());199 auto term2 = cast<ReduceOp>(block2->getTerminator());200 201 b.inlineBlockBefore(block2, newBlock, newBlock->begin(),202 newBlock->getArguments());203 b.inlineBlockBefore(block1, newBlock, newBlock->begin(),204 newBlock->getArguments());205 206 ValueRange results = newSecondPloop.getResults();207 if (!results.empty()) {208 b.setInsertionPointToEnd(newBlock);209 210 ValueRange reduceArgs1 = term1.getOperands();211 ValueRange reduceArgs2 = term2.getOperands();212 SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());213 newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());214 215 auto newReduceOp = scf::ReduceOp::create(b, term2.getLoc(), newReduceArgs);216 217 for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(218 term1.getReductions(), term2.getReductions()))) {219 Block &oldRedBlock = reg.front();220 Block &newRedBlock = newReduceOp.getReductions()[i].front();221 b.inlineBlockBefore(&oldRedBlock, &newRedBlock, newRedBlock.begin(),222 newRedBlock.getArguments());223 }224 225 firstPloop.replaceAllUsesWith(results.take_front(inits1.size()));226 secondPloop.replaceAllUsesWith(results.take_back(inits2.size()));227 }228 term1->erase();229 term2->erase();230 firstPloop.erase();231 secondPloop.erase();232 secondPloop = newSecondPloop;233}234 235void mlir::scf::naivelyFuseParallelOps(236 Region ®ion, llvm::function_ref<bool(Value, Value)> mayAlias) {237 OpBuilder b(region);238 // Consider every single block and attempt to fuse adjacent loops.239 SmallVector<SmallVector<ParallelOp>, 1> ploopChains;240 for (auto &block : region) {241 ploopChains.clear();242 ploopChains.push_back({});243 244 // Not using `walk()` to traverse only top-level parallel loops and also245 // make sure that there are no side-effecting ops between the parallel246 // loops.247 bool noSideEffects = true;248 for (auto &op : block) {249 if (auto ploop = dyn_cast<ParallelOp>(op)) {250 if (noSideEffects) {251 ploopChains.back().push_back(ploop);252 } else {253 ploopChains.push_back({ploop});254 noSideEffects = true;255 }256 continue;257 }258 // TODO: Handle region side effects properly.259 noSideEffects &= isMemoryEffectFree(&op) && op.getNumRegions() == 0;260 }261 for (MutableArrayRef<ParallelOp> ploops : ploopChains) {262 for (int i = 0, e = ploops.size(); i + 1 < e; ++i)263 fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);264 }265 }266}267 268namespace {269struct ParallelLoopFusion270 : public impl::SCFParallelLoopFusionBase<ParallelLoopFusion> {271 void runOnOperation() override {272 auto &aa = getAnalysis<AliasAnalysis>();273 274 auto mayAlias = [&](Value val1, Value val2) -> bool {275 return !aa.alias(val1, val2).isNo();276 };277 278 getOperation()->walk([&](Operation *child) {279 for (Region ®ion : child->getRegions())280 naivelyFuseParallelOps(region, mayAlias);281 });282 }283};284} // namespace285 286std::unique_ptr<Pass> mlir::createParallelLoopFusionPass() {287 return std::make_unique<ParallelLoopFusion>();288}289