From ecc7adc46ce926d51022ff5ab5ee60eb61b80972 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval@gmail.com>
Date: Fri, 30 Jun 2023 10:03:24 -0700
Subject: [PATCH] [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                     | 10 +++++++++-
 mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td |  3 +++
 mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp         |  1 +
 mlir/test/Dialect/OpenACC/ops.mlir              | 14 ++++++++++++++
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index 29f6d85..2e32abc 100644
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/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);
 
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 076faa7..5960dfa 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/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
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 77de452..2ac6899 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/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);
 }
 
diff --git a/mlir/test/Dialect/OpenACC/ops.mlir b/mlir/test/Dialect/OpenACC/ops.mlir
index e07ab8c..a1f9432 100644
--- a/mlir/test/Dialect/OpenACC/ops.mlir
+++ b/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>) -> () {
-- 
2.7.4