[mlir][openacc] Add private representation

[clementval]

Summary

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) {
}

[clementval]

Note that the same operation could be used for firstprivate with some small additions.

[jeanPerier]

Looks great!

[razvanlupusoru]

Awesome! This design what I had envisioned with regards to keeping track of initializer. This same approach should work for scalar reductions also whose private copies have explicit initial values in the spec - not just types with default initializer.

I have a few questions:

1. Should the allocation logic should be part of initialization? For example, 2.5.13 private clause says "; it declares that a copy of each item on the list will be created for each gang in all dimensions.". Depending on how gang parallelism is mapped to the underlying hardware, the allocation may not be a stack allocation for which an alloca is relevant for. However, on the flip side, we need some info from frontend with regards to allocation size/type. So maybe allocation/initialization sections should be separated or we need some other way to encode size info.
2. Should the init section should be enforced to exist? OpenACC spec does not specify or enforce initial values for private. Thus it depends on the source language itself. However, I guess you could argue that for those cases, each dialect should have an initializer with an op specifying the "undefined" state. That said, if you think for uniformity reasons we should always capture this info even if undefined, I would be OK with it.
3. Can we find a better name than "declare"? declare is an OpenACC directive with specific semantics. As I thought about it, I kind of like acc.private.recipe, acc.reduction.recipe, etc since these are operations that provide the instructions which will eventually be used for materializing semantics.

[clementval]

In D150622#4346229, @razvanlupusoru wrote:

Awesome! This design what I had envisioned with regards to keeping track of initializer. This same approach should work for scalar reductions also whose private copies have explicit initial values in the spec - not just types with default initializer.

I have a few questions:

1. Should the allocation logic should be part of initialization? For example, 2.5.13 private clause says "; it declares that a copy of each item on the list will be created for each gang in all dimensions.". Depending on how gang parallelism is mapped to the underlying hardware, the allocation may not be a stack allocation for which an alloca is relevant for. However, on the flip side, we need some info from frontend with regards to allocation size/type. So maybe allocation/initialization sections should be separated or we need some other way to encode size info.

We can for sure have a separate regions for allocation/init or the region can just contains the alloc + init information. If the allocation change from the standard one, we should then probably have a pass that can update the init region or duplicate the whole op with specific allocation.

2. Should the init section should be enforced to exist? OpenACC spec does not specify or enforce initial values for private. Thus it depends on the source language itself. However, I guess you could argue that for those cases, each dialect should have an initializer with an op specifying the "undefined" state. That said, if you think for uniformity reasons we should always capture this info even if undefined, I would be OK with it.

I'm fine making the init region optional if we have a dedicated alloc region otherwise I was thinking that the init region holds at least the information about how to allocate the private copy so that's why I made it non optional.

3. Can we find a better name than "declare"? declare is an OpenACC directive with specific semantics. As I thought about it, I kind of like acc.private.recipe, acc.reduction.recipe, etc since these are operations that provide the instructions which will eventually be used for materializing semantics.

Right declare is a specific directive in OpenACC so we might want to change the name here. I'm fine with recipe since I can't find another name.

[clementval]

• Rename declare to recipe
• Couple of typos fix
• emphasize that init region is for both allocation + initialization if needed

[razvanlupusoru]

Thank you.

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_OPENACC_OPENACC_H_
 #define MLIR_DIALECT_OPENACC_OPENACC_H_
 
 #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"
 #include "mlir/Dialect/OpenACC/OpenACCTypeInterfaces.h.inc"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #define GET_TYPEDEF_CLASSES

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

 #ifndef OPENACC_OPS
 #define OPENACC_OPS
 
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 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"
 
 // AccCommon requires definition of OpenACC_Dialect.
 include "mlir/Dialect/OpenACC/AccCommon.td"
 
 // Base class for OpenACC dialect ops.
let extraClassDeclaration = [{
     /// Check if this is an acc update self.
     bool isSelf() {
       return getDataClause() == acc::DataClause::acc_update_self;
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// 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
 //===----------------------------------------------------------------------===//
 
 def OpenACC_ParallelOp : OpenACC_Op<"parallel",
     [AttrSizedOperandSegments, RecursiveMemoryEffects]> {
   let summary = "parallel construct";
   let description = [{
     The "acc.parallel" operation represents a parallel construct block. It has
let assemblyFormat = [{
     attr-dict-with-keyword
   }];
 
   let hasVerifier = 1;
 }
 
 // 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 = [{
     `acc.yield` is a special terminator operation for block inside regions in
     acc ops (parallel and loop). It returns values to the immediately enclosing
     acc op.
   }];
 

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

LogicalResult matchAndRewrite(OpTy op,
 
     return success();
   }
 };
 
 } // 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
 //===----------------------------------------------------------------------===//
 
 /// Check dataOperands for acc.parallel, acc.serial and acc.kernels.
 template <typename Op>
 static LogicalResult checkDataOperands(Op op,
                                        const mlir::ValueRange &operands) {
   for (mlir::Value operand : operands)

mlir/test/Dialect/OpenACC/invalid.mlir

 
 // -----
 
 %value = memref.alloc() : memref<10xf32>
 // expected-error@+1 {{expect data entry/exit operation or acc.getdeviceptr as defining op}}
 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) -> ()
+}

mlir/test/Dialect/OpenACC/ops.mlir

 
 // CHECK-LABEL: func.func @host_data_ops(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<f32>, %[[IFCOND:.*]]: i1)
 // CHECK: %[[PTR:.*]] = acc.use_device varPtr(%[[A]] : !llvm.ptr<f32>) -> !llvm.ptr<f32>
 // CHECK: acc.host_data dataOperands(%[[PTR]] : !llvm.ptr<f32>)
 // 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