[Mlir-commits] [mlir] [mlir][transf] Traits and interf. of alternatives, foreach, and yield. (PR #112169)

Ingo Müller llvmlistbot at llvm.org
Tue Oct 29 01:13:06 PDT 2024


https://github.com/ingomueller-net updated https://github.com/llvm/llvm-project/pull/112169

>From 2f05cb8c7433ffae1fb2fe23cbc63d1f2274f053 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller at google.com>
Date: Mon, 7 Oct 2024 17:13:20 +0000
Subject: [PATCH 1/2] [mlir][transf] Traits and interf. of alternatives,
 foreach, and yield.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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

Signed-off-by: Ingo Müller <ingomueller at google.com>
---
 .../mlir/Dialect/Transform/IR/TransformOps.td | 14 ++--
 .../lib/Dialect/Transform/IR/TransformOps.cpp | 66 +------------------
 2 files changed, 7 insertions(+), 73 deletions(-)

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(

>From a8db082e059f5f2cd71c471dd0fb36e8e54b3292 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller at google.com>
Date: Tue, 29 Oct 2024 08:12:24 +0000
Subject: [PATCH 2/2] Address @ftynse's comments: document decision from this
 PR.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Ingo Müller <ingomueller at google.com>
---
 .../mlir/Dialect/Transform/IR/TransformOps.td     | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index 0481eb4ef7f0bd..195d794e5a835f 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
@@ -87,6 +87,13 @@ def AlternativesOp : TransformDialectOp<"alternatives",
       transform.yield %arg0 : !transform.any_op
     }
     ```
+
+    Note that this operation does not implement the `RegionBranchOpInterface`.
+    That interface verifies that the operands and results passed across the
+    control flow edges are equal (or compatible). In particular, it expects the
+    result passed from a region to its successor to be the argument of that
+    region; however, the argument of all `alternatives` regions are always
+    provided by the parent op and never by the precedessor region.
   }];
 
   let arguments = (ins Optional<TransformHandleTypeInterface>:$scope);
@@ -641,6 +648,14 @@ def ForeachOp : TransformDialectOp<"foreach",
     sequence fails immediately with the same failure, leaving the payload IR in
     a potentially invalid state, i.e., this operation offers no transformation
     rollback capabilities.
+
+    Note that this operation does not implement the `RegionBranchOpInterface`.
+    That interface verifies that the operands and results passed across the
+    control flow edges are equal (or compatible). In particular, it expects the
+    result passed from a region to the parent to *be* the result of that op;
+    however, the result of the `body` region only *contributes* to the result
+    in that the result of the op is an aggregation of of the results of all
+    iterations of the body.
   }];
 
   let arguments = (ins Variadic<Transform_AnyHandleOrParamType>:$targets,



More information about the Mlir-commits mailing list