Skip to content

[mlir][bufferization] Generalize returns to be ops with ReturnLike trait #124949

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
cathyzhyi wants to merge 1 commit into from
Closed

[mlir][bufferization] Generalize returns to be ops with ReturnLike trait #124949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion mlir/include/mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ class FuncOp;

namespace bufferization {
/// Helper function that returns all func.return ops in the given function.
SmallVector<func::ReturnOp> getReturnOps(func::FuncOp funcOp);
SmallVector<Operation *> getReturnOps(func::FuncOp funcOp);

namespace func_ext {
/// The state of analysis of a FuncOp.
Expand Down
20 changes: 12 additions & 8 deletions mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,13 @@

namespace mlir {
/// Return all func.return ops in the given function.
SmallVector<func::ReturnOp> bufferization::getReturnOps(func::FuncOp funcOp) {
SmallVector<func::ReturnOp> result;
for (Block &b : funcOp.getBody())
if (auto returnOp = dyn_cast<func::ReturnOp>(b.getTerminator()))
result.push_back(returnOp);
SmallVector<Operation *> bufferization::getReturnOps(func::FuncOp funcOp) {
SmallVector<Operation *> result;
for (Block &b : funcOp.getBody()) {
Operation *terminator = b.getTerminator();
if (terminator->hasTrait<OpTrait::ReturnLike>())
result.push_back(b.getTerminator());
}
return result;
}

Expand Down Expand Up @@ -439,7 +441,7 @@ struct FuncOpInterface
return failure();

// 2. Bufferize the operands of the all return op.
for (func::ReturnOp returnOp : getReturnOps(funcOp)) {
for (Operation *returnOp : getReturnOps(funcOp)) {
assert(returnOp->getNumOperands() == retTypes.size() &&
"incorrect number of return values");
SmallVector<Value> returnValues;
Expand All @@ -457,11 +459,13 @@ struct FuncOpInterface
// Note: If `inferFunctionResultLayout = true`, casts are later folded
// away.
Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
returnOp.getLoc(), bufferizedType, returnVal);
returnOp->getLoc(), bufferizedType, returnVal);
returnValues.push_back(toMemrefOp);
}

returnOp.getOperandsMutable().assign(returnValues);
for (auto [i, operand] : enumerate(returnValues)) {
returnOp->setOperand(i, operand);
}
}

// 3. Set the new function type.
Expand Down
12 changes: 6 additions & 6 deletions mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
}

// Find all func.return ops.
SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
SmallVector<Operation *> returnOps = getReturnOps(funcOp);
assert(!returnOps.empty() && "expected at least one ReturnOp");

// Build alias sets. Merge all aliases from all func.return ops.
Expand All @@ -142,7 +142,7 @@ aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
int64_t bbArgIdx = bbArg.getArgNumber();
// Store aliases in a set, so that we don't add the same alias twice.
SetVector<int64_t> aliases;
for (func::ReturnOp returnOp : returnOps) {
for (Operation *returnOp : returnOps) {
for (OpOperand &returnVal : returnOp->getOpOperands()) {
if (isa<RankedTensorType>(returnVal.get().getType())) {
int64_t returnIdx = returnVal.getOperandNumber();
Expand Down Expand Up @@ -192,7 +192,7 @@ aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
// argument for the i-th operand. In contrast to aliasing information,
// which is just "merged", equivalence information must match across all
// func.return ops.
for (func::ReturnOp returnOp : ArrayRef(returnOps).drop_front()) {
for (Operation *returnOp : ArrayRef(returnOps).drop_front()) {
std::optional<int64_t> maybeEquiv =
findEquivalentBlockArgIdx(returnOp->getOpOperand(i));
if (maybeEquiv != bbArgIdx) {
Expand Down Expand Up @@ -398,7 +398,7 @@ static Value unpackCast(Value v) {
/// func.return ops. This function returns as many types as the return ops have
/// operands. If the i-th operand is not the same for all func.return ops, then
/// the i-th returned type is an "empty" type.
static SmallVector<Type> getReturnTypes(SmallVector<func::ReturnOp> returnOps) {
static SmallVector<Type> getReturnTypes(SmallVector<Operation *> returnOps) {
assert(!returnOps.empty() && "expected at least one ReturnOp");
int numOperands = returnOps.front()->getNumOperands();

Expand Down Expand Up @@ -434,11 +434,11 @@ static void foldMemRefCasts(func::FuncOp funcOp) {
return;

// Compute the common result types of all return ops.
SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
SmallVector<Operation *> returnOps = getReturnOps(funcOp);
SmallVector<Type> resultTypes = getReturnTypes(returnOps);

// Remove direct casts.
for (func::ReturnOp returnOp : returnOps) {
for (Operation *returnOp : returnOps) {
for (OpOperand &operand : returnOp->getOpOperands()) {
// Bail if no common result type was found.
if (resultTypes[operand.getOperandNumber()]) {
Expand Down