[mlir][openacc] Add firstprivate representation

Add a representation for firstprivate clause modeled on the
private representation added in D150622.
The firstprivate recipe operation has an additional mandatory
region representing a sequences of operations needed to copy
the initial value to the created private copy.

Depends on D150622

Reviewed By: razvanlupusoru

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

3 files changed, 218 insertions, 23 deletions

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index fea9eaa4a4b6..36314e63c2cd 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -458,6 +458,71 @@ def OpenACC_PrivateRecipeOp : OpenACC_Op<"private.recipe",
 }
 
 //===----------------------------------------------------------------------===//
+// 2.5.14 firstprivate clause
+//===----------------------------------------------------------------------===//
+
+def OpenACC_FirstprivateRecipeOp : OpenACC_Op<"firstprivate.recipe",
+    [IsolatedFromAbove, Symbol]> {
+  let summary = "privatization recipe";
+
+  let description = [{
+    Declares an OpenACC privatization recipe with copy of the initial value.
+    The operation requires two mandatory regions and one optional.
+
+      1. The initializer region specifies how to allocate and initialize a new
+         private value. For example in Fortran, a derived-type might have a
+         default initialization. The region has an argument that contains the
+         value that need to be privatized. This is useful if the type is not
+         known at compile time and the private value is needed to create its
+         copy.
+      2. The copy region specifies how to copy the initial value to the newly
+         created private value. It takes the initial value and the privatized
+         value as arguments.
+      3. The destroy region specifies how to destruct the value when it reaches
+         its end of life. It takes the privatized value as argument.
+
+    A single privatization recipe can be used for multiple operand if they have
+    the same type and do not require a specific default initialization.
+
+    Example:
+
+    ```mlir
+    acc.firstprivate.recipe @privatization_f32 : f32 init {
+    ^bb0(%0: f32):
+      // init region contains a sequence of operations to create and
+      // initialize the copy if needed. It yields the create copy.
+    } copy {
+    ^bb0(%0: f32, %1: !llvm.ptr<f32>):
+      // copy region contains a sequence of operations to copy the initial value
+      // of the firstprivate value to the newly created value.
+    } destroy {
+    ^bb0(%0: f32)
+      // destroy region contains a sequences of operations to destruct the
+      // created copy.
+    }
+
+    // The privatization symbol is then used in the corresponding operation.
+    acc.parallel firstprivate(@privatization_f32 -> %a : f32) {
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttr:$type);
+
+  let regions = (region AnyRegion:$initRegion, AnyRegion:$copyRegion,
+                        AnyRegion:$destroyRegion);
+
+  let assemblyFormat = [{
+    $sym_name `:` $type attr-dict-with-keyword `init` $initRegion
+    `copy` $copyRegion
+    (`destroy` $destroyRegion^)?
+  }];
+
+  let hasRegionVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // 2.5.1 parallel Construct
 //===----------------------------------------------------------------------===//
 
@@ -921,7 +986,7 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
 
 // Yield operation for the acc.loop and acc.parallel operations.
 def OpenACC_YieldOp : OpenACC_Op<"yield", [ReturnLike, Terminator,
-    ParentOneOf<["ParallelOp, LoopOp, SerialOp, PrivateRecipeOp"]>]> {
+    ParentOneOf<["ParallelOp, LoopOp, SerialOp, PrivateRecipeOp, FirstprivateRecipeOp"]>]> {
   let summary = "Acc yield and termination operation";
 
   let description = [{
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 0a93447ef04e..a7c12b253d3f 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -338,31 +338,70 @@ struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> {
 // PrivateRecipeOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult acc::PrivateRecipeOp::verifyRegions() {
-  if (getInitRegion().empty())
-    return emitOpError() << "expects non-empty init region";
-  Block &initBlock = getInitRegion().front();
-  if (initBlock.getNumArguments() != 1 ||
-      initBlock.getArgument(0).getType() != getType())
-    return emitOpError() << "expects init region with one argument of the "
-                         << "privatization type";
-
-  for (YieldOp yieldOp : getInitRegion().getOps<YieldOp>()) {
-    if (yieldOp.getOperands().size() != 1 ||
-        yieldOp.getOperands().getTypes()[0] != getType())
-      return emitOpError() << "expects init region to yield a value "
-                              "of the privatization type";
+static LogicalResult verifyPrivateLikeRegion(Operation *op, Region &region,
+                                             StringRef regionName, Type type,
+                                             unsigned expectNbArg,
+                                             bool optionalRegion,
+                                             bool verifyYield) {
+  if (optionalRegion && region.empty())
+    return success();
+
+  if (region.empty())
+    return op->emitOpError() << "expects non-empty " << regionName << " region";
+  Block &firstBlock = region.front();
+  if (expectNbArg == 1 && (firstBlock.getNumArguments() != 1 ||
+                           firstBlock.getArgument(0).getType() != type))
+    return op->emitOpError() << "expects " << regionName
+                             << " region with one "
+                                "argument of the privatization type";
+  if (expectNbArg == 2 && (firstBlock.getNumArguments() != 2 ||
+                           firstBlock.getArgument(0).getType() != type))
+    return op->emitOpError() << "expects " << regionName
+                             << " region with two "
+                                "arguments of the privatization type";
+
+  if (verifyYield) {
+    for (YieldOp yieldOp : region.getOps<acc::YieldOp>()) {
+      if (yieldOp.getOperands().size() != 1 ||
+          yieldOp.getOperands().getTypes()[0] != type)
+        return op->emitOpError() << "expects " << regionName
+                                 << " region to "
+                                    "yield a value of the privatization type";
+    }
   }
+  return success();
+}
 
-  // Destroy region is optional.
-  if (getDestroyRegion().empty())
-    return success();
+LogicalResult acc::PrivateRecipeOp::verifyRegions() {
+  if (failed(verifyPrivateLikeRegion(*this, getInitRegion(), "init", getType(),
+                                     1, /*optional=*/false,
+                                     /*verifyYield=*/true)))
+    return failure();
+  if (failed(verifyPrivateLikeRegion(*this, getDestroyRegion(), "destroy",
+                                     getType(), 1, /*optional=*/true,
+                                     /*verifyYield=*/false)))
+    return failure();
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// FirstprivateRecipeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult acc::FirstprivateRecipeOp::verifyRegions() {
+  if (failed(verifyPrivateLikeRegion(*this, getInitRegion(), "init", getType(),
+                                     1, /*optional=*/false,
+                                     /*verifyYield=*/true)))
+    return failure();
 
-  Block &destroyBlock = getDestroyRegion().front();
-  if (destroyBlock.getNumArguments() != 1 ||
-      destroyBlock.getArgument(0).getType() != getType())
-    return emitOpError() << "expects destroy region with one argument of the "
-                         << "privatization type";
+  if (failed(verifyPrivateLikeRegion(*this, getCopyRegion(), "copy", getType(),
+                                     2, /*optional=*/false,
+                                     /*verifyYield=*/false)))
+    return failure();
+  if (failed(verifyPrivateLikeRegion(*this, getDestroyRegion(), "destroy",
+                                     getType(), 1, /*optional=*/true,
+                                     /*verifyYield=*/false)))
+    return failure();
   return success();
 }
 
diff --git a/mlir/test/Dialect/OpenACC/invalid.mlir b/mlir/test/Dialect/OpenACC/invalid.mlir
index 257435c9c4c7..02ddae197b46 100644
--- a/mlir/test/Dialect/OpenACC/invalid.mlir
+++ b/mlir/test/Dialect/OpenACC/invalid.mlir
@@ -303,3 +303,94 @@ acc.private.recipe @privatization_i32 : !llvm.ptr<i32> init {
 ^bb0(%arg0 : f32):
   "test.openacc_dummy_op"(%arg0) : (f32) -> ()
 }
+
+// -----
+
+// expected-error@+1 {{expects non-empty init region}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
+} copy {}
+
+// -----
+
+// expected-error@+1 {{expects init region with one argument of the privatization type}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
+^bb0(%arg0 : !llvm.ptr<f32>):
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
+  llvm.store %c0, %0 : !llvm.ptr<i32>
+  acc.yield %0 : !llvm.ptr<i32>
+} copy {}
+
+// -----
+
+// expected-error@+1 {{expects init region to yield a value of the privatization type}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<f32> init {
+^bb0(%arg0 : !llvm.ptr<f32>):
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
+  llvm.store %c0, %0 : !llvm.ptr<i32>
+  acc.yield %0 : !llvm.ptr<i32>
+} copy {}
+
+// -----
+
+// expected-error@+1 {{expects non-empty copy region}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
+^bb0(%arg0 : !llvm.ptr<i32>):
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
+  llvm.store %c0, %0 : !llvm.ptr<i32>
+  acc.yield %0 : !llvm.ptr<i32>
+} copy {
+}
+
+// -----
+
+// expected-error@+1 {{expects copy region with two arguments of the privatization type}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
+^bb0(%arg0 : !llvm.ptr<i32>):
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
+  llvm.store %c0, %0 : !llvm.ptr<i32>
+  acc.yield %0 : !llvm.ptr<i32>
+} copy {
+^bb0(%arg0 : f32):
+  "test.openacc_dummy_op"(%arg0) : (f32) -> ()
+}
+
+// -----
+
+// expected-error@+1 {{expects copy region with two arguments of the privatization type}}
+acc.firstprivate.recipe @privatization_i32 : !llvm.ptr<i32> init {
+^bb0(%arg0 : !llvm.ptr<i32>):
+  %c1 = arith.constant 1 : i32
+  %c0 = arith.constant 0 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<i32>
+  llvm.store %c0, %0 : !llvm.ptr<i32>
+  acc.yield %0 : !llvm.ptr<i32>
+} copy {
+^bb0(%arg0 : f32, %arg1 : i32):
+  "test.openacc_dummy_op"(%arg0) : (f32) -> ()
+}
+
+// -----
+
+// expected-error@+1 {{destroy region with one argument of the privatization type}}
+acc.firstprivate.recipe @privatization_i32 : i32 init {
+^bb0(%arg0 : i32):
+  %0 = arith.constant 1 : i32
+  acc.yield %0 : i32
+} copy {
+^bb0(%arg0 : i32, %arg1 : !llvm.ptr<i32>):
+  llvm.store %arg0, %arg1 : !llvm.ptr<i32>
+  acc.yield
+} destroy {
+^bb0(%arg0 : f32):
+  acc.yield
+}
+
+