[mlir][openacc] Add async attribute and optional operand

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

Reviewed By: razvanlupusoru

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

flang/lib/Lower/OpenACC.cpp

@@ -1353,10 +1353,15 @@ static void genACCDataOp(Fortran::lower::AbstractConverter &converter,
                          Fortran::semantics::SemanticsContext &semanticsContext,
                          Fortran::lower::StatementContext &stmtCtx,
                          const Fortran::parser::AccClauseList &accClauseList) {
-  mlir::Value ifCond;
+  mlir::Value ifCond, async;
   llvm::SmallVector<mlir::Value> attachEntryOperands, createEntryOperands,
       copyEntryOperands, copyoutEntryOperands, dataClauseOperands;
 
+  // Async has 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;
+
   fir::FirOpBuilder &builder = converter.getFirOpBuilder();
 
   // Lower clauses values mapped to operands.
@@ -1444,11 +1449,14 @@ static void genACCDataOp(Fortran::lower::AbstractConverter &converter,
   llvm::SmallVector<mlir::Value> operands;
   llvm::SmallVector<int32_t> operandSegments;
   addOperand(operands, operandSegments, ifCond);
+  addOperand(operands, operandSegments, async);
   addOperands(operands, operandSegments, dataClauseOperands);
 
   auto dataOp = createRegionOp<mlir::acc::DataOp, mlir::acc::TerminatorOp>(
       builder, currentLocation, operands, operandSegments);
 
+  dataOp.setAsyncAttr(addAsyncAttr);
+
   auto insPt = builder.saveInsertionPoint();
   builder.setInsertionPointAfter(dataOp);
 

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

@@ -865,6 +865,8 @@ def OpenACC_DataOp : OpenACC_Op<"data",
 
 
   let arguments = (ins Optional<I1>:$ifCond,
+                       Optional<IntOrIndex>:$async,
+                       UnitAttr:$asyncAttr,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$dataClauseOperands,
                        OptionalAttr<DefaultValueAttr>:$defaultAttr);
 
@@ -881,6 +883,7 @@ def OpenACC_DataOp : OpenACC_Op<"data",
   let assemblyFormat = [{
     oilist(
         `if` `(` $ifCond `)`
+      | `async` `(` $async `:` type($async) `)`
       | `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
     )
     $region attr-dict-with-keyword

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

@@ -867,6 +867,7 @@ unsigned DataOp::getNumDataOperands() { return getDataClauseOperands().size(); }
 
 Value DataOp::getDataOperand(unsigned i) {
   unsigned numOptional = getIfCond() ? 1 : 0;
+  numOptional += getAsync() ? 1 : 0;
   return getOperand(numOptional + i);
 }
 

mlir/test/Dialect/OpenACC/ops.mlir

@@ -814,6 +814,14 @@ func.func @testdataop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
 
   acc.data {
   } attributes { defaultAttr = #acc<defaultvalue none> }
+
+  acc.data {
+  } attributes { defaultAttr = #acc<defaultvalue none>, async }
+
+  %a1 = arith.constant 1 : i64
+  acc.data async(%a1 : i64) {
+  } attributes { defaultAttr = #acc<defaultvalue none>, async }
+
   return
 }
 
@@ -913,6 +921,12 @@ func.func @testdataop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
 // CHECK:      acc.data {
 // 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>}
+
 // -----
 
 func.func @testupdateop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {