[mlir][transf] Traits and interf. of alternatives, foreach, and yield.

This PR revisits the traits and interfaces of the `alternatives`,
`foreach`, and `yield`ops. First, it adds the `ReturnLike` trait to
`transform.yield`. This is required in the one-shot bufferization pass
since the merging of llvm#110332, which analyses any `FunctionOpInterface`
and expects them to have a `ReturnLike` terminator.

This uncovered problems with `alternatives` and `foreach`, which
implemented the `BranchRegionOpInterface`. That interface verifies 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`. Making `yield` a
`ReturnLike` op, those checks fail without other changes.

We explored and rejected the possibility to fix the implementation of
`RegionBranchOpInterface` of the two concerned ops in llvm#111408. The
problem is that the interface expects the result passed from a region to
the parent op to *be* the result of the op, whereas in the two ops they
*contribute* to the result: in `alternatives` only the operand yielded
from one of the regions becomes the result whereas the others are
ignored, and in `foreach` the operands from all iterations are fused
into a new value that becomes the result of the op.

Signed-off-by: Ingo Müller <[email protected]>

Author: ingomueller-net

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

@@ -24,10 +24,7 @@ include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
 
 def AlternativesOp : TransformDialectOp<"alternatives",
-    [DeclareOpInterfaceMethods<RegionBranchOpInterface,
-        ["getEntrySuccessorOperands", "getSuccessorRegions",
-         "getRegionInvocationBounds"]>,
-     DeclareOpInterfaceMethods<TransformOpInterface>,
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      IsolatedFromAbove, PossibleTopLevelTransformOpTrait,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
@@ -37,8 +34,8 @@ def AlternativesOp : TransformDialectOp<"alternatives",
     sequence of transform operations to be applied to the same payload IR. The
     regions are visited in order of appearance, and transforms in them are
     applied in their respective order of appearance. If one of these transforms
-    fails to apply, the remaining ops in the same region are skipped an the next
-    region is attempted. If all transformations in a region succeed, the
+    fails to apply, the remaining ops in the same region are skipped and the
+    next region is attempted. If all transformations in a region succeed, the
     remaining regions are skipped and the entire "alternatives" transformation
     succeeds. If all regions contained a failing transformation, the entire
     "alternatives" transformation fails.
@@ -610,8 +607,6 @@ def ForeachMatchOp : TransformDialectOp<"foreach_match", [
 def ForeachOp : TransformDialectOp<"foreach",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-     DeclareOpInterfaceMethods<RegionBranchOpInterface, [
-         "getSuccessorRegions", "getEntrySuccessorOperands"]>,
      SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">
     ]> {
   let summary = "Executes the body for each element of the payload";
@@ -1358,7 +1353,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

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

@@ -94,38 +94,6 @@ ensurePayloadIsSeparateFromTransform(transform::TransformOpInterface transform,
 // AlternativesOp
 //===----------------------------------------------------------------------===//
 
-OperandRange
-transform::AlternativesOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  if (!point.isParent() && getOperation()->getNumOperands() == 1)
-    return getOperation()->getOperands();
-  return OperandRange(getOperation()->operand_end(),
-                      getOperation()->operand_end());
-}
-
-void transform::AlternativesOp::getSuccessorRegions(
-    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
-  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());
-}
-
-void transform::AlternativesOp::getRegionInvocationBounds(
-    ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
-  (void)operands;
-  // The region corresponding to the first alternative is always executed, the
-  // remaining may or may not be executed.
-  bounds.reserve(getNumRegions());
-  bounds.emplace_back(1, 1);
-  bounds.resize(getNumRegions(), InvocationBounds(0, 1));
-}
-
 static void forwardEmptyOperands(Block *block, transform::TransformState &state,
                                  transform::TransformResults &results) {
   for (const auto &res : block->getParentOp()->getOpResults())
@@ -1500,28 +1468,6 @@ void transform::ForeachOp::getEffects(
   producesHandle(getOperation()->getOpResults(), effects);
 }
 
-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();
-}
-
-OperandRange
-transform::ForeachOp::getEntrySuccessorOperands(RegionBranchPoint point) {
-  // Each block argument handle is mapped to a subset (one op to be precise)
-  // of the payload of the corresponding `targets` operand of ForeachOp.
-  assert(point == getBody() && "unexpected region index");
-  return getOperation()->getOperands();
-}
-
 transform::YieldOp transform::ForeachOp::getYieldOp() {
   return cast<transform::YieldOp>(getBody().front().getTerminator());
 }
@@ -2702,16 +2648,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(