Mehdi Amini
6e6c296f20
Replace the RegionSuccessor parent sentinel with an explicit Operation * target. Operation successors now model continuing after that operation, while region successors continue to point at Region. Update RegionBranchOpInterface helpers, analyses, and dialect implementations to use operation successors. This removes the ambiguity of interpreting parent relative to whichever operation produced the successor when introducing early-exits. For downstream integration for this API change: - `successor.isParent()` -> `successor.isOperation()` - `RegionSuccessor::parent()` -> RegionSuccessor(op)` ; that means that you need to have access to the target op. In general it is either `getOperation()` or `getParentOp()` (depending if you're updating a method on the `RegionBranchOpInterface` operation or on the terminator. Assisted-by: Codex
102 lines
4.0 KiB
C++
102 lines
4.0 KiB
C++
#include "mlir/Analysis/SliceWalk.h"
|
|
#include "mlir/Interfaces/ControlFlowInterfaces.h"
|
|
|
|
using namespace mlir;
|
|
|
|
WalkContinuation mlir::walkSlice(ValueRange rootValues,
|
|
WalkCallback walkCallback) {
|
|
// Search the backward slice starting from the root values.
|
|
SmallVector<Value> workList = rootValues;
|
|
llvm::SmallDenseSet<Value, 16> seenValues;
|
|
while (!workList.empty()) {
|
|
// Search the backward slice of the current value.
|
|
Value current = workList.pop_back_val();
|
|
|
|
// Skip the current value if it has already been seen.
|
|
if (!seenValues.insert(current).second)
|
|
continue;
|
|
|
|
// Call the walk callback with the current value.
|
|
WalkContinuation continuation = walkCallback(current);
|
|
if (continuation.wasInterrupted())
|
|
return continuation;
|
|
if (continuation.wasSkipped())
|
|
continue;
|
|
|
|
assert(continuation.wasAdvancedTo());
|
|
// Add the next values to the work list if the walk should continue.
|
|
workList.append(continuation.getNextValues().begin(),
|
|
continuation.getNextValues().end());
|
|
}
|
|
|
|
return WalkContinuation::skip();
|
|
}
|
|
|
|
/// Returns the predecessor branch operands that match `blockArg`, or nullopt if
|
|
/// some of the predecessor terminators do not implement the BranchOpInterface.
|
|
static std::optional<SmallVector<Value>>
|
|
getBlockPredecessorOperands(BlockArgument blockArg) {
|
|
Block *block = blockArg.getOwner();
|
|
|
|
// Search the predecessor operands for all predecessor terminators.
|
|
SmallVector<Value> predecessorOperands;
|
|
for (auto it = block->pred_begin(); it != block->pred_end(); ++it) {
|
|
Block *predecessor = *it;
|
|
auto branchOp = dyn_cast<BranchOpInterface>(predecessor->getTerminator());
|
|
if (!branchOp)
|
|
return std::nullopt;
|
|
SuccessorOperands successorOperands =
|
|
branchOp.getSuccessorOperands(it.getSuccessorIndex());
|
|
// Store the predecessor operand if the block argument matches an operand
|
|
// and is not produced by the terminator.
|
|
if (Value operand = successorOperands[blockArg.getArgNumber()])
|
|
predecessorOperands.push_back(operand);
|
|
}
|
|
|
|
return predecessorOperands;
|
|
}
|
|
|
|
std::optional<SmallVector<Value>>
|
|
mlir::getControlFlowPredecessors(Value value) {
|
|
if (OpResult opResult = dyn_cast<OpResult>(value)) {
|
|
if (auto selectOp = opResult.getDefiningOp<SelectLikeOpInterface>())
|
|
return SmallVector<Value>(
|
|
{selectOp.getTrueValue(), selectOp.getFalseValue()});
|
|
auto regionOp = opResult.getDefiningOp<RegionBranchOpInterface>();
|
|
// If the interface is not implemented, there are no control flow
|
|
// predecessors to work with.
|
|
if (!regionOp)
|
|
return std::nullopt;
|
|
// Add the control flow predecessor operands to the work list.
|
|
RegionSuccessor region = RegionSuccessor(regionOp.getOperation());
|
|
SmallVector<Value> predecessorOperands;
|
|
// TODO (#175168): This assumes that there are no non-successor-inputs
|
|
// in front of the op result.
|
|
regionOp.getPredecessorValues(region, opResult.getResultNumber(),
|
|
predecessorOperands);
|
|
return predecessorOperands;
|
|
}
|
|
|
|
auto blockArg = cast<BlockArgument>(value);
|
|
Block *block = blockArg.getOwner();
|
|
// Search the region predecessor operands for structured control flow.
|
|
if (block->isEntryBlock()) {
|
|
if (auto regionBranchOp =
|
|
dyn_cast<RegionBranchOpInterface>(block->getParentOp())) {
|
|
RegionSuccessor region(blockArg.getParentRegion());
|
|
SmallVector<Value> predecessorOperands;
|
|
// TODO (#175168): This assumes that there are no non-successor-inputs
|
|
// in front of the block argument.
|
|
regionBranchOp.getPredecessorValues(region, blockArg.getArgNumber(),
|
|
predecessorOperands);
|
|
return predecessorOperands;
|
|
}
|
|
// If the interface is not implemented, there are no control flow
|
|
// predecessors to work with.
|
|
return std::nullopt;
|
|
}
|
|
|
|
// Search the block predecessor operands for unstructured control flow.
|
|
return getBlockPredecessorOperands(blockArg);
|
|
}
|