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

[ingomueller-net]

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

This uncovered problems with alternatives and foreach, which implemented the RegionBranchOpInterface. 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 #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.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ingo Müller (ingomueller-net)

Changes

This PR revisits the traits and interfaces of the alternatives, foreach, and yieldops. First, it 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.

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 #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.

---

Full diff: https://github.com/llvm/llvm-project/pull/112169.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Transform/IR/TransformOps.td (+5-9)
• (modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+2-64)

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index b946fc8875860b..0481eb4ef7f0bd 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/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
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 590cae9aa0d667..02b5c312d69d9c 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/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(

[ingomueller-net]

I am not sure whether this is the right approach. I have pushed this here as a basis for discussion.

I hadn't previously understood getEntrySuccessorOperands. Maybe the problems we found in #110332 can actually be fixed with that...

[ftynse]

Please add documentation explicitly stating that these interfaces should not be added to these ops and explain why. This will prevent somebody from coming and trying to add them back a month from now.

[ingomueller-net]

Please add documentation explicitly stating that these interfaces should not be added to these ops and explain why. This will prevent somebody from coming and trying to add them back a month from now.

Yeah, good idea. Done.

@matthias-springer: Do you want to have a look as well before I merge?

[ftynse]

Actually, I missed this and noticed on the other commit. I don't think yield is return-like. Is this necessary?

[ingomueller-net]

This change was proposed after #110322 broke things. But it got since then reversed, so it currently isn't "needed." The more I think about it, the more I doubt that #110322 is actually a good idea: it assumes that FunctionOpInterface has a ReturnLike terminator, which it shouldn't.