[mlir][OpenMP] Add omp.single

[shraiysh]

Summary

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

Also added tests for the same.

Co-authored-by: Kiran Kumar T P <[e-mail]>

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

228

I think copyprivate is missed. __kmpc_copyprivate is needed when lowering from MLIR to LLVMIR. I guess all the directives and clauses related to global thread variables need to be specified in OpenMP dialect.

mlir/test/Dialect/OpenMP/invalid.mlir

974

Remove this line.

[shraiysh]

Addressed comments.

[shraiysh]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

228

Isn't privatization being handled during translation from PFT to MLIR at the moment? I did not put it in the operation because of that.

[shraiysh]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

228

Data sharing and data copying - both have not been handled in MLIR so far. While data sharing (privatization) has been handled during translation from PFT to MLIR, data copying (copyin and copyprivate) have not been handled anywhere. We need to have a discussion where these should go. (I think their handling is going to be very similar to data sharing clauses and hence should go into translation from PFT to MLIR but I am not sure about this). @kiranchandramohan can you please advise on where these should be handled?

[kiranchandramohan]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

228

Yes, as @shraiysh says, at the moment we do not have handling for data-sharing or data-copying clauses in the OpenMP dialect. @peixin is it OK for you to add them (if necessary) when we start handling those clauses?

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

228

The data sharing clauses for local data can be handled when lowering to MLIR such as private, shared, firstprivate. For global data, they will be needed when lowering to LLVM IR such as copyin and copyprivate.

Sure. Let's add them when we start working on them.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let assemblyFormat = [{
     ) $region attr-dict
   }];
 
   let hasVerifier = 1;
   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
 //===----------------------------------------------------------------------===//
 
 def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
                          AllTypesMatch<["lowerBound", "upperBound", "step"]>,
                          RecursiveSideEffects]> {
   let summary = "workshare loop construct";
   let description = [{

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

if (!(isa<SectionOp>(inst) || isa<TerminatorOp>(inst))) {
       return emitOpError()
              << "expected omp.section op or terminator op inside region";
     }
   }
 
   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
 //===----------------------------------------------------------------------===//
 
 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` inclusive? steps
 /// steps := `step` `(`ssa-id-list`)`
 ParseResult

mlir/test/Dialect/OpenMP/invalid.mlir

omp.section {
       omp.terminator
     }
     omp.terminator
   ^bb2:
     omp.terminator
   }
   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
+}

mlir/test/Dialect/OpenMP/ops.mlir

omp.section {
       // CHECK: "test.payload"(%{{.*}}) : (!llvm.ptr<f32>) -> ()
       "test.payload"(%redn_var) : (!llvm.ptr<f32>) -> ()
     }
     // CHECK: omp.terminator
     omp.terminator
   }
   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
+}