[mlir][OpenMP] Add omp.single

This patch adds omp.single according to Section 2.8.2 of OpenMP 5.0.

Also added tests for the same.

Reviewed By: peixin

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

Co-authored-by: Kiran Kumar T P <kirankumar.tp@amd.com>

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 5f973e6ca318e538c843917d83b3fbf7b135f734..efc471833d1c7c8d359e9d8b36b82714d711ddde 100644 (file)
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -209,6 +209,38 @@ def SectionsOp : OpenMP_Op<"sections", [AttrSizedOperandSegments]> {
   let hasRegionVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// 2.8.2 Single Construct
+//===----------------------------------------------------------------------===//
+
+def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
+  let summary = "single directive";
+  let description = [{
+    The single construct specifies that the associated structured block is 
+    executed by only one of the threads in the team (not necessarily the
+    master thread), in the context of its implicit task. The other threads
+    in the team, which do not execute the block, wait at an implicit barrier
+    at the end of the single construct unless a nowait clause is specified.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$allocate_vars,
+                       Variadic<AnyType>:$allocators_vars,
+                       UnitAttr:$nowait);
+
+  let regions = (region SizedRegion<1>:$region);
+
+  let assemblyFormat = [{
+    oilist(`allocate` `(`
+              custom<AllocateAndAllocator>(
+                $allocate_vars, type($allocate_vars),
+                $allocators_vars, type($allocators_vars)
+              ) `)`
+          |`nowait` $nowait
+    ) $region attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.9.2 Workshare Loop Construct
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 62f9498b11d949ceedde64a376330ae58d6b62b0..c43dea42aebc079e6f7bcd55742a9f51ad08aabe 100644 (file)
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -492,6 +492,15 @@ LogicalResult SectionsOp::verifyRegions() {
   return success();
 }
 
+LogicalResult SingleOp::verify() {
+  // Check for allocate clause restrictions
+  if (allocate_vars().size() != allocators_vars().size())
+    return emitError(
+        "expected equal sizes for allocate and allocator variables");
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//

diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index c5006f3eddc933c38e7ab58ebed4ef08394422e7..8c42052211dbc51bdc83be6220273be2262a75dd 100644 (file)
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -970,3 +970,13 @@ func @omp_sections() {
   }
   return
 }
+
+// -----
+
+func @omp_single(%data_var : memref<i32>) -> () {
+  // expected-error @below {{expected equal sizes for allocate and allocator variables}}
+  "omp.single" (%data_var) ({
+    omp.barrier
+  }) {operand_segment_sizes = dense<[1,0]> : vector<2xi32>} : (memref<i32>) -> ()
+  return
+}

diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index 70166a2328f9068371d52927f8a264c46fb3a7dc..f272f56626fd23d5c86f7419069f95f1ff8d3325 100644 (file)
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -720,3 +720,63 @@ func @omp_sectionsop(%data_var1 : memref<i32>, %data_var2 : memref<i32>,
   }
   return
 }
+
+// CHECK-LABEL: func @omp_single
+func @omp_single() {
+  omp.parallel {
+    // CHECK: omp.single {
+    omp.single {
+      "test.payload"() : () -> ()
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func @omp_single_nowait
+func @omp_single_nowait() {
+  omp.parallel {
+    // CHECK: omp.single nowait {
+    omp.single nowait {
+      "test.payload"() : () -> ()
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func @omp_single_allocate
+func @omp_single_allocate(%data_var: memref<i32>) {
+  omp.parallel {
+    // CHECK: omp.single allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>) {
+    omp.single allocate(%data_var : memref<i32> -> %data_var : memref<i32>) {
+      "test.payload"() : () -> ()
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func @omp_single_allocate_nowait
+func @omp_single_allocate_nowait(%data_var: memref<i32>) {
+  omp.parallel {
+    // CHECK: omp.single allocate(%{{.*}} : memref<i32> -> %{{.*}} : memref<i32>) nowait {
+    omp.single allocate(%data_var : memref<i32> -> %data_var : memref<i32>) nowait {
+      "test.payload"() : () -> ()
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}