Skip to content

[mlir][transform] Overhaul RegionBranchOpInterface implementations. #111408

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[mlir][transform] Overhaul RegionBranchOpInterface implementations. #111408

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e7be852
[mlir][transform] Make `yield` a `ReturnLike` op.
ingomueller-net Oct 7, 2024
d626f95
Overhaul `RegionBranchOpInterface` impls of ops using `yield`.
ingomueller-net Oct 10, 2024
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
3 changes: 2 additions & 1 deletion mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1358,7 +1358,8 @@ def VerifyOp : TransformDialectOp<"verify",
}

def YieldOp : TransformDialectOp<"yield",
[Terminator, DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
[Terminator, ReturnLike,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
let summary = "Yields operation handles from a transform IR region";
let description = [{
This terminator operation yields operation handles from regions of the
Expand Down
34 changes: 5 additions & 29 deletions mlir/lib/Dialect/Transform/IR/TransformOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,16 +104,8 @@ transform::AlternativesOp::getEntrySuccessorOperands(RegionBranchPoint point) {

void transform::AlternativesOp::getSuccessorRegions(
RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
for (Region &alternative : llvm::drop_begin(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing this does not seem right. When a region fails, the op jumps to the next region.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In particular and unlike scf.for, transform.foreach does not pass control from one iteration of its body to the next directly; it rather passes the control back to the parent op, which then passes it back to the body for the next iteration. That can be seen by the fact that the body always gets arguments of the same type as the operands of the parent op (and none of the yielded types) and the types that are yielded correspond exactly to the result types of the parent op.

I don't follow this part. Where is the type mismatch here?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I guess you are saying that the successor operands do not match with the block arguments of the next region.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the RegionBranchOpInterface does not support the use case here. When you branch back to the parent op, there is no way to go back into the op a second time.

    A "region successor" indicates the target of a branch. It can indicate
    either a region of this op or this op. In the former case, the region
    successor is a region pointer and a range of block arguments to which the
    "successor operands" are forwarded to. In the latter case, the control flow
    leaves this op and the region successor is a range of results of this op to
    which the successor operands are forwarded to.

This sentence talks about results of the op, so if we would allow going back into the op, you could set the same results multiple times.

    In the latter case, the control flow
    leaves this op and the region successor is a range of results of this op to
    which the successor operands are forwarded to.

I think we should clarify this in the documentation of RegionBranchOpInterface. (We could also allow it but most transformations on top of RegionBranchOpInterface probably assume that you cannot go back into the op.)

Do we have to implement the RegionBranchOpInterface on these transform ops at all?

Copy link
Contributor Author

@ingomueller-net ingomueller-net Oct 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In sequence, I think the sentence matches, no? When control reaches yield, the operands of yield are forwarded to the parent op, they become the result of the op, and "control flow leaves [the] op," right?

What doesn't quite fit for foreach and alternatives is the "control flow leaves [the] op" part. One could argue that the "region successor is a [part of or a potential] range of results of [the] op," but then the control flow does not leave the op; instead it may go to the next region or re-enter the body. Plus there are rules about how to combine the different results yielded by different regions/iterations (though the types always match).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In practice, what's the difference between "no arguments are passed" and "it isn't specified which arguments are passed"? The former does not imply that the region must have 0 arguments.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe that "it isn't specified which arguments are passed" cannot be expressed currently. An empty argument list is always interpreted as "no arguments are passed."

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that's correct. But I am wondering if it makes any difference for analyses and transformations.

If we a region has an argument, but the RegionBranchOpInterface says that no arguments are forwarded to that region: what does that mean? It means that we have no idea where the data for that block argument is coming from. Maybe the op itself produces it.

Or maybe the value was actually forwarded from another region but we did not account for it in the RegionBranchOpInterface (and the terminator interface).

Does it matter which one is the case for an analysis that checks the RegionBranchOpInterface? It has to be conservative around such cases anyway.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we a region has an argument, but the RegionBranchOpInterface says that no arguments are forwarded to that region: what does that mean?

Are you sure that that can even exist? I think I ran into failed type check that seem to check exactly that while working on the last commit...

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think scf.forall is an example.

getAlternatives(),
point.isParent() ? 0
: point.getRegionOrNull()->getRegionNumber() + 1)) {
regions.emplace_back(&alternative, !getOperands().empty()
? alternative.getArguments()
: Block::BlockArgListType());
}
if (!point.isParent())
regions.emplace_back(getOperation()->getResults());
regions.emplace_back(getResults());
}

void transform::AlternativesOp::getRegionInvocationBounds(
Expand Down Expand Up @@ -1502,16 +1494,8 @@ void transform::ForeachOp::getEffects(

void transform::ForeachOp::getSuccessorRegions(
RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
Region *bodyRegion = &getBody();
if (point.isParent()) {
regions.emplace_back(bodyRegion, bodyRegion->getArguments());
return;
}

// Branch back to the region or the parent.
assert(point == getBody() && "unexpected region index");
regions.emplace_back(bodyRegion, bodyRegion->getArguments());
regions.emplace_back();
if (point.getRegionOrNull() == &getBody())
regions.emplace_back(getResults());
}

OperandRange
Expand Down Expand Up @@ -2702,16 +2686,8 @@ transform::SequenceOp::getEntrySuccessorOperands(RegionBranchPoint point) {

void transform::SequenceOp::getSuccessorRegions(
RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
if (point.isParent()) {
Region *bodyRegion = &getBody();
regions.emplace_back(bodyRegion, getNumOperands() != 0
? bodyRegion->getArguments()
: Block::BlockArgListType());
return;
}

assert(point == getBody() && "unexpected region index");
regions.emplace_back(getOperation()->getResults());
if (point.getRegionOrNull() == &getBody())
regions.emplace_back(getResults());
}

void transform::SequenceOp::getRegionInvocationBounds(
Expand Down
Loading