[mlir][OpenACC][OpenMP] Modify atomic capture to allow update/write

OpenACC's C++ version has a variant of capture that permits a update
followed by a write. Therefore the verifier was overly strict in this
case.

According to our reading of the OpenMP 6.0 spec, it appears that `atomic
captured update` (page 495) also requires this form, so it seems
reasonable to allow this for both languages.

Author: erichkeane

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

@@ -2999,6 +2999,12 @@ def AtomicCaptureOp : OpenACC_Op<"atomic.capture",
         acc.atomic.write ...
         acc.terminator
       }
+
+      acc.atomic.capture {
+        acc.atomic.update ...
+        acc.atomic.write ...
+        acc.terminator
+      }
     ```
 
   }];

mlir/include/mlir/Dialect/OpenACCMPCommon/Interfaces/AtomicInterfaces.td

@@ -239,6 +239,7 @@ def AtomicCaptureOpInterface : OpInterface<"AtomicCaptureOpInterface"> {
     implement one of the atomic interfaces. It can be found in one of these
     forms:
       `{ atomic.update, atomic.read }`
+      `{ atomic.update, atomic.write }`
       `{ atomic.read, atomic.update }`
       `{ atomic.read, atomic.write }`
   }];
@@ -291,12 +292,15 @@ def AtomicCaptureOpInterface : OpInterface<"AtomicCaptureOpInterface"> {
         auto secondWriteStmt = dyn_cast<AtomicWriteOpInterface>(secondOp);
 
         if (!((firstUpdateStmt && secondReadStmt) ||
+              (firstUpdateStmt && secondWriteStmt) ||
               (firstReadStmt && secondUpdateStmt) ||
               (firstReadStmt && secondWriteStmt)))
           return ops.front().emitError()
                 << "invalid sequence of operations in the capture region";
-        if (firstUpdateStmt && secondReadStmt &&
-            firstUpdateStmt.getX() != secondReadStmt.getX())
+        if ((firstUpdateStmt && secondReadStmt &&
+             firstUpdateStmt.getX() != secondReadStmt.getX()) ||
+            (firstUpdateStmt && secondWriteStmt &&
+             firstUpdateStmt.getX() != secondWriteStmt.getX()))
           return firstUpdateStmt.emitError()
                 << "updated variable in atomic.update must be captured in "
                     "second operation";

mlir/test/Dialect/OpenACC/invalid.mlir

@@ -690,7 +690,6 @@ func.func @acc_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
 
 func.func @acc_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   acc.atomic.capture {
-    // expected-error @below {{invalid sequence of operations in the capture region}}
     acc.atomic.update %x : memref<i32> {
     ^bb0(%xval: i32):
       %newval = llvm.add %xval, %expr : i32
@@ -704,6 +703,23 @@ func.func @acc_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
 
 // -----
 
+func.func @acc_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
+  acc.atomic.capture {
+    // expected-error @below {{updated variable in atomic.update must be captured in second operation}}
+    acc.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      acc.yield %newval : i32
+    }
+    acc.atomic.write %v = %expr : memref<i32>, i32
+
+    acc.terminator
+  }
+  return
+}
+
+// -----
+
 func.func @acc_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   acc.atomic.capture {
     // expected-error @below {{invalid sequence of operations in the capture region}}

mlir/test/Dialect/OpenMP/invalid.mlir

@@ -1263,7 +1263,22 @@ func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
 
 func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
-    // expected-error @below {{invalid sequence of operations in the capture region}}
+    // expected-error @below {{updated variable in atomic.update must be captured in second operation}}
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
+    omp.atomic.write %v = %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
+  omp.atomic.capture {
     omp.atomic.update %x : memref<i32> {
     ^bb0(%xval: i32):
       %newval = llvm.add %xval, %expr : i32
@@ -1289,6 +1304,22 @@ func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
 
 // -----
 
+func.func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
+  omp.atomic.capture {
+    omp.atomic.update %x : memref<i32> {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield (%newval : i32)
+    }
+    omp.atomic.write %x = %expr : memref<i32>, i32
+
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
 func.func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
   omp.atomic.capture {
     // expected-error @below {{updated variable in atomic.update must be captured in second operation}}