1213 lines · cpp
1//===- RegionUtils.cpp - Region-related transformation utilities ----------===//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#include "mlir/Transforms/RegionUtils.h"10 11#include "mlir/Analysis/SliceAnalysis.h"12#include "mlir/Analysis/TopologicalSortUtils.h"13#include "mlir/IR/Block.h"14#include "mlir/IR/Dominance.h"15#include "mlir/IR/IRMapping.h"16#include "mlir/IR/Operation.h"17#include "mlir/IR/PatternMatch.h"18#include "mlir/IR/Value.h"19#include "mlir/Interfaces/ControlFlowInterfaces.h"20#include "mlir/Interfaces/SideEffectInterfaces.h"21#include "mlir/Support/LogicalResult.h"22 23#include "llvm/ADT/DepthFirstIterator.h"24#include "llvm/ADT/PostOrderIterator.h"25#include "llvm/ADT/STLExtras.h"26#include "llvm/Support/DebugLog.h"27 28#include <deque>29#include <iterator>30 31using namespace mlir;32 33#define DEBUG_TYPE "region-utils"34 35void mlir::replaceAllUsesInRegionWith(Value orig, Value replacement,36 Region ®ion) {37 for (auto &use : llvm::make_early_inc_range(orig.getUses())) {38 if (region.isAncestor(use.getOwner()->getParentRegion()))39 use.set(replacement);40 }41}42 43void mlir::visitUsedValuesDefinedAbove(44 Region ®ion, Region &limit, function_ref<void(OpOperand *)> callback) {45 assert(limit.isAncestor(®ion) &&46 "expected isolation limit to be an ancestor of the given region");47 48 // Collect proper ancestors of `limit` upfront to avoid traversing the region49 // tree for every value.50 SmallPtrSet<Region *, 4> properAncestors;51 for (auto *reg = limit.getParentRegion(); reg != nullptr;52 reg = reg->getParentRegion()) {53 properAncestors.insert(reg);54 }55 56 region.walk([callback, &properAncestors](Operation *op) {57 for (OpOperand &operand : op->getOpOperands())58 // Callback on values defined in a proper ancestor of region.59 if (properAncestors.count(operand.get().getParentRegion()))60 callback(&operand);61 });62}63 64void mlir::visitUsedValuesDefinedAbove(65 MutableArrayRef<Region> regions, function_ref<void(OpOperand *)> callback) {66 for (Region ®ion : regions)67 visitUsedValuesDefinedAbove(region, region, callback);68}69 70void mlir::getUsedValuesDefinedAbove(Region ®ion, Region &limit,71 SetVector<Value> &values) {72 visitUsedValuesDefinedAbove(region, limit, [&](OpOperand *operand) {73 values.insert(operand->get());74 });75}76 77void mlir::getUsedValuesDefinedAbove(MutableArrayRef<Region> regions,78 SetVector<Value> &values) {79 for (Region ®ion : regions)80 getUsedValuesDefinedAbove(region, region, values);81}82 83//===----------------------------------------------------------------------===//84// Make block isolated from above.85//===----------------------------------------------------------------------===//86 87SmallVector<Value> mlir::makeRegionIsolatedFromAbove(88 RewriterBase &rewriter, Region ®ion,89 llvm::function_ref<bool(Operation *)> cloneOperationIntoRegion) {90 91 // Get initial list of values used within region but defined above.92 llvm::SetVector<Value> initialCapturedValues;93 mlir::getUsedValuesDefinedAbove(region, initialCapturedValues);94 95 std::deque<Value> worklist(initialCapturedValues.begin(),96 initialCapturedValues.end());97 llvm::DenseSet<Value> visited;98 llvm::DenseSet<Operation *> visitedOps;99 100 llvm::SetVector<Value> finalCapturedValues;101 SmallVector<Operation *> clonedOperations;102 while (!worklist.empty()) {103 Value currValue = worklist.front();104 worklist.pop_front();105 if (visited.count(currValue))106 continue;107 visited.insert(currValue);108 109 Operation *definingOp = currValue.getDefiningOp();110 if (!definingOp || visitedOps.count(definingOp)) {111 finalCapturedValues.insert(currValue);112 continue;113 }114 visitedOps.insert(definingOp);115 116 if (!cloneOperationIntoRegion(definingOp)) {117 // Defining operation isnt cloned, so add the current value to final118 // captured values list.119 finalCapturedValues.insert(currValue);120 continue;121 }122 123 // Add all operands of the operation to the worklist and mark the op as to124 // be cloned.125 for (Value operand : definingOp->getOperands()) {126 if (visited.count(operand))127 continue;128 worklist.push_back(operand);129 }130 clonedOperations.push_back(definingOp);131 }132 133 // The operations to be cloned need to be ordered in topological order134 // so that they can be cloned into the region without violating use-def135 // chains.136 mlir::computeTopologicalSorting(clonedOperations);137 138 OpBuilder::InsertionGuard g(rewriter);139 // Collect types of existing block140 Block *entryBlock = ®ion.front();141 SmallVector<Type> newArgTypes =142 llvm::to_vector(entryBlock->getArgumentTypes());143 SmallVector<Location> newArgLocs = llvm::to_vector(llvm::map_range(144 entryBlock->getArguments(), [](BlockArgument b) { return b.getLoc(); }));145 146 // Append the types of the captured values.147 for (auto value : finalCapturedValues) {148 newArgTypes.push_back(value.getType());149 newArgLocs.push_back(value.getLoc());150 }151 152 // Create a new entry block.153 Block *newEntryBlock =154 rewriter.createBlock(®ion, region.begin(), newArgTypes, newArgLocs);155 auto newEntryBlockArgs = newEntryBlock->getArguments();156 157 // Create a mapping between the captured values and the new arguments added.158 IRMapping map;159 auto replaceIfFn = [&](OpOperand &use) {160 return use.getOwner()->getBlock()->getParent() == ®ion;161 };162 for (auto [arg, capturedVal] :163 llvm::zip(newEntryBlockArgs.take_back(finalCapturedValues.size()),164 finalCapturedValues)) {165 map.map(capturedVal, arg);166 rewriter.replaceUsesWithIf(capturedVal, arg, replaceIfFn);167 }168 rewriter.setInsertionPointToStart(newEntryBlock);169 for (auto *clonedOp : clonedOperations) {170 Operation *newOp = rewriter.clone(*clonedOp, map);171 rewriter.replaceOpUsesWithIf(clonedOp, newOp->getResults(), replaceIfFn);172 }173 rewriter.mergeBlocks(174 entryBlock, newEntryBlock,175 newEntryBlock->getArguments().take_front(entryBlock->getNumArguments()));176 return llvm::to_vector(finalCapturedValues);177}178 179//===----------------------------------------------------------------------===//180// Unreachable Block Elimination181//===----------------------------------------------------------------------===//182 183/// Erase the unreachable blocks within the provided regions. Returns success184/// if any blocks were erased, failure otherwise.185// TODO: We could likely merge this with the DCE algorithm below.186LogicalResult mlir::eraseUnreachableBlocks(RewriterBase &rewriter,187 MutableArrayRef<Region> regions) {188 LDBG() << "Starting eraseUnreachableBlocks with " << regions.size()189 << " regions";190 191 // Set of blocks found to be reachable within a given region.192 llvm::df_iterator_default_set<Block *, 16> reachable;193 // If any blocks were found to be dead.194 int erasedDeadBlocks = 0;195 196 SmallVector<Region *, 1> worklist;197 worklist.reserve(regions.size());198 for (Region ®ion : regions)199 worklist.push_back(®ion);200 201 LDBG(2) << "Initial worklist size: " << worklist.size();202 203 while (!worklist.empty()) {204 Region *region = worklist.pop_back_val();205 if (region->empty()) {206 LDBG(2) << "Skipping empty region";207 continue;208 }209 210 LDBG(2) << "Processing region with " << region->getBlocks().size()211 << " blocks";212 if (region->getParentOp())213 LDBG(2) << " -> for operation: "214 << OpWithFlags(region->getParentOp(),215 OpPrintingFlags().skipRegions());216 217 // If this is a single block region, just collect the nested regions.218 if (region->hasOneBlock()) {219 for (Operation &op : region->front())220 for (Region ®ion : op.getRegions())221 worklist.push_back(®ion);222 continue;223 }224 225 // Mark all reachable blocks.226 reachable.clear();227 for (Block *block : depth_first_ext(®ion->front(), reachable))228 (void)block /* Mark all reachable blocks */;229 230 LDBG(2) << "Found " << reachable.size() << " reachable blocks out of "231 << region->getBlocks().size() << " total blocks";232 233 // Collect all of the dead blocks and push the live regions onto the234 // worklist.235 for (Block &block : llvm::make_early_inc_range(*region)) {236 if (!reachable.count(&block)) {237 LDBG() << "Erasing unreachable block: " << █238 block.dropAllDefinedValueUses();239 rewriter.eraseBlock(&block);240 ++erasedDeadBlocks;241 continue;242 }243 244 // Walk any regions within this block.245 for (Operation &op : block)246 for (Region ®ion : op.getRegions())247 worklist.push_back(®ion);248 }249 }250 251 LDBG() << "Finished eraseUnreachableBlocks, erased " << erasedDeadBlocks252 << " dead blocks";253 254 return success(erasedDeadBlocks > 0);255}256 257//===----------------------------------------------------------------------===//258// Dead Code Elimination259//===----------------------------------------------------------------------===//260 261namespace {262/// Data structure used to track which values have already been proved live.263///264/// Because Operation's can have multiple results, this data structure tracks265/// liveness for both Value's and Operation's to avoid having to look through266/// all Operation results when analyzing a use.267///268/// This data structure essentially tracks the dataflow lattice.269/// The set of values/ops proved live increases monotonically to a fixed-point.270class LiveMap {271public:272 /// Value methods.273 bool wasProvenLive(Value value) {274 // TODO: For results that are removable, e.g. for region based control flow,275 // we could allow for these values to be tracked independently.276 if (OpResult result = dyn_cast<OpResult>(value))277 return wasProvenLive(result.getOwner());278 return wasProvenLive(cast<BlockArgument>(value));279 }280 bool wasProvenLive(BlockArgument arg) { return liveValues.count(arg); }281 void setProvedLive(Value value) {282 // TODO: For results that are removable, e.g. for region based control flow,283 // we could allow for these values to be tracked independently.284 if (OpResult result = dyn_cast<OpResult>(value))285 return setProvedLive(result.getOwner());286 setProvedLive(cast<BlockArgument>(value));287 }288 void setProvedLive(BlockArgument arg) {289 changed |= liveValues.insert(arg).second;290 }291 292 /// Operation methods.293 bool wasProvenLive(Operation *op) { return liveOps.count(op); }294 void setProvedLive(Operation *op) { changed |= liveOps.insert(op).second; }295 296 /// Methods for tracking if we have reached a fixed-point.297 void resetChanged() { changed = false; }298 bool hasChanged() { return changed; }299 300private:301 bool changed = false;302 DenseSet<Value> liveValues;303 DenseSet<Operation *> liveOps;304};305} // namespace306 307static bool isUseSpeciallyKnownDead(OpOperand &use, LiveMap &liveMap) {308 Operation *owner = use.getOwner();309 unsigned operandIndex = use.getOperandNumber();310 // This pass generally treats all uses of an op as live if the op itself is311 // considered live. However, for successor operands to terminators we need a312 // finer-grained notion where we deduce liveness for operands individually.313 // The reason for this is easiest to think about in terms of a classical phi314 // node based SSA IR, where each successor operand is really an operand to a315 // *separate* phi node, rather than all operands to the branch itself as with316 // the block argument representation that MLIR uses.317 //318 // And similarly, because each successor operand is really an operand to a phi319 // node, rather than to the terminator op itself, a terminator op can't e.g.320 // "print" the value of a successor operand.321 if (owner->hasTrait<OpTrait::IsTerminator>()) {322 if (BranchOpInterface branchInterface = dyn_cast<BranchOpInterface>(owner))323 if (auto arg = branchInterface.getSuccessorBlockArgument(operandIndex))324 return !liveMap.wasProvenLive(*arg);325 return false;326 }327 return false;328}329 330static void processValue(Value value, LiveMap &liveMap) {331 bool provedLive = llvm::any_of(value.getUses(), [&](OpOperand &use) {332 if (isUseSpeciallyKnownDead(use, liveMap))333 return false;334 return liveMap.wasProvenLive(use.getOwner());335 });336 if (provedLive)337 liveMap.setProvedLive(value);338}339 340static void propagateLiveness(Region ®ion, LiveMap &liveMap);341 342static void propagateTerminatorLiveness(Operation *op, LiveMap &liveMap) {343 // Terminators are always live.344 liveMap.setProvedLive(op);345 346 // Check to see if we can reason about the successor operands and mutate them.347 BranchOpInterface branchInterface = dyn_cast<BranchOpInterface>(op);348 if (!branchInterface) {349 for (Block *successor : op->getSuccessors())350 for (BlockArgument arg : successor->getArguments())351 liveMap.setProvedLive(arg);352 return;353 }354 355 // If we can't reason about the operand to a successor, conservatively mark356 // it as live.357 for (unsigned i = 0, e = op->getNumSuccessors(); i != e; ++i) {358 SuccessorOperands successorOperands =359 branchInterface.getSuccessorOperands(i);360 for (unsigned opI = 0, opE = successorOperands.getProducedOperandCount();361 opI != opE; ++opI)362 liveMap.setProvedLive(op->getSuccessor(i)->getArgument(opI));363 }364}365 366static void propagateLiveness(Operation *op, LiveMap &liveMap) {367 // Recurse on any regions the op has.368 for (Region ®ion : op->getRegions())369 propagateLiveness(region, liveMap);370 371 // Process terminator operations.372 if (op->hasTrait<OpTrait::IsTerminator>())373 return propagateTerminatorLiveness(op, liveMap);374 375 // Don't reprocess live operations.376 if (liveMap.wasProvenLive(op))377 return;378 379 // Process the op itself.380 if (!wouldOpBeTriviallyDead(op))381 return liveMap.setProvedLive(op);382 383 // If the op isn't intrinsically alive, check it's results.384 for (Value value : op->getResults())385 processValue(value, liveMap);386}387 388static void propagateLiveness(Region ®ion, LiveMap &liveMap) {389 if (region.empty())390 return;391 392 for (Block *block : llvm::post_order(®ion.front())) {393 // We process block arguments after the ops in the block, to promote394 // faster convergence to a fixed point (we try to visit uses before defs).395 for (Operation &op : llvm::reverse(block->getOperations()))396 propagateLiveness(&op, liveMap);397 398 // We currently do not remove entry block arguments, so there is no need to399 // track their liveness.400 // TODO: We could track these and enable removing dead operands/arguments401 // from region control flow operations.402 if (block->isEntryBlock())403 continue;404 405 for (Value value : block->getArguments()) {406 if (!liveMap.wasProvenLive(value))407 processValue(value, liveMap);408 }409 }410}411 412static void eraseTerminatorSuccessorOperands(Operation *terminator,413 LiveMap &liveMap) {414 BranchOpInterface branchOp = dyn_cast<BranchOpInterface>(terminator);415 if (!branchOp)416 return;417 418 for (unsigned succI = 0, succE = terminator->getNumSuccessors();419 succI < succE; succI++) {420 // Iterating successors in reverse is not strictly needed, since we421 // aren't erasing any successors. But it is slightly more efficient422 // since it will promote later operands of the terminator being erased423 // first, reducing the quadratic-ness.424 unsigned succ = succE - succI - 1;425 SuccessorOperands succOperands = branchOp.getSuccessorOperands(succ);426 Block *successor = terminator->getSuccessor(succ);427 428 for (unsigned argI = 0, argE = succOperands.size(); argI < argE; ++argI) {429 // Iterating args in reverse is needed for correctness, to avoid430 // shifting later args when earlier args are erased.431 unsigned arg = argE - argI - 1;432 if (!liveMap.wasProvenLive(successor->getArgument(arg)))433 succOperands.erase(arg);434 }435 }436}437 438static LogicalResult deleteDeadness(RewriterBase &rewriter,439 MutableArrayRef<Region> regions,440 LiveMap &liveMap) {441 bool erasedAnything = false;442 for (Region ®ion : regions) {443 if (region.empty())444 continue;445 bool hasSingleBlock = region.hasOneBlock();446 447 // Delete every operation that is not live. Graph regions may have cycles448 // in the use-def graph, so we must explicitly dropAllUses() from each449 // operation as we erase it. Visiting the operations in post-order450 // guarantees that in SSA CFG regions value uses are removed before defs,451 // which makes dropAllUses() a no-op.452 for (Block *block : llvm::post_order(®ion.front())) {453 if (!hasSingleBlock)454 eraseTerminatorSuccessorOperands(block->getTerminator(), liveMap);455 for (Operation &childOp :456 llvm::make_early_inc_range(llvm::reverse(block->getOperations()))) {457 if (!liveMap.wasProvenLive(&childOp)) {458 erasedAnything = true;459 childOp.dropAllUses();460 rewriter.eraseOp(&childOp);461 } else {462 erasedAnything |= succeeded(463 deleteDeadness(rewriter, childOp.getRegions(), liveMap));464 }465 }466 }467 // Delete block arguments.468 // The entry block has an unknown contract with their enclosing block, so469 // skip it.470 for (Block &block : llvm::drop_begin(region.getBlocks(), 1)) {471 block.eraseArguments(472 [&](BlockArgument arg) { return !liveMap.wasProvenLive(arg); });473 }474 }475 return success(erasedAnything);476}477 478// This function performs a simple dead code elimination algorithm over the479// given regions.480//481// The overall goal is to prove that Values are dead, which allows deleting ops482// and block arguments.483//484// This uses an optimistic algorithm that assumes everything is dead until485// proved otherwise, allowing it to delete recursively dead cycles.486//487// This is a simple fixed-point dataflow analysis algorithm on a lattice488// {Dead,Alive}. Because liveness flows backward, we generally try to489// iterate everything backward to speed up convergence to the fixed-point. This490// allows for being able to delete recursively dead cycles of the use-def graph,491// including block arguments.492//493// This function returns success if any operations or arguments were deleted,494// failure otherwise.495LogicalResult mlir::runRegionDCE(RewriterBase &rewriter,496 MutableArrayRef<Region> regions) {497 LiveMap liveMap;498 do {499 liveMap.resetChanged();500 501 for (Region ®ion : regions)502 propagateLiveness(region, liveMap);503 } while (liveMap.hasChanged());504 505 return deleteDeadness(rewriter, regions, liveMap);506}507 508//===----------------------------------------------------------------------===//509// Block Merging510//===----------------------------------------------------------------------===//511 512//===----------------------------------------------------------------------===//513// BlockEquivalenceData514//===----------------------------------------------------------------------===//515 516namespace {517/// This class contains the information for comparing the equivalencies of two518/// blocks. Blocks are considered equivalent if they contain the same operations519/// in the same order. The only allowed divergence is for operands that come520/// from sources outside of the parent block, i.e. the uses of values produced521/// within the block must be equivalent.522/// e.g.,523/// Equivalent:524/// ^bb1(%arg0: i32)525/// return %arg0, %foo : i32, i32526/// ^bb2(%arg1: i32)527/// return %arg1, %bar : i32, i32528/// Not Equivalent:529/// ^bb1(%arg0: i32)530/// return %foo, %arg0 : i32, i32531/// ^bb2(%arg1: i32)532/// return %arg1, %bar : i32, i32533struct BlockEquivalenceData {534 BlockEquivalenceData(Block *block);535 536 /// Return the order index for the given value that is within the block of537 /// this data.538 unsigned getOrderOf(Value value) const;539 540 /// The block this data refers to.541 Block *block;542 /// A hash value for this block.543 llvm::hash_code hash;544 /// A map of result producing operations to their relative orders within this545 /// block. The order of an operation is the number of defined values that are546 /// produced within the block before this operation.547 DenseMap<Operation *, unsigned> opOrderIndex;548};549} // namespace550 551BlockEquivalenceData::BlockEquivalenceData(Block *block)552 : block(block), hash(0) {553 unsigned orderIt = block->getNumArguments();554 for (Operation &op : *block) {555 if (unsigned numResults = op.getNumResults()) {556 opOrderIndex.try_emplace(&op, orderIt);557 orderIt += numResults;558 }559 auto opHash = OperationEquivalence::computeHash(560 &op, OperationEquivalence::ignoreHashValue,561 OperationEquivalence::ignoreHashValue,562 OperationEquivalence::IgnoreLocations);563 hash = llvm::hash_combine(hash, opHash);564 }565}566 567unsigned BlockEquivalenceData::getOrderOf(Value value) const {568 assert(value.getParentBlock() == block && "expected value of this block");569 570 // Arguments use the argument number as the order index.571 if (BlockArgument arg = dyn_cast<BlockArgument>(value))572 return arg.getArgNumber();573 574 // Otherwise, the result order is offset from the parent op's order.575 OpResult result = cast<OpResult>(value);576 auto opOrderIt = opOrderIndex.find(result.getDefiningOp());577 assert(opOrderIt != opOrderIndex.end() && "expected op to have an order");578 return opOrderIt->second + result.getResultNumber();579}580 581//===----------------------------------------------------------------------===//582// BlockMergeCluster583//===----------------------------------------------------------------------===//584 585namespace {586/// This class represents a cluster of blocks to be merged together.587class BlockMergeCluster {588public:589 BlockMergeCluster(BlockEquivalenceData &&leaderData)590 : leaderData(std::move(leaderData)) {}591 592 /// Attempt to add the given block to this cluster. Returns success if the593 /// block was merged, failure otherwise.594 LogicalResult addToCluster(BlockEquivalenceData &blockData);595 596 /// Try to merge all of the blocks within this cluster into the leader block.597 LogicalResult merge(RewriterBase &rewriter);598 599private:600 /// The equivalence data for the leader of the cluster.601 BlockEquivalenceData leaderData;602 603 /// The set of blocks that can be merged into the leader.604 llvm::SmallSetVector<Block *, 1> blocksToMerge;605 606 /// A set of operand+index pairs that correspond to operands that need to be607 /// replaced by arguments when the cluster gets merged.608 std::set<std::pair<int, int>> operandsToMerge;609};610} // namespace611 612LogicalResult BlockMergeCluster::addToCluster(BlockEquivalenceData &blockData) {613 if (leaderData.hash != blockData.hash)614 return failure();615 Block *leaderBlock = leaderData.block, *mergeBlock = blockData.block;616 if (leaderBlock->getArgumentTypes() != mergeBlock->getArgumentTypes())617 return failure();618 619 // A set of operands that mismatch between the leader and the new block.620 SmallVector<std::pair<int, int>, 8> mismatchedOperands;621 auto lhsIt = leaderBlock->begin(), lhsE = leaderBlock->end();622 auto rhsIt = blockData.block->begin(), rhsE = blockData.block->end();623 for (int opI = 0; lhsIt != lhsE && rhsIt != rhsE; ++lhsIt, ++rhsIt, ++opI) {624 // Check that the operations are equivalent.625 if (!OperationEquivalence::isEquivalentTo(626 &*lhsIt, &*rhsIt, OperationEquivalence::ignoreValueEquivalence,627 /*markEquivalent=*/nullptr,628 OperationEquivalence::Flags::IgnoreLocations))629 return failure();630 631 // Compare the operands of the two operations. If the operand is within632 // the block, it must refer to the same operation.633 auto lhsOperands = lhsIt->getOperands(), rhsOperands = rhsIt->getOperands();634 for (int operand : llvm::seq<int>(0, lhsIt->getNumOperands())) {635 Value lhsOperand = lhsOperands[operand];636 Value rhsOperand = rhsOperands[operand];637 if (lhsOperand == rhsOperand)638 continue;639 // Check that the types of the operands match.640 if (lhsOperand.getType() != rhsOperand.getType())641 return failure();642 643 // Check that these uses are both external, or both internal.644 bool lhsIsInBlock = lhsOperand.getParentBlock() == leaderBlock;645 bool rhsIsInBlock = rhsOperand.getParentBlock() == mergeBlock;646 if (lhsIsInBlock != rhsIsInBlock)647 return failure();648 // Let the operands differ if they are defined in a different block. These649 // will become new arguments if the blocks get merged.650 if (!lhsIsInBlock) {651 652 // Check whether the operands aren't the result of an immediate653 // predecessors terminator. In that case we are not able to use it as a654 // successor operand when branching to the merged block as it does not655 // dominate its producing operation.656 auto isValidSuccessorArg = [](Block *block, Value operand) {657 if (operand.getDefiningOp() !=658 operand.getParentBlock()->getTerminator())659 return true;660 return !llvm::is_contained(block->getPredecessors(),661 operand.getParentBlock());662 };663 664 if (!isValidSuccessorArg(leaderBlock, lhsOperand) ||665 !isValidSuccessorArg(mergeBlock, rhsOperand))666 return failure();667 668 mismatchedOperands.emplace_back(opI, operand);669 continue;670 }671 672 // Otherwise, these operands must have the same logical order within the673 // parent block.674 if (leaderData.getOrderOf(lhsOperand) != blockData.getOrderOf(rhsOperand))675 return failure();676 }677 678 // If the lhs or rhs has external uses, the blocks cannot be merged as the679 // merged version of this operation will not be either the lhs or rhs680 // alone (thus semantically incorrect), but some mix dependending on which681 // block preceeded this.682 // TODO allow merging of operations when one block does not dominate the683 // other684 if (rhsIt->isUsedOutsideOfBlock(mergeBlock) ||685 lhsIt->isUsedOutsideOfBlock(leaderBlock)) {686 return failure();687 }688 }689 // Make sure that the block sizes are equivalent.690 if (lhsIt != lhsE || rhsIt != rhsE)691 return failure();692 693 // If we get here, the blocks are equivalent and can be merged.694 operandsToMerge.insert(mismatchedOperands.begin(), mismatchedOperands.end());695 blocksToMerge.insert(blockData.block);696 return success();697}698 699/// Returns true if the predecessor terminators of the given block can not have700/// their operands updated.701static bool ableToUpdatePredOperands(Block *block) {702 for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) {703 if (!isa<BranchOpInterface>((*it)->getTerminator()))704 return false;705 }706 return true;707}708 709/// Prunes the redundant list of new arguments. E.g., if we are passing an710/// argument list like [x, y, z, x] this would return [x, y, z] and it would711/// update the `block` (to whom the argument are passed to) accordingly. The new712/// arguments are passed as arguments at the back of the block, hence we need to713/// know how many `numOldArguments` were before, in order to correctly replace714/// the new arguments in the block715static SmallVector<SmallVector<Value, 8>, 2> pruneRedundantArguments(716 const SmallVector<SmallVector<Value, 8>, 2> &newArguments,717 RewriterBase &rewriter, unsigned numOldArguments, Block *block) {718 719 SmallVector<SmallVector<Value, 8>, 2> newArgumentsPruned(720 newArguments.size(), SmallVector<Value, 8>());721 722 if (newArguments.empty())723 return newArguments;724 725 // `newArguments` is a 2D array of size `numLists` x `numArgs`726 unsigned numLists = newArguments.size();727 unsigned numArgs = newArguments[0].size();728 729 // Map that for each arg index contains the index that we can use in place of730 // the original index. E.g., if we have newArgs = [x, y, z, x], we will have731 // idxToReplacement[3] = 0732 llvm::DenseMap<unsigned, unsigned> idxToReplacement;733 734 // This is a useful data structure to track the first appearance of a Value735 // on a given list of arguments736 DenseMap<Value, unsigned> firstValueToIdx;737 for (unsigned j = 0; j < numArgs; ++j) {738 Value newArg = newArguments[0][j];739 firstValueToIdx.try_emplace(newArg, j);740 }741 742 // Go through the first list of arguments (list 0).743 for (unsigned j = 0; j < numArgs; ++j) {744 // Look back to see if there are possible redundancies in list 0. Please745 // note that we are using a map to annotate when an argument was seen first746 // to avoid a O(N^2) algorithm. This has the drawback that if we have two747 // lists like:748 // list0: [%a, %a, %a]749 // list1: [%c, %b, %b]750 // We cannot simplify it, because firstValueToIdx[%a] = 0, but we cannot751 // point list1[1](==%b) or list1[2](==%b) to list1[0](==%c). However, since752 // the number of arguments can be potentially unbounded we cannot afford a753 // O(N^2) algorithm (to search to all the possible pairs) and we need to754 // accept the trade-off.755 unsigned k = firstValueToIdx[newArguments[0][j]];756 if (k == j)757 continue;758 759 bool shouldReplaceJ = true;760 unsigned replacement = k;761 // If a possible redundancy is found, then scan the other lists: we762 // can prune the arguments if and only if they are redundant in every763 // list.764 for (unsigned i = 1; i < numLists; ++i)765 shouldReplaceJ =766 shouldReplaceJ && (newArguments[i][k] == newArguments[i][j]);767 // Save the replacement.768 if (shouldReplaceJ)769 idxToReplacement[j] = replacement;770 }771 772 // Populate the pruned argument list.773 for (unsigned i = 0; i < numLists; ++i)774 for (unsigned j = 0; j < numArgs; ++j)775 if (!idxToReplacement.contains(j))776 newArgumentsPruned[i].push_back(newArguments[i][j]);777 778 // Replace the block's redundant arguments.779 SmallVector<unsigned> toErase;780 for (auto [idx, arg] : llvm::enumerate(block->getArguments())) {781 if (idxToReplacement.contains(idx)) {782 Value oldArg = block->getArgument(numOldArguments + idx);783 Value newArg =784 block->getArgument(numOldArguments + idxToReplacement[idx]);785 rewriter.replaceAllUsesWith(oldArg, newArg);786 toErase.push_back(numOldArguments + idx);787 }788 }789 790 // Erase the block's redundant arguments.791 for (unsigned idxToErase : llvm::reverse(toErase))792 block->eraseArgument(idxToErase);793 return newArgumentsPruned;794}795 796LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {797 // Don't consider clusters that don't have blocks to merge.798 if (blocksToMerge.empty())799 return failure();800 801 Block *leaderBlock = leaderData.block;802 if (!operandsToMerge.empty()) {803 // If the cluster has operands to merge, verify that the predecessor804 // terminators of each of the blocks can have their successor operands805 // updated.806 // TODO: We could try and sub-partition this cluster if only some blocks807 // cause the mismatch.808 if (!ableToUpdatePredOperands(leaderBlock) ||809 !llvm::all_of(blocksToMerge, ableToUpdatePredOperands))810 return failure();811 812 // Collect the iterators for each of the blocks to merge. We will walk all813 // of the iterators at once to avoid operand index invalidation.814 SmallVector<Block::iterator, 2> blockIterators;815 blockIterators.reserve(blocksToMerge.size() + 1);816 blockIterators.push_back(leaderBlock->begin());817 for (Block *mergeBlock : blocksToMerge)818 blockIterators.push_back(mergeBlock->begin());819 820 // Update each of the predecessor terminators with the new arguments.821 SmallVector<SmallVector<Value, 8>, 2> newArguments(822 1 + blocksToMerge.size(),823 SmallVector<Value, 8>(operandsToMerge.size()));824 unsigned curOpIndex = 0;825 unsigned numOldArguments = leaderBlock->getNumArguments();826 for (const auto &it : llvm::enumerate(operandsToMerge)) {827 unsigned nextOpOffset = it.value().first - curOpIndex;828 curOpIndex = it.value().first;829 830 // Process the operand for each of the block iterators.831 for (unsigned i = 0, e = blockIterators.size(); i != e; ++i) {832 Block::iterator &blockIter = blockIterators[i];833 std::advance(blockIter, nextOpOffset);834 auto &operand = blockIter->getOpOperand(it.value().second);835 newArguments[i][it.index()] = operand.get();836 837 // Update the operand and insert an argument if this is the leader.838 if (i == 0) {839 Value operandVal = operand.get();840 operand.set(leaderBlock->addArgument(operandVal.getType(),841 operandVal.getLoc()));842 }843 }844 }845 846 // Prune redundant arguments and update the leader block argument list847 newArguments = pruneRedundantArguments(newArguments, rewriter,848 numOldArguments, leaderBlock);849 850 // Update the predecessors for each of the blocks.851 auto updatePredecessors = [&](Block *block, unsigned clusterIndex) {852 for (auto predIt = block->pred_begin(), predE = block->pred_end();853 predIt != predE; ++predIt) {854 auto branch = cast<BranchOpInterface>((*predIt)->getTerminator());855 unsigned succIndex = predIt.getSuccessorIndex();856 branch.getSuccessorOperands(succIndex).append(857 newArguments[clusterIndex]);858 }859 };860 updatePredecessors(leaderBlock, /*clusterIndex=*/0);861 for (unsigned i = 0, e = blocksToMerge.size(); i != e; ++i)862 updatePredecessors(blocksToMerge[i], /*clusterIndex=*/i + 1);863 }864 865 // Replace all uses of the merged blocks with the leader and erase them.866 for (Block *block : blocksToMerge) {867 block->replaceAllUsesWith(leaderBlock);868 rewriter.eraseBlock(block);869 }870 return success();871}872 873/// Identify identical blocks within the given region and merge them, inserting874/// new block arguments as necessary. Returns success if any blocks were merged,875/// failure otherwise.876static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,877 Region ®ion) {878 if (region.empty() || region.hasOneBlock())879 return failure();880 881 // Identify sets of blocks, other than the entry block, that branch to the882 // same successors. We will use these groups to create clusters of equivalent883 // blocks.884 DenseMap<SuccessorRange, SmallVector<Block *, 1>> matchingSuccessors;885 for (Block &block : llvm::drop_begin(region, 1))886 matchingSuccessors[block.getSuccessors()].push_back(&block);887 888 bool mergedAnyBlocks = false;889 for (ArrayRef<Block *> blocks : llvm::make_second_range(matchingSuccessors)) {890 if (blocks.size() == 1)891 continue;892 893 SmallVector<BlockMergeCluster, 1> clusters;894 for (Block *block : blocks) {895 BlockEquivalenceData data(block);896 897 // Don't allow merging if this block has any regions.898 // TODO: Add support for regions if necessary.899 bool hasNonEmptyRegion = llvm::any_of(*block, [](Operation &op) {900 return llvm::any_of(op.getRegions(),901 [](Region ®ion) { return !region.empty(); });902 });903 if (hasNonEmptyRegion)904 continue;905 906 // Don't allow merging if this block's arguments are used outside of the907 // original block.908 bool argHasExternalUsers = llvm::any_of(909 block->getArguments(), [block](mlir::BlockArgument &arg) {910 return arg.isUsedOutsideOfBlock(block);911 });912 if (argHasExternalUsers)913 continue;914 915 // Try to add this block to an existing cluster.916 bool addedToCluster = false;917 for (auto &cluster : clusters)918 if ((addedToCluster = succeeded(cluster.addToCluster(data))))919 break;920 if (!addedToCluster)921 clusters.emplace_back(std::move(data));922 }923 for (auto &cluster : clusters)924 mergedAnyBlocks |= succeeded(cluster.merge(rewriter));925 }926 927 return success(mergedAnyBlocks);928}929 930/// Identify identical blocks within the given regions and merge them, inserting931/// new block arguments as necessary.932static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,933 MutableArrayRef<Region> regions) {934 llvm::SmallSetVector<Region *, 1> worklist;935 for (auto ®ion : regions)936 worklist.insert(®ion);937 bool anyChanged = false;938 while (!worklist.empty()) {939 Region *region = worklist.pop_back_val();940 if (succeeded(mergeIdenticalBlocks(rewriter, *region))) {941 worklist.insert(region);942 anyChanged = true;943 }944 945 // Add any nested regions to the worklist.946 for (Block &block : *region)947 for (auto &op : block)948 for (auto &nestedRegion : op.getRegions())949 worklist.insert(&nestedRegion);950 }951 952 return success(anyChanged);953}954 955/// If a block's argument is always the same across different invocations, then956/// drop the argument and use the value directly inside the block957static LogicalResult dropRedundantArguments(RewriterBase &rewriter,958 Block &block) {959 SmallVector<size_t> argsToErase;960 961 // Go through the arguments of the block.962 for (auto [argIdx, blockOperand] : llvm::enumerate(block.getArguments())) {963 bool sameArg = true;964 Value commonValue;965 966 // Go through the block predecessor and flag if they pass to the block967 // different values for the same argument.968 for (Block::pred_iterator predIt = block.pred_begin(),969 predE = block.pred_end();970 predIt != predE; ++predIt) {971 auto branch = dyn_cast<BranchOpInterface>((*predIt)->getTerminator());972 if (!branch) {973 sameArg = false;974 break;975 }976 unsigned succIndex = predIt.getSuccessorIndex();977 SuccessorOperands succOperands = branch.getSuccessorOperands(succIndex);978 auto branchOperands = succOperands.getForwardedOperands();979 if (!commonValue) {980 commonValue = branchOperands[argIdx];981 continue;982 }983 if (branchOperands[argIdx] != commonValue) {984 sameArg = false;985 break;986 }987 }988 989 // If they are passing the same value, drop the argument.990 if (commonValue && sameArg) {991 argsToErase.push_back(argIdx);992 993 // Remove the argument from the block.994 rewriter.replaceAllUsesWith(blockOperand, commonValue);995 }996 }997 998 // Remove the arguments.999 for (size_t argIdx : llvm::reverse(argsToErase)) {1000 block.eraseArgument(argIdx);1001 1002 // Remove the argument from the branch ops.1003 for (auto predIt = block.pred_begin(), predE = block.pred_end();1004 predIt != predE; ++predIt) {1005 auto branch = cast<BranchOpInterface>((*predIt)->getTerminator());1006 unsigned succIndex = predIt.getSuccessorIndex();1007 SuccessorOperands succOperands = branch.getSuccessorOperands(succIndex);1008 succOperands.erase(argIdx);1009 }1010 }1011 return success(!argsToErase.empty());1012}1013 1014/// This optimization drops redundant argument to blocks. I.e., if a given1015/// argument to a block receives the same value from each of the block1016/// predecessors, we can remove the argument from the block and use directly the1017/// original value. This is a simple example:1018///1019/// %cond = llvm.call @rand() : () -> i11020/// %val0 = llvm.mlir.constant(1 : i64) : i641021/// %val1 = llvm.mlir.constant(2 : i64) : i641022/// %val2 = llvm.mlir.constant(3 : i64) : i641023/// llvm.cond_br %cond, ^bb1(%val0 : i64, %val1 : i64), ^bb2(%val0 : i64, %val21024/// : i64)1025///1026/// ^bb1(%arg0 : i64, %arg1 : i64):1027/// llvm.call @foo(%arg0, %arg1)1028///1029/// The previous IR can be rewritten as:1030/// %cond = llvm.call @rand() : () -> i11031/// %val0 = llvm.mlir.constant(1 : i64) : i641032/// %val1 = llvm.mlir.constant(2 : i64) : i641033/// %val2 = llvm.mlir.constant(3 : i64) : i641034/// llvm.cond_br %cond, ^bb1(%val1 : i64), ^bb2(%val2 : i64)1035///1036/// ^bb1(%arg0 : i64):1037/// llvm.call @foo(%val0, %arg0)1038///1039static LogicalResult dropRedundantArguments(RewriterBase &rewriter,1040 MutableArrayRef<Region> regions) {1041 llvm::SmallSetVector<Region *, 1> worklist;1042 for (Region ®ion : regions)1043 worklist.insert(®ion);1044 bool anyChanged = false;1045 while (!worklist.empty()) {1046 Region *region = worklist.pop_back_val();1047 1048 // Add any nested regions to the worklist.1049 for (Block &block : *region) {1050 anyChanged =1051 succeeded(dropRedundantArguments(rewriter, block)) || anyChanged;1052 1053 for (Operation &op : block)1054 for (Region &nestedRegion : op.getRegions())1055 worklist.insert(&nestedRegion);1056 }1057 }1058 return success(anyChanged);1059}1060 1061//===----------------------------------------------------------------------===//1062// Region Simplification1063//===----------------------------------------------------------------------===//1064 1065/// Run a set of structural simplifications over the given regions. This1066/// includes transformations like unreachable block elimination, dead argument1067/// elimination, as well as some other DCE. This function returns success if any1068/// of the regions were simplified, failure otherwise.1069LogicalResult mlir::simplifyRegions(RewriterBase &rewriter,1070 MutableArrayRef<Region> regions,1071 bool mergeBlocks) {1072 bool eliminatedBlocks = succeeded(eraseUnreachableBlocks(rewriter, regions));1073 bool eliminatedOpsOrArgs = succeeded(runRegionDCE(rewriter, regions));1074 bool mergedIdenticalBlocks = false;1075 bool droppedRedundantArguments = false;1076 if (mergeBlocks) {1077 mergedIdenticalBlocks = succeeded(mergeIdenticalBlocks(rewriter, regions));1078 droppedRedundantArguments =1079 succeeded(dropRedundantArguments(rewriter, regions));1080 }1081 return success(eliminatedBlocks || eliminatedOpsOrArgs ||1082 mergedIdenticalBlocks || droppedRedundantArguments);1083}1084 1085//===---------------------------------------------------------------------===//1086// Move operation dependencies1087//===---------------------------------------------------------------------===//1088 1089LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,1090 Operation *op,1091 Operation *insertionPoint,1092 DominanceInfo &dominance) {1093 // Currently unsupported case where the op and insertion point are1094 // in different basic blocks.1095 if (op->getBlock() != insertionPoint->getBlock()) {1096 return rewriter.notifyMatchFailure(1097 op, "unsupported case where operation and insertion point are not in "1098 "the same basic block");1099 }1100 // If `insertionPoint` does not dominate `op`, do nothing1101 if (!dominance.properlyDominates(insertionPoint, op)) {1102 return rewriter.notifyMatchFailure(op,1103 "insertion point does not dominate op");1104 }1105 1106 // Find the backward slice of operation for each `Value` the operation1107 // depends on. Prune the slice to only include operations not already1108 // dominated by the `insertionPoint`1109 BackwardSliceOptions options;1110 options.inclusive = false;1111 options.omitUsesFromAbove = false;1112 // Since current support is to only move within a same basic block,1113 // the slices dont need to look past block arguments.1114 options.omitBlockArguments = true;1115 options.filter = [&](Operation *sliceBoundaryOp) {1116 return !dominance.properlyDominates(sliceBoundaryOp, insertionPoint);1117 };1118 llvm::SetVector<Operation *> slice;1119 LogicalResult result = getBackwardSlice(op, &slice, options);1120 assert(result.succeeded() && "expected a backward slice");1121 (void)result;1122 1123 // If the slice contains `insertionPoint` cannot move the dependencies.1124 if (slice.contains(insertionPoint)) {1125 return rewriter.notifyMatchFailure(1126 op,1127 "cannot move dependencies before operation in backward slice of op");1128 }1129 1130 // We should move the slice in topological order, but `getBackwardSlice`1131 // already does that. So no need to sort again.1132 for (Operation *op : slice) {1133 rewriter.moveOpBefore(op, insertionPoint);1134 }1135 return success();1136}1137 1138LogicalResult mlir::moveOperationDependencies(RewriterBase &rewriter,1139 Operation *op,1140 Operation *insertionPoint) {1141 DominanceInfo dominance(op);1142 return moveOperationDependencies(rewriter, op, insertionPoint, dominance);1143}1144 1145LogicalResult mlir::moveValueDefinitions(RewriterBase &rewriter,1146 ValueRange values,1147 Operation *insertionPoint,1148 DominanceInfo &dominance) {1149 // Remove the values that already dominate the insertion point.1150 SmallVector<Value> prunedValues;1151 for (auto value : values) {1152 if (dominance.properlyDominates(value, insertionPoint)) {1153 continue;1154 }1155 // Block arguments are not supported.1156 if (isa<BlockArgument>(value)) {1157 return rewriter.notifyMatchFailure(1158 insertionPoint,1159 "unsupported case of moving block argument before insertion point");1160 }1161 // Check for currently unsupported case if the insertion point is in a1162 // different block.1163 if (value.getDefiningOp()->getBlock() != insertionPoint->getBlock()) {1164 return rewriter.notifyMatchFailure(1165 insertionPoint,1166 "unsupported case of moving definition of value before an insertion "1167 "point in a different basic block");1168 }1169 prunedValues.push_back(value);1170 }1171 1172 // Find the backward slice of operation for each `Value` the operation1173 // depends on. Prune the slice to only include operations not already1174 // dominated by the `insertionPoint`1175 BackwardSliceOptions options;1176 options.inclusive = true;1177 options.omitUsesFromAbove = false;1178 // Since current support is to only move within a same basic block,1179 // the slices dont need to look past block arguments.1180 options.omitBlockArguments = true;1181 options.filter = [&](Operation *sliceBoundaryOp) {1182 return !dominance.properlyDominates(sliceBoundaryOp, insertionPoint);1183 };1184 llvm::SetVector<Operation *> slice;1185 for (auto value : prunedValues) {1186 LogicalResult result = getBackwardSlice(value, &slice, options);1187 assert(result.succeeded() && "expected a backward slice");1188 (void)result;1189 }1190 1191 // If the slice contains `insertionPoint` cannot move the dependencies.1192 if (slice.contains(insertionPoint)) {1193 return rewriter.notifyMatchFailure(1194 insertionPoint,1195 "cannot move dependencies before operation in backward slice of op");1196 }1197 1198 // Sort operations topologically before moving.1199 mlir::topologicalSort(slice);1200 1201 for (Operation *op : slice) {1202 rewriter.moveOpBefore(op, insertionPoint);1203 }1204 return success();1205}1206 1207LogicalResult mlir::moveValueDefinitions(RewriterBase &rewriter,1208 ValueRange values,1209 Operation *insertionPoint) {1210 DominanceInfo dominance(insertionPoint);1211 return moveValueDefinitions(rewriter, values, insertionPoint, dominance);1212}1213