159 lines · cpp
1//===- ControlFlowSinkUtils.cpp - Code to perform control-flow sinking ----===//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 utilities for control-flow sinking. Control-flow10// sinking moves operations whose only uses are in conditionally-executed blocks11// into those blocks so that they aren't executed on paths where their results12// are not needed.13//14// Control-flow sinking is not implemented on BranchOpInterface because15// sinking ops into the successors of branch operations may move ops into loops.16// It is idiomatic MLIR to perform optimizations at IR levels that readily17// provide the necessary information.18//19//===----------------------------------------------------------------------===//20 21#include "mlir/Transforms/ControlFlowSinkUtils.h"22#include "mlir/IR/Dominance.h"23#include "mlir/IR/Matchers.h"24#include "mlir/IR/Operation.h"25#include "mlir/IR/OperationSupport.h"26#include "mlir/Interfaces/ControlFlowInterfaces.h"27#include "llvm/Support/DebugLog.h"28#include <vector>29 30#define DEBUG_TYPE "cf-sink"31 32using namespace mlir;33 34namespace {35/// A helper struct for control-flow sinking.36class Sinker {37public:38 /// Create an operation sinker with given dominance info.39 Sinker(function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,40 function_ref<void(Operation *, Region *)> moveIntoRegion,41 DominanceInfo &domInfo)42 : shouldMoveIntoRegion(shouldMoveIntoRegion),43 moveIntoRegion(moveIntoRegion), domInfo(domInfo) {}44 45 /// Given a list of regions, find operations to sink and sink them. Return the46 /// number of operations sunk.47 size_t sinkRegions(RegionRange regions);48 49private:50 /// Given a region and an op which dominates the region, returns true if all51 /// users of the given op are dominated by the entry block of the region, and52 /// thus the operation can be sunk into the region.53 bool allUsersDominatedBy(Operation *op, Region *region);54 55 /// Given a region and a top-level op (an op whose parent region is the given56 /// region), determine whether the defining ops of the op's operands can be57 /// sunk into the region.58 ///59 /// Add moved ops to the work queue.60 void tryToSinkPredecessors(Operation *user, Region *region,61 std::vector<Operation *> &stack);62 63 /// Iterate over all the ops in a region and try to sink their predecessors.64 /// Recurse on subgraphs using a work queue.65 void sinkRegion(Region *region);66 67 /// The callback to determine whether an op should be moved in to a region.68 function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion;69 /// The calback to move an operation into the region.70 function_ref<void(Operation *, Region *)> moveIntoRegion;71 /// Dominance info to determine op user dominance with respect to regions.72 DominanceInfo &domInfo;73 /// The number of operations sunk.74 size_t numSunk = 0;75};76} // end anonymous namespace77 78bool Sinker::allUsersDominatedBy(Operation *op, Region *region) {79 assert(region->findAncestorOpInRegion(*op) == nullptr &&80 "expected op to be defined outside the region");81 return llvm::all_of(op->getUsers(), [&](Operation *user) {82 // The user is dominated by the region if its containing block is dominated83 // by the region's entry block.84 return domInfo.dominates(®ion->front(), user->getBlock());85 });86}87 88void Sinker::tryToSinkPredecessors(Operation *user, Region *region,89 std::vector<Operation *> &stack) {90 LDBG() << "Contained op: "91 << OpWithFlags(user, OpPrintingFlags().skipRegions());92 for (Value value : user->getOperands()) {93 Operation *op = value.getDefiningOp();94 // Ignore block arguments and ops that are already inside the region.95 if (!op || op->getParentRegion() == region)96 continue;97 LDBG() << "Try to sink:\n"98 << OpWithFlags(op, OpPrintingFlags().skipRegions());99 100 // If the op's users are all in the region and it can be moved, then do so.101 if (allUsersDominatedBy(op, region) && shouldMoveIntoRegion(op, region)) {102 moveIntoRegion(op, region);103 ++numSunk;104 // Add the op to the work queue.105 stack.push_back(op);106 }107 }108}109 110void Sinker::sinkRegion(Region *region) {111 // Initialize the work queue with all the ops in the region.112 std::vector<Operation *> stack;113 for (Operation &op : region->getOps())114 stack.push_back(&op);115 116 // Process all the ops depth-first. This ensures that nodes of subgraphs are117 // sunk in the correct order.118 while (!stack.empty()) {119 Operation *op = stack.back();120 stack.pop_back();121 tryToSinkPredecessors(op, region, stack);122 }123}124 125size_t Sinker::sinkRegions(RegionRange regions) {126 for (Region *region : regions)127 if (!region->empty())128 sinkRegion(region);129 return numSunk;130}131 132size_t mlir::controlFlowSink(133 RegionRange regions, DominanceInfo &domInfo,134 function_ref<bool(Operation *, Region *)> shouldMoveIntoRegion,135 function_ref<void(Operation *, Region *)> moveIntoRegion) {136 return Sinker(shouldMoveIntoRegion, moveIntoRegion, domInfo)137 .sinkRegions(regions);138}139 140void mlir::getSinglyExecutedRegionsToSink(RegionBranchOpInterface branch,141 SmallVectorImpl<Region *> ®ions) {142 // Collect constant operands.143 SmallVector<Attribute> operands(branch->getNumOperands(), Attribute());144 for (auto [idx, operand] : llvm::enumerate(branch->getOperands()))145 (void)matchPattern(operand, m_Constant(&operands[idx]));146 147 // Get the invocation bounds.148 SmallVector<InvocationBounds> bounds;149 branch.getRegionInvocationBounds(operands, bounds);150 151 // For a simple control-flow sink, only consider regions that are executed at152 // most once.153 for (auto it : llvm::zip(branch->getRegions(), bounds)) {154 const InvocationBounds &bound = std::get<1>(it);155 if (bound.getUpperBound() && *bound.getUpperBound() <= 1)156 regions.push_back(&std::get<0>(it));157 }158}159