[mlir][openacc] Add private representation

Currently privatization is not well represented in the OpenACC dialect. This
patch is a prototype to model privatization with a dedicated operation that
declares how the private value is created and destroyed.

The newly introduced operation is `acc.private.recipe` and it declares
an OpenACC privatization. The operation requires one mandatory and
one optional region.

  1. The initializer region specifies how to create 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
     a known at compile time and the private value is needed to create its
     copy.
  2. The destroy region specifies how to destruct the value when it reaches
     its end of life. It takes the privatized value as argument.

A same privatization can be used for multiple operand if they have the same
type and do not require a specific default initialization.

Example:

```mlir
acc.private.recipe @privatization_f32 : f32 init {
^bb0(%0: f32):
  // init region contains a sequence of operations to create and initialize the
  // copy if needed.
} 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 private(@privatization_f32 -> %a : f32) {
}
```

Reviewed By: razvanlupusoru, jeanPerier

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

5 files changed, 185 insertions, 1 deletions

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
index aadcfff93238..00fe7faf416f 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
@@ -16,6 +16,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/SymbolTable.h"
 
 #include "mlir/Dialect/OpenACC/OpenACCOpsDialect.h.inc"
 #include "mlir/Dialect/OpenACC/OpenACCOpsEnums.h.inc"
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 7beec7f47011..fea9eaa4a4b6 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -18,6 +18,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/OpenACC/OpenACCBase.td"
 include "mlir/Dialect/OpenACC/OpenACCOpsTypes.td"
 include "mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td"
@@ -400,6 +401,63 @@ def OpenACC_UpdateHostOp : OpenACC_DataExitOp<"update_host",
 }
 
 //===----------------------------------------------------------------------===//
+// 2.5.13 private clause
+//===----------------------------------------------------------------------===//
+
+def OpenACC_PrivateRecipeOp : OpenACC_Op<"private.recipe",
+    [IsolatedFromAbove, Symbol]> {
+  let summary = "privatization recipe";
+
+  let description = [{
+    Declares an OpenACC privatization recipe. The operation requires one
+    mandatory and one optional region.
+
+      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 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.private.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.
+    } 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 private(@privatization_f32 -> %a : f32) {
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttr:$type);
+
+  let regions = (region AnyRegion:$initRegion,
+                        AnyRegion:$destroyRegion);
+
+  let assemblyFormat = [{
+    $sym_name `:` $type attr-dict-with-keyword `init` $initRegion
+    (`destroy` $destroyRegion^)?
+  }];
+
+  let hasRegionVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // 2.5.1 parallel Construct
 //===----------------------------------------------------------------------===//
 
@@ -863,7 +921,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"]>]> {
+    ParentOneOf<["ParallelOp, LoopOp, SerialOp, PrivateRecipeOp"]>]> {
   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 4013657a4b28..0a93447ef04e 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -335,6 +335,38 @@ struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> {
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// 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";
+  }
+
+  // Destroy region is optional.
+  if (getDestroyRegion().empty())
+    return success();
+
+  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";
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/test/Dialect/OpenACC/invalid.mlir b/mlir/test/Dialect/OpenACC/invalid.mlir
index 83ef73cf577f..257435c9c4c7 100644
--- a/mlir/test/Dialect/OpenACC/invalid.mlir
+++ b/mlir/test/Dialect/OpenACC/invalid.mlir
@@ -258,3 +258,48 @@ acc.serial dataOperands(%value : memref<10xf32>) {
 acc.kernels dataOperands(%value : memref<10xf32>) {
   acc.yield
 }
+
+// -----
+
+// expected-error@+1 {{expects non-empty init region}}
+acc.private.recipe @privatization_i32 : !llvm.ptr<i32> init {
+}
+
+// -----
+
+// expected-error@+1 {{expects init region with one argument of the privatization type}}
+acc.private.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>
+}
+
+// -----
+
+// expected-error@+1 {{expects init region to yield a value of the privatization type}}
+acc.private.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>
+}
+
+// -----
+
+// expected-error@+1 {{expects destroy region with one argument of the privatization type}}
+acc.private.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>
+} destroy {
+^bb0(%arg0 : f32):
+  "test.openacc_dummy_op"(%arg0) : (f32) -> ()
+}
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index 3c35219de03b..dd396a2a3def 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -1287,3 +1287,51 @@ func.func @host_data_ops(%a: !llvm.ptr<f32>, %ifCond: i1) -> () {
 // CHECK: acc.host_data dataOperands(%[[PTR]] : !llvm.ptr<f32>) {
 // CHECK: } attributes {if_present}
 // CHECK: acc.host_data if(%[[IFCOND]]) dataOperands(%[[PTR]] : !llvm.ptr<f32>)
+
+// -----
+
+acc.private.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>
+}
+
+// CHECK: acc.private.recipe @privatization_i32 : !llvm.ptr<i32> init {
+// CHECK: %[[C1:.*]] = arith.constant 1 : i32
+// CHECK: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[C1]] x i32 : (i32) -> !llvm.ptr<i32>
+// CHECK: llvm.store %[[C0]], %[[ALLOCA]] : !llvm.ptr<i32>
+// CHECK: acc.yield %[[ALLOCA]] : !llvm.ptr<i32>
+
+// -----
+
+func.func private @destroy_struct(!llvm.struct<(i32, i32)>) -> ()
+
+acc.private.recipe @privatization_struct_i32_i64 : !llvm.struct<(i32, i32)> init {
+^bb0(%arg0 : !llvm.struct<(i32, i32)>):
+  %c1 = arith.constant 1 : i32
+  %0 = llvm.mlir.undef : !llvm.struct<(i32, i32)>
+  %1 = llvm.insertvalue %c1, %0[0] : !llvm.struct<(i32, i32)>
+  %2 = llvm.insertvalue %c1, %1[1] : !llvm.struct<(i32, i32)>
+  acc.yield %2 : !llvm.struct<(i32, i32)>
+} destroy {
+^bb0(%arg0: !llvm.struct<(i32, i32)>):
+  func.call @destroy_struct(%arg0) : (!llvm.struct<(i32, i32)>) -> ()
+  acc.terminator
+}
+
+// CHECK: func.func private @destroy_struct(!llvm.struct<(i32, i32)>)
+
+// CHECK: acc.private.recipe @privatization_struct_i32_i64 : !llvm.struct<(i32, i32)> init {
+// CHECK:   %[[C1:.*]] = arith.constant 1 : i32
+// CHECK:   %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32, i32)>
+// CHECK:   %[[UNDEF1:.*]] = llvm.insertvalue %[[C1]], %[[UNDEF]][0] : !llvm.struct<(i32, i32)> 
+// CHECK:   %[[UNDEF2:.*]] = llvm.insertvalue %[[C1]], %[[UNDEF1]][1] : !llvm.struct<(i32, i32)> 
+// CHECK:   acc.yield %[[UNDEF2]] : !llvm.struct<(i32, i32)>
+// CHECK: } destroy {
+// CHECK: ^bb0(%[[ARG0:.*]]: !llvm.struct<(i32, i32)>):
+// CHECK:   func.call @destroy_struct(%[[ARG0]]) : (!llvm.struct<(i32, i32)>) -> ()
+// CHECK:   acc.terminator