[clang-tools-extra] [mlir][OpenMP] Added `omp.structured_region` operation (PR #68825)

[via cfe-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Shraiysh (shraiysh)

<details>
<summary>Changes</summary>

This patch adds `omp.structured_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.

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


4 Files Affected:

- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+34) 
- (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+13) 
- (modified) mlir/test/Dialect/OpenMP/invalid.mlir (+22) 
- (added) mlir/test/Dialect/OpenMP/structured_region.mlir (+53) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 5da05476a683725..b9e041a22b5bffb 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1987,4 +1987,38 @@ def ClauseRequiresAttr :
   EnumAttr<OpenMP_Dialect, ClauseRequires, "clause_requires"> {
 }
 
+
+def StructuredRegionOp : OpenMP_Op<"structured_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.structured_region` operation and a terminator operation as
+    shown below.
+
+    ```
+    // Example with `omp.sections`
+    omp.sections {
+      omp.structured_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:$region);
+  let assemblyFormat = [{
+    $region 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 2bf9355ed62676b..70811a230b36ccd 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -26,6 +26,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
+#include "llvm/Support/Debug.h"
 #include <cstddef>
 #include <optional>
 
@@ -1468,6 +1469,18 @@ LogicalResult CancellationPointOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// RegionOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult StructuredRegionOp::verify() {
+  Operation *parentOp = (*this)->getParentOp();
+  assert(parentOp && "'omp.region' op must have a parent");
+  if (!isa<OpenMPDialect>(parentOp->getDialect()))
+    return emitOpError() << "must be used under an OpenMP Dialect operation";
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // DataBoundsOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index c8025249e27004e..94814a4d398b8c7 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -1657,3 +1657,25 @@ func.func @omp_target_exit_data(%map1: memref<?xi32>) {
 }
 
 llvm.mlir.global internal @_QFsubEx() : i32
+
+// -----
+
+func.func @omp_region_invalid() {
+  // expected-error @below {{'omp.structured_region' op must be used under an OpenMP Dialect operation}}
+  omp.structured_region {
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func.func @omp_region_invalid(%c: i1) {
+  scf.if %c {
+    // expected-error @below {{'omp.structured_region' op must be used under an OpenMP Dialect operation}}
+    omp.structured_region {
+      omp.terminator
+    }
+  }
+  return
+}
diff --git a/mlir/test/Dialect/OpenMP/structured_region.mlir b/mlir/test/Dialect/OpenMP/structured_region.mlir
new file mode 100644
index 000000000000000..fc1e0edb07388eb
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/structured_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.structured_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.structured_region {
+      "test.foo"() : () -> ()
+      omp.terminator
+    }
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: @omp_region_with_branch
+// CHECK-NEXT: omp.task {
+// CHECK-NEXT:   omp.structured_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.structured_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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/68825