142 lines · cpp
1#include "mlir/Analysis/SliceWalk.h"2#include "mlir/Interfaces/ControlFlowInterfaces.h"3 4using namespace mlir;5 6WalkContinuation mlir::walkSlice(ValueRange rootValues,7 WalkCallback walkCallback) {8 // Search the backward slice starting from the root values.9 SmallVector<Value> workList = rootValues;10 llvm::SmallDenseSet<Value, 16> seenValues;11 while (!workList.empty()) {12 // Search the backward slice of the current value.13 Value current = workList.pop_back_val();14 15 // Skip the current value if it has already been seen.16 if (!seenValues.insert(current).second)17 continue;18 19 // Call the walk callback with the current value.20 WalkContinuation continuation = walkCallback(current);21 if (continuation.wasInterrupted())22 return continuation;23 if (continuation.wasSkipped())24 continue;25 26 assert(continuation.wasAdvancedTo());27 // Add the next values to the work list if the walk should continue.28 workList.append(continuation.getNextValues().begin(),29 continuation.getNextValues().end());30 }31 32 return WalkContinuation::skip();33}34 35/// Returns the operands from all predecessor regions that match `operandNumber`36/// for the `successor` region within `regionOp`.37static SmallVector<Value>38getRegionPredecessorOperands(RegionBranchOpInterface regionOp,39 RegionSuccessor successor,40 unsigned operandNumber) {41 SmallVector<Value> predecessorOperands;42 43 // Returns true if `successors` contains `successor`.44 auto isContained = [](ArrayRef<RegionSuccessor> successors,45 RegionSuccessor successor) {46 auto *it = llvm::find_if(successors, [&successor](RegionSuccessor curr) {47 return curr.getSuccessor() == successor.getSuccessor();48 });49 return it != successors.end();50 };51 52 // Search the operand ranges on the region operation itself.53 SmallVector<Attribute> operandAttributes(regionOp->getNumOperands());54 SmallVector<RegionSuccessor> successors;55 regionOp.getEntrySuccessorRegions(operandAttributes, successors);56 if (isContained(successors, successor)) {57 OperandRange operands = regionOp.getEntrySuccessorOperands(successor);58 predecessorOperands.push_back(operands[operandNumber]);59 }60 61 // Search the operand ranges on region terminators.62 for (Region ®ion : regionOp->getRegions()) {63 for (Block &block : region) {64 auto terminatorOp =65 dyn_cast<RegionBranchTerminatorOpInterface>(block.getTerminator());66 if (!terminatorOp)67 continue;68 SmallVector<Attribute> operandAttributes(terminatorOp->getNumOperands());69 SmallVector<RegionSuccessor> successors;70 terminatorOp.getSuccessorRegions(operandAttributes, successors);71 if (isContained(successors, successor)) {72 OperandRange operands = terminatorOp.getSuccessorOperands(successor);73 predecessorOperands.push_back(operands[operandNumber]);74 }75 }76 }77 78 return predecessorOperands;79}80 81/// Returns the predecessor branch operands that match `blockArg`, or nullopt if82/// some of the predecessor terminators do not implement the BranchOpInterface.83static std::optional<SmallVector<Value>>84getBlockPredecessorOperands(BlockArgument blockArg) {85 Block *block = blockArg.getOwner();86 87 // Search the predecessor operands for all predecessor terminators.88 SmallVector<Value> predecessorOperands;89 for (auto it = block->pred_begin(); it != block->pred_end(); ++it) {90 Block *predecessor = *it;91 auto branchOp = dyn_cast<BranchOpInterface>(predecessor->getTerminator());92 if (!branchOp)93 return std::nullopt;94 SuccessorOperands successorOperands =95 branchOp.getSuccessorOperands(it.getSuccessorIndex());96 // Store the predecessor operand if the block argument matches an operand97 // and is not produced by the terminator.98 if (Value operand = successorOperands[blockArg.getArgNumber()])99 predecessorOperands.push_back(operand);100 }101 102 return predecessorOperands;103}104 105std::optional<SmallVector<Value>>106mlir::getControlFlowPredecessors(Value value) {107 if (OpResult opResult = dyn_cast<OpResult>(value)) {108 if (auto selectOp = opResult.getDefiningOp<SelectLikeOpInterface>())109 return SmallVector<Value>(110 {selectOp.getTrueValue(), selectOp.getFalseValue()});111 auto regionOp = opResult.getDefiningOp<RegionBranchOpInterface>();112 // If the interface is not implemented, there are no control flow113 // predecessors to work with.114 if (!regionOp)115 return std::nullopt;116 // Add the control flow predecessor operands to the work list.117 RegionSuccessor region(regionOp, regionOp->getResults());118 SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(119 regionOp, region, opResult.getResultNumber());120 return predecessorOperands;121 }122 123 auto blockArg = cast<BlockArgument>(value);124 Block *block = blockArg.getOwner();125 // Search the region predecessor operands for structured control flow.126 if (block->isEntryBlock()) {127 if (auto regionBranchOp =128 dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {129 RegionSuccessor region(blockArg.getParentRegion());130 SmallVector<Value> predecessorOperands = getRegionPredecessorOperands(131 regionBranchOp, region, blockArg.getArgNumber());132 return predecessorOperands;133 }134 // If the interface is not implemented, there are no control flow135 // predecessors to work with.136 return std::nullopt;137 }138 139 // Search the block predecessor operands for unstructured control flow.140 return getBlockPredecessorOperands(blockArg);141}142