[Mlir-commits] [mlir] [mlir][OpenMP] Added omp.region operation (PR #65243)

[llvmlistbot at llvm.org]

https://github.com/shraiysh created https://github.com/llvm/llvm-project/pull/65243:

This patch adds omp.region operation. This is equivalent to a code block surrounded by an OpenMP construct in C/C++/Fortran. This is a part of the effort to implement the canonical loop operation in OpenMP dialect.


>From ff635ce0ce910f0cde248a4babb3c27333ddc108 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Sun, 3 Sep 2023 22:40:10 -0500
Subject: [PATCH] [mlir][OpenMP] Added omp.region operation

This patch adds omp.region operation. This is equivalent to a code block surrounded by an OpenMP construct in C/C++/Fortran. This is a part of the effort to implement the canonical loop operation in OpenMP dialect.
---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 34 ++++++++++++
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  | 15 ++++++
 mlir/test/Dialect/OpenMP/region.mlir          | 53 +++++++++++++++++++
 3 files changed, 102 insertions(+)
 create mode 100644 mlir/test/Dialect/OpenMP/region.mlir

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index b1e1fe00b8594a..cf5b246178f5d7 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1787,4 +1787,38 @@ def ClauseRequiresAttr :
   EnumAttr<OpenMP_Dialect, ClauseRequires, "clause_requires"> {
 }
 
+
+def RegionOp : OpenMP_Op<"region"> {
+  let summary = "Encapsulates a region in an OpenMP construct.";
+  let description = [{
+    Encapsulates a region surrounded by an OpenMP Construct. The intended use
+    of this operation is that within an OpenMP operation's region, there would
+    be a single `omp.region` operation and a terminator operation as shown
+    below.
+
+    ```
+    // Example with `omp.task`
+    omp.task {
+      omp.region {
+        call @foo : () -> ()
+      }
+      omp.terminator
+    }
+    ```
+
+    This operation is especially more useful in operations that use `omp.yield`
+    as a terminator. For example, in the proposed canonical loop operation,
+    this operation would help fix the arguments of the yield operation and it
+    is not affected by branches in the region assosciated with the canonical
+    loop. In cases where the yielded value has to be a compile time constant,
+    this provides a mechanism to avoid complex static analysis for the constant
+    value.
+  }];
+  let regions = (region AnyRegion:$block);
+  let assemblyFormat = [{
+    $block attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // OPENMP_OPS
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 2ba5f1aca9cf6b..2a2fcdb788cb99 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1524,6 +1524,21 @@ LogicalResult CancellationPointOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// RegionOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult RegionOp::verify() {
+  Operation *parentOp = (*this)->getParentOp();
+  if (!parentOp)
+    return emitOpError() << "`omp.region` must have a parent";
+
+  if (!isa<OpenMPDialect>(parentOp->getDialect()))
+    return emitOpError()
+           << "`omp.region` must be used under an OpenMP Dialect operation.";
+  return success();
+}
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 
diff --git a/mlir/test/Dialect/OpenMP/region.mlir b/mlir/test/Dialect/OpenMP/region.mlir
new file mode 100644
index 00000000000000..4e0ddbc07e9ec9
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/region.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: @basic_omp_region
+// CHECK-NEXT: omp.parallel {
+// CHECK-NEXT:   omp.region {
+// CHECK-NEXT:     "test.foo"() : () -> ()
+// CHECK-NEXT:     omp.terminator
+// CHECK-NEXT:   }
+// CHECK-NEXT:   omp.terminator
+// CHECK-NEXT: }
+// CHECK-NEXT: return
+func.func @basic_omp_region() {
+  omp.parallel {
+    omp.region {
+      "test.foo"() : () -> ()
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: @omp_region_with_branch
+// CHECK-NEXT: omp.task {
+// CHECK-NEXT:   omp.region {
+// CHECK-NEXT:     %[[c:.*]] = "test.foo"() : () -> i1
+// CHECK-NEXT:     cf.cond_br %[[c]], ^[[bb1:.*]](%[[c]] : i1), ^[[bb2:.*]](%[[c]] : i1)
+// CHECK-NEXT:   ^[[bb1]](%[[arg:.*]]: i1):
+// CHECK-NEXT:     "test.bar"() : () -> ()
+// CHECK-NEXT:     omp.terminator
+// CHECK-NEXT:   ^[[bb2]](%[[arg2:.*]]: i1):
+// CHECK-NEXT:     "test.baz"() : () -> ()
+// CHECK-NEXT:     omp.terminator
+// CHECK-NEXT:   }
+// CHECK-NEXT:   omp.terminator
+// CHECK-NEXT: }
+// CHECK-NEXT: return
+func.func @omp_region_with_branch(%a: i32) {
+  omp.task {
+    omp.region {
+      %c = "test.foo"() : () -> i1
+      cf.cond_br %c, ^bb1(%c: i1), ^bb2(%c: i1)
+    ^bb1(%arg: i1):
+      "test.bar"() : () -> ()
+      omp.terminator
+    ^bb2(%arg2: i1):
+      "test.baz"() : () -> ()
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}