From e7be852053771fb1c7fa84f4105b675e42201d7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ingo=20M=C3=BCller?= Date: Mon, 7 Oct 2024 17:13:20 +0000 Subject: [PATCH 1/2] [mlir][transform] Make `yield` a `ReturnLike` op. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds the `ReturnLike` trait to `transform.yield`. This is required in the one-shot bufferization pass since the merging of #110332, which analyses any `FunctionOpInterface` and expects them to have a `ReturnLike` terminator. Signed-off-by: Ingo Müller --- mlir/include/mlir/Dialect/Transform/IR/TransformOps.td | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td index b946fc8875860..d3933cad920a3 100644 --- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td +++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td @@ -1358,7 +1358,8 @@ def VerifyOp : TransformDialectOp<"verify", } def YieldOp : TransformDialectOp<"yield", - [Terminator, DeclareOpInterfaceMethods]> { + [Terminator, ReturnLike, + DeclareOpInterfaceMethods]> { let summary = "Yields operation handles from a transform IR region"; let description = [{ This terminator operation yields operation handles from regions of the From d626f95564da4f82313001405437f8cbff30511e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ingo=20M=C3=BCller?= Date: Thu, 10 Oct 2024 08:34:04 +0000 Subject: [PATCH 2/2] Overhaul `RegionBranchOpInterface` impls of ops using `yield`. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit If we give `transform.yield` the `ReturnLike` trait, then the checks made by `RegionBranchOpInterface` behave differently. More concretely, the check that the operands and results passed across the control flow edges from the parent op to its regions and back are equal (or compatible). Which results are passed back from regions to the parent op depend on the terminator and/or whether it is `ReturnLike`. This commit radically changes the implementations of the `alternatives`, `foreach`, and `sequence` ops, which are the ops that use `yield` as their terminator. In fact, all of these ops only ever pass control from the parent op to the region and from there back to the parent op. 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. It is unclear why the previous implementation worked. Clearly, the type checks have not been executed because the terminator was not `ReturnLike`, so the mismatching types passed unnoticed. I suppose that there simply is/was no other use case of using the `RegionBranchOpInterface` of the affected ops yet. Signed-off-by: Ingo Müller --- .../lib/Dialect/Transform/IR/TransformOps.cpp | 34 +++---------------- 1 file changed, 5 insertions(+), 29 deletions(-) diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp index 590cae9aa0d66..48b25d19d7dc3 100644 --- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp +++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp @@ -104,16 +104,8 @@ transform::AlternativesOp::getEntrySuccessorOperands(RegionBranchPoint point) { void transform::AlternativesOp::getSuccessorRegions( RegionBranchPoint point, SmallVectorImpl ®ions) { - for (Region &alternative : llvm::drop_begin( - 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( @@ -1502,16 +1494,8 @@ void transform::ForeachOp::getEffects( void transform::ForeachOp::getSuccessorRegions( RegionBranchPoint point, SmallVectorImpl ®ions) { - 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 @@ -2702,16 +2686,8 @@ transform::SequenceOp::getEntrySuccessorOperands(RegionBranchPoint point) { void transform::SequenceOp::getSuccessorRegions( RegionBranchPoint point, SmallVectorImpl ®ions) { - 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(