[mlir][openacc] Add wait attribute and variadic operand

[clementval]

Summary

OpenACC 3.2 allowed the wait clause to the data construct. This patch
adds a unit attribute and a variadic operand to the data operation to model
the wait clause in a similar way it was added to other data operation.
The attribute models the presence of the clause without any argument. When
arguments are provided they are placed in the wait operand.

Depends on D154111

[clementval]

Add wiatDevnum

[clementval]

Add missing waitDevnum

[razvanlupusoru]

Looks great. Thank you!

flang/lib/Lower/OpenACC.cpp

createComputeOp(Fortran::lower::AbstractConverter &converter,
   return computeOp;
 }
 
 static void genACCDataOp(Fortran::lower::AbstractConverter &converter,
                          mlir::Location currentLocation,
                          Fortran::semantics::SemanticsContext &semanticsContext,
                          Fortran::lower::StatementContext &stmtCtx,
                          const Fortran::parser::AccClauseList &accClauseList) {
-  mlir::Value ifCond, async;
+  mlir::Value ifCond, async, waitDevnum;
   llvm::SmallVector<mlir::Value> attachEntryOperands, createEntryOperands,
-      copyEntryOperands, copyoutEntryOperands, dataClauseOperands;
+      copyEntryOperands, copyoutEntryOperands, dataClauseOperands, waitOperands;
 
-  // Async has an optional value but can be present with
+  // Async and wait have an optional value but can be present with
   // no value as well. When there is no value, the op has an attribute to
   // represent the clause.
   bool addAsyncAttr = false;
+  bool addWaitAttr = false;
 
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
 
   // Lower clauses values mapped to operands.
-  // Keep track of each group of operands separatly as clauses can appear
+  // Keep track of each group of operands separately as clauses can appear
   // more than once.
   for (const Fortran::parser::AccClause &clause : accClauseList.v) {
     mlir::Location clauseLocation = converter.genLocation(clause.source);
     if (const auto *ifClause =
             std::get_if<Fortran::parser::AccClause::If>(&clause.u)) {
       genIfClause(converter, clauseLocation, ifClause, ifCond, stmtCtx);
     } else if (const auto *copyClause =
                    std::get_if<Fortran::parser::AccClause::Copy>(&clause.u)) {
for (const Fortran::parser::AccClause &clause : accClauseList.v) {
     }
   }
 
   // Prepare the operand segment size attribute and the operands value range.
   llvm::SmallVector<mlir::Value> operands;
   llvm::SmallVector<int32_t> operandSegments;
   addOperand(operands, operandSegments, ifCond);
   addOperand(operands, operandSegments, async);
+  addOperand(operands, operandSegments, waitDevnum);
+  addOperands(operands, operandSegments, waitOperands);
   addOperands(operands, operandSegments, dataClauseOperands);
 
   auto dataOp = createRegionOp<mlir::acc::DataOp, mlir::acc::TerminatorOp>(
       builder, currentLocation, operands, operandSegments);
 
   dataOp.setAsyncAttr(addAsyncAttr);
+  dataOp.setAsyncAttr(addWaitAttr);
 
   auto insPt = builder.saveInsertionPoint();
   builder.setInsertionPointAfter(dataOp);
 
   // Create the exit operations after the region.
   genDataExitOperations<mlir::acc::CopyinOp, mlir::acc::CopyoutOp>(
       builder, copyEntryOperands, /*structured=*/true, /*implicit=*/false);
   genDataExitOperations<mlir::acc::CreateOp, mlir::acc::CopyoutOp>(

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

let description = [{
     }
     ```
   }];
 
 
   let arguments = (ins Optional<I1>:$ifCond,
                        Optional<IntOrIndex>:$async,
                        UnitAttr:$asyncAttr,
+                       Optional<IntOrIndex>:$waitDevnum,
+                       Variadic<IntOrIndex>:$waitOperands,
+                       UnitAttr:$waitAttr,
                        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(
         `if` `(` $ifCond `)`
       | `async` `(` $async `:` type($async) `)`
       | `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
+      | `wait_devnum` `(` $waitDevnum `:` type($waitDevnum) `)`
+      | `wait` `(` $waitOperands `:` type($waitOperands) `)`
     )
     $region attr-dict-with-keyword
   }];
   let hasVerifier = 1;
 }
 
 def OpenACC_TerminatorOp : OpenACC_Op<"terminator", [Terminator]> {
   let summary = "Generic terminator for OpenACC regions";

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

LogicalResult acc::DataOp::verify() {
   return success();
 }
 
 unsigned DataOp::getNumDataOperands() { return getDataClauseOperands().size(); }
 
 Value DataOp::getDataOperand(unsigned i) {
   unsigned numOptional = getIfCond() ? 1 : 0;
   numOptional += getAsync() ? 1 : 0;
+  numOptional += getWaitOperands().size();
   return getOperand(numOptional + i);
 }
 
 //===----------------------------------------------------------------------===//
 // ExitDataOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::ExitDataOp::verify() {

mlir/test/Dialect/OpenACC/ops.mlir

func.func @testdataop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
 
   acc.data {
   } attributes { defaultAttr = #acc<defaultvalue none>, async }
 
   %a1 = arith.constant 1 : i64
   acc.data async(%a1 : i64) {
   } attributes { defaultAttr = #acc<defaultvalue none>, async }
 
+  acc.data {
+  } attributes { defaultAttr = #acc<defaultvalue none>, wait }
+
+  %w1 = arith.constant 1 : i64
+  acc.data wait(%w1 : i64) {
+  } attributes { defaultAttr = #acc<defaultvalue none>, wait }
+
+  %wd1 = arith.constant 1 : i64
+  acc.data wait_devnum(%wd1 : i64) wait(%w1 : i64) {
+  } attributes { defaultAttr = #acc<defaultvalue none>, wait }
+
   return
 }
 
 // CHECK:      func @testdataop(%[[ARGA:.*]]: memref<f32>, %[[ARGB:.*]]: memref<f32>, %[[ARGC:.*]]: memref<f32>) {
 
 // CHECK:      %[[IFCOND1:.*]] = arith.constant true
 // CHECK:      %[[PRESENT_A:.*]] = acc.present varPtr(%[[ARGA]] : memref<f32>) -> memref<f32>
 // CHECK:      acc.data if(%[[IFCOND1]]) dataOperands(%[[PRESENT_A]] : memref<f32>) {

 // CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue none>}
 
 // CHECK:      acc.data {
 // CHECK-NEXT: } attributes {async, defaultAttr = #acc<defaultvalue none>}
 
 // CHECK:      acc.data async(%{{.*}} : i64) {
 // CHECK-NEXT: } attributes {async, defaultAttr = #acc<defaultvalue none>}
 
+// CHECK:      acc.data {
+// CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue none>, wait}
+
+// CHECK:      acc.data wait(%{{.*}} : i64) {
+// CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue none>, wait}
+
+// CHECK:      acc.data wait_devnum(%{{.*}} : i64) wait(%{{.*}} : i64) {
+// CHECK-NEXT: } attributes {defaultAttr = #acc<defaultvalue none>, wait}
+
 // -----
 
 func.func @testupdateop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
   %i64Value = arith.constant 1 : i64
   %i32Value = arith.constant 1 : i32
   %idxValue = arith.constant 1 : index
   %ifCond = arith.constant true
   %0 = acc.update_device varPtr(%a : memref<f32>) -> memref<f32>