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

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

Reviewed By: razvanlupusoru

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

diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index 2e32abcc4204985fd1fefd0587bb664b193a6c97..4313413065be34323ac2a2a25d6e82134af3449f 100644 (file)
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/flang/lib/Lower/OpenACC.cpp
@@ -1353,19 +1353,20 @@ static void genACCDataOp(Fortran::lower::AbstractConverter &converter,
                          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);
@@ -1450,12 +1451,15 @@ static void genACCDataOp(Fortran::lower::AbstractConverter &converter,
   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);

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 5960dfadbc44f28dd1e1c9310cf572c317d21cb6..a1af1cf7fea8958de41b2167c3447ad8033e011a 100644 (file)
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -867,6 +867,9 @@ def OpenACC_DataOp : OpenACC_Op<"data",
   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);
 
@@ -885,6 +888,8 @@ def OpenACC_DataOp : OpenACC_Op<"data",
         `if` `(` $ifCond `)`
       | `async` `(` $async `:` type($async) `)`
       | `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
+      | `wait_devnum` `(` $waitDevnum `:` type($waitDevnum) `)`
+      | `wait` `(` $waitOperands `:` type($waitOperands) `)`
     )
     $region attr-dict-with-keyword
   }];

diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 2ac6899ae7ab8d0dfe36f90d05f970c2418fceec..9c6cffa8399dcb7e45a103fab591caae25cc4890 100644 (file)
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -868,6 +868,7 @@ 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);
 }
 

diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index a1f94323864d6040067cbdb7844ade44a656b8af..fa2f8c9f90ca25070a0771a6a8583913c0401e40 100644 (file)
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/mlir/test/Dialect/OpenACC/ops.mlir
@@ -822,6 +822,17 @@ func.func @testdataop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
   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
 }
 
@@ -927,6 +938,15 @@ func.func @testdataop(%a: memref<f32>, %b: memref<f32>, %c: memref<f32>) -> () {
 // 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>) -> () {