From 7951ab12b9674c5999f269aef8fa93b7994f27f0 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval@gmail.com>
Date: Tue, 11 Apr 2023 10:42:44 -0700
Subject: [PATCH] [mlir][openacc] Relax the single block constraint on acc.loop

The acc.loop operation was constrained by the SingleBlockImplicitTerminator.
This patch relax this constraint to allow multiple block in the loop.

Reviewed By: razvanlupusoru

Differential Revision: https://reviews.llvm.org/D148025
---
 mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td |  5 +---
 mlir/test/Dialect/OpenACC/ops.mlir              | 31 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 9203e8b..9a6e942 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -343,9 +343,7 @@ def OpenACC_ExitDataOp : OpenACC_Op<"exit_data", [AttrSizedOperandSegments]> {
 // 2.9 loop Construct
 //===----------------------------------------------------------------------===//
 
-def OpenACC_LoopOp : OpenACC_Op<"loop",
-      [AttrSizedOperandSegments,
-       SingleBlockImplicitTerminator<"acc::YieldOp">]> {
+def OpenACC_LoopOp : OpenACC_Op<"loop", [AttrSizedOperandSegments]> {
   let summary = "loop construct";
 
   let description = [{
@@ -367,7 +365,6 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
     ```
   }];
 
-
   let arguments = (ins OptionalAttr<I64Attr>:$collapse,
                        Optional<IntOrIndex>:$gangNum,
                        Optional<IntOrIndex>:$gangStatic,
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index 48d343b..de5c3b4 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -323,6 +323,37 @@ func.func @testloopop() -> () {
 
 // -----
 
+func.func @acc_loop_multiple_block() {
+  acc.parallel {
+    acc.loop {
+      %c1 = arith.constant 1 : index
+      cf.br ^bb1(%c1 : index)
+    ^bb1(%9: index):
+      %c0 = arith.constant 0 : index
+      %12 = arith.cmpi sgt, %9, %c0 : index
+      cf.cond_br %12, ^bb2, ^bb3
+    ^bb2:
+      %c1_0 = arith.constant 1 : index
+      %c10 = arith.constant 10 : index
+      %22 = arith.subi %c10, %c1_0 : index
+      cf.br ^bb1(%22 : index)
+    ^bb3:
+      acc.yield
+    }
+    acc.yield
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @acc_loop_multiple_block()
+// CHECK: acc.parallel
+// CHECK: acc.loop
+// CHECK-3: ^bb{{.*}}
+// CHECK: acc.yield
+// CHECK: acc.yield
+
+// -----
+
 func.func @testparallelop(%a: memref<10xf32>, %b: memref<10xf32>, %c: memref<10x10xf32>) -> () {
   %i64value = arith.constant 1 : i64
   %i32value = arith.constant 1 : i32
-- 
2.7.4