[Mlir-commits] [mlir] 9775c0c - [mlir] Fix ControlFlowInterfaces implementation for Async dialect

[Vladislav Vinogradov]

Author: Vladislav Vinogradov
Date: 2021-08-20T12:14:45+03:00
New Revision: 9775c0c9f0bb8e72d48912d5dd572dd38d7ca819

URL: https://github.com/llvm/llvm-project/commit/9775c0c9f0bb8e72d48912d5dd572dd38d7ca819
DIFF: https://github.com/llvm/llvm-project/commit/9775c0c9f0bb8e72d48912d5dd572dd38d7ca819.diff

LOG: [mlir] Fix ControlFlowInterfaces implementation for Async dialect

* Add `RegionBranchTerminatorOpInterface` to `YieldOp`.
* Implement `getSuccessorEntryOperands` in `ExecuteOp`.
* Fix `getSuccessorRegions` implementation in `ExecuteOp`.

Reviewed By: ezhulenev

Differential Revision: https://reviews.llvm.org/D108373

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
    mlir/lib/Dialect/Async/IR/Async.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
index d168b8cefad8a..ca7aa32646059 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
+++ b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
@@ -29,7 +29,8 @@ class Async_Op<string mnemonic, list<OpTrait> traits = []> :
 def Async_ExecuteOp :
   Async_Op<"execute", [SingleBlockImplicitTerminator<"YieldOp">,
                        DeclareOpInterfaceMethods<RegionBranchOpInterface,
-                                                 ["getNumRegionInvocations"]>,
+                                                 ["getSuccessorEntryOperands",
+                                                  "getNumRegionInvocations"]>,
                        AttrSizedOperandSegments]> {
   let summary = "Asynchronous execute operation";
   let description = [{
@@ -99,7 +100,9 @@ def Async_ExecuteOp :
 }
 
 def Async_YieldOp :
-    Async_Op<"yield", [HasParent<"ExecuteOp">, NoSideEffect, Terminator]> {
+    Async_Op<"yield", [
+      HasParent<"ExecuteOp">, NoSideEffect, Terminator,
+      DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
   let summary = "terminator for Async execute operation";
   let description = [{
     The `async.yield` is a special terminator operation for the block inside

diff  --git a/mlir/lib/Dialect/Async/IR/Async.cpp b/mlir/lib/Dialect/Async/IR/Async.cpp
index 67b4096122745..030c8787b7939 100644
--- a/mlir/lib/Dialect/Async/IR/Async.cpp
+++ b/mlir/lib/Dialect/Async/IR/Async.cpp
@@ -48,6 +48,12 @@ static LogicalResult verify(YieldOp op) {
   return success();
 }
 
+MutableOperandRange
+YieldOp::getMutableSuccessorOperands(Optional<unsigned> index) {
+  assert(!index.hasValue());
+  return operandsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 /// ExecuteOp
 //===----------------------------------------------------------------------===//
@@ -55,24 +61,28 @@ static LogicalResult verify(YieldOp op) {
 constexpr char kOperandSegmentSizesAttr[] = "operand_segment_sizes";
 
 void ExecuteOp::getNumRegionInvocations(
-    ArrayRef<Attribute> operands, SmallVectorImpl<int64_t> &countPerRegion) {
-  (void)operands;
+    ArrayRef<Attribute>, SmallVectorImpl<int64_t> &countPerRegion) {
   assert(countPerRegion.empty());
   countPerRegion.push_back(1);
 }
 
+OperandRange ExecuteOp::getSuccessorEntryOperands(unsigned index) {
+  assert(index == 0 && "invalid region index");
+  return operands();
+}
+
 void ExecuteOp::getSuccessorRegions(Optional<unsigned> index,
-                                    ArrayRef<Attribute> operands,
+                                    ArrayRef<Attribute>,
                                     SmallVectorImpl<RegionSuccessor> &regions) {
   // The `body` region branch back to the parent operation.
   if (index.hasValue()) {
-    assert(*index == 0);
-    regions.push_back(RegionSuccessor(getResults()));
+    assert(*index == 0 && "invalid region index");
+    regions.push_back(RegionSuccessor(results()));
     return;
   }
 
   // Otherwise the successor is the body region.
-  regions.push_back(RegionSuccessor(&body()));
+  regions.push_back(RegionSuccessor(&body(), body().getArguments()));
 }
 
 void ExecuteOp::build(OpBuilder &builder, OperationState &result,