[Mlir-commits] [mlir] 61807f5 - [mlir][openacc] Add wait attribute and variadic operand

[Valentin Clement]

Author: Valentin Clement
Date: 2023-06-30T10:04:49-07:00
New Revision: 61807f5c295814207f71aebf6438b9c00aedce9f

URL: https://github.com/llvm/llvm-project/commit/61807f5c295814207f71aebf6438b9c00aedce9f
DIFF: https://github.com/llvm/llvm-project/commit/61807f5c295814207f71aebf6438b9c00aedce9f.diff

LOG: [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

Added: 
    

Modified: 
    flang/lib/Lower/OpenACC.cpp
    mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
    mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
    mlir/test/Dialect/OpenACC/ops.mlir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index 2e32abcc420498..4313413065be34 100644
--- 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 5960dfadbc44f2..a1af1cf7fea895 100644
--- 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 2ac6899ae7ab8d..9c6cffa8399dcb 100644
--- 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 a1f94323864d60..fa2f8c9f90ca25 100644
--- 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>) -> () {