[mlir][openacc] Use new private representation in acc.parallel

[clementval]

Summary

Update acc.parallel private operands list to use the new design
introduced in D150622.

Test in flang/test/Lower/OpenACC/acc-parallel.f90 and
flang/test/Lower/OpenACC/acc-parallel-loop.f90 are temporarly
disabled and will be enabled with updated lowering in the follow-up
patch.

[razvanlupusoru]

Thank you! I enjoy seeing the privatization recipe - it makes it clearer how it is intended to be materialized.

flang/test/Lower/OpenACC/acc-parallel-loop.f90

 ! CHECK:      acc.parallel dataOperands(%[[ATTACH_F]], %[[ATTACH_G]] : !fir.ptr<f32>, !fir.ptr<f32>) {
 ! CHECK:        acc.loop {
 ! CHECK:          fir.do_loop
 ! CHECK:          acc.yield
 ! CHECK-NEXT:   }{{$}}
 ! CHECK:        acc.yield
 ! CHECK-NEXT: }{{$}}
 
-  !$acc parallel loop private(a) firstprivate(b)
-  DO i = 1, n
-    a(i) = b(i)
-  END DO
-
-! CHECK:      acc.parallel firstprivate(%[[B]] : !fir.ref<!fir.array<10xf32>>) private(%[[A]] : !fir.ref<!fir.array<10xf32>>) {
-! CHECK:        acc.loop private(%[[A]] : !fir.ref<!fir.array<10xf32>>) {
-! CHECK:          fir.do_loop
-! CHECK:          acc.yield
-! CHECK-NEXT:   }{{$}}
-! CHECK:        acc.yield
-! CHECK-NEXT: }{{$}}
+! TODO: will be updated after lowering change in privatization to MLIR
+!  !$acc parallel loop private(a) firstprivate(b)
+!  DO i = 1, n
+!    a(i) = b(i)
+!  END DO
+
+! TODO:      acc.parallel firstprivate(%[[B]] : !fir.ref<!fir.array<10xf32>>) private(%[[A]] : !fir.ref<!fir.array<10xf32>>) {
+! TODO:        acc.loop private(%[[A]] : !fir.ref<!fir.array<10xf32>>) {
+! TODO:          fir.do_loop
+! TODO:          acc.yield
+! TODO-NEXT:   }{{$}}
+! TODO:        acc.yield
+! TODO-NEXT: }{{$}}
 
   !$acc parallel loop seq
   DO i = 1, n
     a(i) = b(i)
   END DO
 
 ! CHECK:      acc.parallel {
 ! CHECK:        acc.loop {

flang/test/Lower/OpenACC/acc-parallel.f90

 !CHECK: %[[BOX_ADDR_E:.*]] = fir.box_addr %[[BOX_E]] : (!fir.box<!fir.ptr<f32>>) -> !fir.ptr<f32> 
 !CHECK: %[[ATTACH_E:.*]] = acc.attach varPtr(%[[BOX_ADDR_E]] : !fir.ptr<f32>) -> !fir.ptr<f32> {name = "e"} 
 !CHECK: acc.parallel dataOperands(%[[ATTACH_D]], %[[ATTACH_E]] : !fir.ptr<f32>, !fir.ptr<f32>) { 
 !CHECK:        acc.yield
 !CHECK-NEXT: }{{$}}
 !CHECK: acc.detach accPtr(%[[ATTACH_D]] : !fir.ptr<f32>) {dataClause = 10 : i64, name = "d"}
 !CHECK: acc.detach accPtr(%[[ATTACH_E]] : !fir.ptr<f32>) {dataClause = 10 : i64, name = "e"}
 
-  !$acc parallel private(a) firstprivate(b) private(c)
-  !$acc end parallel
-
-!CHECK:      acc.parallel firstprivate(%[[B]] : !fir.ref<!fir.array<10x10xf32>>) private(%[[A]], %[[C]] : !fir.ref<!fir.array<10x10xf32>>, !fir.ref<!fir.array<10x10xf32>>) {
-!CHECK:        acc.yield
-!CHECK-NEXT: }{{$}}
+! TODO: will be updated after lowering change in privatization to MLIR
+!  !$acc parallel private(a) firstprivate(b) private(c)
+!  !$acc end parallel
+
+!TODO:      acc.parallel firstprivate(%[[B]] : !fir.ref<!fir.array<10x10xf32>>) private(%[[A]], %[[C]] : !fir.ref<!fir.array<10x10xf32>>, !fir.ref<!fir.array<10x10xf32>>) {
+!TODO:        acc.yield
+!TODO-NEXT: }{{$}}
 
 end subroutine acc_parallel

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

let arguments = (ins Optional<IntOrIndex>:$async,
                        Optional<IntOrIndex>:$numGangs,
                        Optional<IntOrIndex>:$numWorkers,
                        Optional<IntOrIndex>:$vectorLength,
                        Optional<I1>:$ifCond,
                        Optional<I1>:$selfCond,
                        UnitAttr:$selfAttr,
                        OptionalAttr<OpenACC_ReductionOperatorAttr>:$reductionOp,
                        Variadic<AnyType>:$reductionOperands,
-                       Variadic<AnyType>:$gangPrivateOperands,
+                       Variadic<OpenACC_PointerLikeTypeInterface>:$gangPrivateOperands,
+                       OptionalAttr<SymbolRefArrayAttr>:$privatizations,
                        Variadic<AnyType>:$gangFirstPrivateOperands,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$dataClauseOperands,
                        OptionalAttr<DefaultValueAttr>:$defaultAttr);
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
     /// The number of data operands.
     unsigned getNumDataOperands();
 
     /// The i-th data operand passed.
     Value getDataOperand(unsigned i);
   }];
 
   let assemblyFormat = [{
     oilist(
         `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
       | `async` `(` $async `:` type($async) `)`
       | `firstprivate` `(` $gangFirstPrivateOperands `:`
             type($gangFirstPrivateOperands) `)`
       | `num_gangs` `(` $numGangs `:` type($numGangs) `)`
       | `num_workers` `(` $numWorkers `:` type($numWorkers) `)`
-      | `private` `(` $gangPrivateOperands `:` type($gangPrivateOperands) `)`
+      | `private` `(` custom<PrivatizationList>(
+            $gangPrivateOperands, type($gangPrivateOperands), $privatizations)
+        `)`
       | `vector_length` `(` $vectorLength `:` type($vectorLength) `)`
       | `wait` `(` $waitOperands `:` type($waitOperands) `)`
       | `self` `(` $selfCond `)`
       | `if` `(` $ifCond `)`
       | `reduction` `(` $reductionOperands `:` type($reductionOperands) `)`
     )
     $region attr-dict-with-keyword
   }];

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

if (yieldOp.getOperands().size() != 1 ||
       return emitOpError() << "expects combiner region to yield a value "
                               "of the reduction type";
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// Custom parser and printer verifier for private clause
+//===----------------------------------------------------------------------===//
+
+static ParseResult parsePrivatizationList(
+    mlir::OpAsmParser &parser,
+    llvm::SmallVectorImpl<mlir::OpAsmParser::UnresolvedOperand> &operands,
+    llvm::SmallVectorImpl<Type> &types, mlir::ArrayAttr &privatizationSymbols) {
+  llvm::SmallVector<SymbolRefAttr> privatizationVec;
+  if (failed(parser.parseCommaSeparatedList([&]() {
+        if (parser.parseAttribute(privatizationVec.emplace_back()) ||
+            parser.parseArrow() ||
+            parser.parseOperand(operands.emplace_back()) ||
+            parser.parseColonType(types.emplace_back()))
+          return failure();
+        return success();
+      })))
+    return failure();
+  llvm::SmallVector<mlir::Attribute> privatizations(privatizationVec.begin(),
+                                                    privatizationVec.end());
+  privatizationSymbols = ArrayAttr::get(parser.getContext(), privatizations);
+  return success();
+}
+
+static void
+printPrivatizationList(mlir::OpAsmPrinter &p, mlir::Operation *op,
+                       mlir::OperandRange privateOperands,
+                       mlir::TypeRange privateTypes,
+                       std::optional<mlir::ArrayAttr> privatizations) {
+  for (unsigned i = 0, e = privatizations->size(); i < e; ++i) {
+    if (i != 0)
+      p << ", ";
+    p << (*privatizations)[i] << " -> " << privateOperands[i] << " : "
+      << privateOperands[i].getType();
+  }
+}
+
+//===----------------------------------------------------------------------===//
 // 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)
     if (!mlir::isa<acc::AttachOp, acc::CopyinOp, acc::CopyoutOp, acc::CreateOp,
                    acc::DeleteOp, acc::DetachOp, acc::DevicePtrOp,
                    acc::GetDevicePtrOp, acc::NoCreateOp, acc::PresentOp>(
             operand.getDefiningOp()))
       return op.emitError(
           "expect data entry/exit operation or acc.getdeviceptr "
           "as defining op");
   return success();
 }
 
+static LogicalResult
+checkPrivatizationList(Operation *op,
+                       std::optional<mlir::ArrayAttr> privatizations,
+                       mlir::OperandRange privateOperands) {
+  if (!privateOperands.empty()) {
+    if (!privatizations || privatizations->size() != privateOperands.size())
+      return op->emitOpError() << "expected as many privatizations symbol "
+                                  "reference as private operands";
+  } else {
+    if (privatizations)
+      return op->emitOpError() << "unexpected privatizations symbol reference";
+    return success();
+  }
+
+  llvm::DenseSet<Value> privates;
+  for (auto args : llvm::zip(privateOperands, *privatizations)) {
+    mlir::Value privateOperand = std::get<0>(args);
+
+    if (!privates.insert(privateOperand).second)
+      return op->emitOpError() << "private operand appears more than once";
+
+    mlir::Type varType = privateOperand.getType();
+    auto symbolRef = std::get<1>(args).cast<SymbolRefAttr>();
+    auto decl =
+        SymbolTable::lookupNearestSymbolFrom<PrivateRecipeOp>(op, symbolRef);
+    if (!decl)
+      return op->emitOpError() << "expected symbol reference " << symbolRef
+                               << " to point to a private declaration";
+
+    if (decl.getType() && decl.getType() != varType)
+      return op->emitOpError()
+             << "expected private (" << varType
+             << ") to be the same type as private declaration ("
+             << decl.getType() << ")";
+  }
+
+  return success();
+}
+
 unsigned ParallelOp::getNumDataOperands() {
   return getReductionOperands().size() + getGangPrivateOperands().size() +
          getGangFirstPrivateOperands().size() + getDataClauseOperands().size();
 }
 
 Value ParallelOp::getDataOperand(unsigned i) {
   unsigned numOptional = getAsync() ? 1 : 0;
   numOptional += getNumGangs() ? 1 : 0;
   numOptional += getNumWorkers() ? 1 : 0;
   numOptional += getVectorLength() ? 1 : 0;
   numOptional += getIfCond() ? 1 : 0;
   numOptional += getSelfCond() ? 1 : 0;
   return getOperand(getWaitOperands().size() + numOptional + i);
 }
 
 LogicalResult acc::ParallelOp::verify() {
+  if (failed(checkPrivatizationList(*this, getPrivatizations(),
+                                    getGangPrivateOperands())))
+    return failure();
   return checkDataOperands<acc::ParallelOp>(*this, getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
 // SerialOp
 //===----------------------------------------------------------------------===//
 
 unsigned SerialOp::getNumDataOperands() {

mlir/test/Dialect/OpenACC/ops.mlir

 //  CHECK-NEXT:     } attributes {seq}
 //  CHECK-NEXT:     acc.yield
 //  CHECK-NEXT:   }
 //  CHECK-NEXT:   return %{{.*}} : memref<10x10xf32>
 //  CHECK-NEXT: }
 
 // -----
 
+acc.private.recipe @privatization_memref_10_f32 : memref<10xf32> init {
+^bb0(%arg0: memref<10xf32>):
+  %0 = memref.alloc() : memref<10xf32>
+  acc.yield %0 : memref<10xf32>
+} destroy {
+^bb0(%arg0: memref<10xf32>):
+  memref.dealloc %arg0 : memref<10xf32> 
+  acc.terminator
+}
+
 func.func @compute3(%a: memref<10x10xf32>, %b: memref<10x10xf32>, %c: memref<10xf32>, %d: memref<10xf32>) -> memref<10xf32> {
   %lb = arith.constant 0 : index
   %st = arith.constant 1 : index
   %c10 = arith.constant 10 : index
   %numGangs = arith.constant 10 : i64
   %numWorkers = arith.constant 10 : i64
 
   %pa = acc.present varPtr(%a : memref<10x10xf32>) -> memref<10x10xf32>
   %pb = acc.present varPtr(%b : memref<10x10xf32>) -> memref<10x10xf32>
   %pc = acc.present varPtr(%c : memref<10xf32>) -> memref<10xf32>
   %pd = acc.present varPtr(%d : memref<10xf32>) -> memref<10xf32>
   acc.data dataOperands(%pa, %pb, %pc, %pd: memref<10x10xf32>, memref<10x10xf32>, memref<10xf32>, memref<10xf32>) {
-    acc.parallel num_gangs(%numGangs: i64) num_workers(%numWorkers: i64) private(%c : memref<10xf32>) {
+    acc.parallel num_gangs(%numGangs: i64) num_workers(%numWorkers: i64) private(@privatization_memref_10_f32 -> %c : memref<10xf32>) {
       acc.loop gang {
         scf.for %x = %lb to %c10 step %st {
           acc.loop worker {
             scf.for %y = %lb to %c10 step %st {
               %axy = memref.load %a[%x, %y] : memref<10x10xf32>
               %bxy = memref.load %b[%x, %y] : memref<10x10xf32>
               %tmp = arith.addf %axy, %bxy : f32
               memref.store %tmp, %c[%y] : memref<10xf32>

 
 // CHECK:      func @compute3({{.*}}: memref<10x10xf32>, {{.*}}: memref<10x10xf32>, [[ARG2:%.*]]: memref<10xf32>, {{.*}}: memref<10xf32>) -> memref<10xf32> {
 // CHECK-NEXT:   [[C0:%.*]] = arith.constant 0 : index
 // CHECK-NEXT:   [[C1:%.*]] = arith.constant 1 : index
 // CHECK-NEXT:   [[C10:%.*]] = arith.constant 10 : index
 // CHECK-NEXT:   [[NUMGANG:%.*]] = arith.constant 10 : i64
 // CHECK-NEXT:   [[NUMWORKERS:%.*]] = arith.constant 10 : i64
 // CHECK:        acc.data dataOperands(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : memref<10x10xf32>, memref<10x10xf32>, memref<10xf32>, memref<10xf32>) {
-// CHECK-NEXT:     acc.parallel num_gangs([[NUMGANG]] : i64) num_workers([[NUMWORKERS]] : i64) private([[ARG2]] : memref<10xf32>) {
+// CHECK-NEXT:     acc.parallel num_gangs([[NUMGANG]] : i64) num_workers([[NUMWORKERS]] : i64) private(@privatization_memref_10_f32 -> [[ARG2]] : memref<10xf32>) {
 // CHECK-NEXT:       acc.loop gang {
 // CHECK-NEXT:         scf.for %{{.*}} = [[C0]] to [[C10]] step [[C1]] {
 // CHECK-NEXT:           acc.loop worker {
 // CHECK-NEXT:             scf.for %{{.*}} = [[C0]] to [[C10]] step [[C1]] {
 // CHECK-NEXT:               %{{.*}} = memref.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:               %{{.*}} = memref.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:               %{{.*}} = arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:               memref.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>

 // CHECK: acc.parallel
 // CHECK: acc.loop
 // CHECK-3: ^bb{{.*}}
 // CHECK: acc.yield
 // CHECK: acc.yield
 
 // -----
 
+acc.private.recipe @privatization_memref_10_f32 : memref<10xf32> init {
+^bb0(%arg0: memref<10xf32>):
+  %0 = memref.alloc() : memref<10xf32>
+  acc.yield %0 : memref<10xf32>
+} destroy {
+^bb0(%arg0: memref<10xf32>):
+  memref.dealloc %arg0 : memref<10xf32> 
+  acc.terminator
+}
+
+acc.private.recipe @privatization_memref_10_10_f32 : memref<10x10xf32> init {
+^bb0(%arg0: memref<10x10xf32>):
+  %0 = memref.alloc() : memref<10x10xf32>
+  acc.yield %0 : memref<10x10xf32>
+} destroy {
+^bb0(%arg0: memref<10x10xf32>):
+  memref.dealloc %arg0 : memref<10x10xf32> 
+  acc.terminator
+}
+
 func.func @testparallelop(%a: memref<10xf32>, %b: memref<10xf32>, %c: memref<10x10xf32>) -> () {
   %i64value = arith.constant 1 : i64
   %i32value = arith.constant 1 : i32
   %idxValue = arith.constant 1 : index
   acc.parallel async(%i64value: i64) {
   }
   acc.parallel async(%i32value: i32) {
   }
func.func @testparallelop(%a: memref<10xf32>, %b: memref<10xf32>, %c: memref<10x10xf32>) -> () {
   acc.parallel num_workers(%idxValue: index) {
   }
   acc.parallel vector_length(%i64value: i64) {
   }
   acc.parallel vector_length(%i32value: i32) {
   }
   acc.parallel vector_length(%idxValue: index) {
   }
-  acc.parallel private(%a, %c : memref<10xf32>, memref<10x10xf32>) firstprivate(%b: memref<10xf32>) {
+  acc.parallel private(@privatization_memref_10_f32 -> %a : memref<10xf32>, @privatization_memref_10_10_f32 -> %c : memref<10x10xf32>) firstprivate(%b: memref<10xf32>) {
   }
   acc.parallel {
   } attributes {defaultAttr = #acc<defaultvalue none>}
   acc.parallel {
   } attributes {defaultAttr = #acc<defaultvalue present>}
   acc.parallel {
   } attributes {asyncAttr}
   acc.parallel {

 // CHECK:      acc.parallel num_workers([[IDXVALUE]] : index) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel vector_length([[I64VALUE]] : i64) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel vector_length([[I32VALUE]] : i32) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel vector_length([[IDXVALUE]] : index) {
 // CHECK-NEXT: }
-// CHECK:      acc.parallel firstprivate([[ARGB]] : memref<10xf32>) private([[ARGA]], [[ARGC]] : memref<10xf32>, memref<10x10xf32>) {
+// CHECK:      acc.parallel firstprivate([[ARGB]] : memref<10xf32>) private(@privatization_memref_10_f32 -> [[ARGA]] : memref<10xf32>, @privatization_memref_10_10_f32 -> [[ARGC]] : memref<10x10xf32>) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel {
 // CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue none>}
 // CHECK:      acc.parallel {
 // CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue present>}
 // CHECK:      acc.parallel {
 // CHECK-NEXT: } attributes {asyncAttr}
 // CHECK:      acc.parallel {