[mlir][OpenMP] Added omp.atomic.capture operation

[shraiysh]

Summary

This patch supports the atomic construct (capture) following section 2.17.7 of OpenMP 5.0 standard. Also added tests for the same.

[peixin]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

726

For fortran code, expr here refers to both expr and expr_list.

730

The evaluation of `expr` need not be atomic w.r.t the read or write of the location designated by `x`.

This is for update clause. For capture clause, the standard states that Neither the evaluation of expr or expr_list, nor the write to the location designated by v, need be atomic with respect to the read or write of the location designated by x.

730

In general, type(x) must dereference to type(expr).

Not clear what this means.

732

Not only binary operation. Since this is used for both c/c++ and fortran, binop is the binary operator or intrinsic (MAX, MIN, IAND, IOR, or IEOR).

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1536–1537

The restriction states that If atomic-clause is update or not present then memory-order-clause must not be acq_rel or acquire, which is not for capture clause. One interesting thing is that the description in the standard states that If the write, update, or capture clause is specified and the release, acq_rel, or seq_cst clause is specified then the strong flush on entry to the atomic operation is also a release flush. Checked with OpenMP 5.1, and there is no this kind of restriction. Don't understand why acq_rel and acquire are not allowed in update or capture clause.

1549

Nit: How about x_type , v_type , expr_type?

1567

%x = %expr binop %x | %x = %x binop %expr would be clear than %x = ...

1601

The same as above.

1632

Nit: remove the brace.

1673

Remove the brace. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

mlir/test/Dialect/OpenMP/ops.mlir

601

lines 593-602 ->

v = x = x binop expr;
x = x binop expr; v = x;

607

v = x = expr binop x;
x = expr binop x; v = x;

619

lines 613-620 ->

v = x; x = x binop expr;
v = x; x binop= expr;

633

Redundant test cases. Move the comment fortran expressions in front of the above test cases.

[shraiysh]

Thanks for the detailed review @peixin. I will work on this patch and address comments based on the discussion from patch D116396.

[shraiysh]

Updated patch to use omp.atomic.read and omp.atomic.update

[shraiysh]

Depends on D116396

[shraiysh]

Fix missing semi-colon error

[peixin]

This design is much easier to review. Nice job. LGTM.

[kiranchandramohan]

For discussion,
-> Generally, in a single block region, if the terminator has no special semantics then it is better to use SingleBlockImplicitTerminator. In this case more so since it is just a couple of instructions.
-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

[peixin]

In D115851#3224776, @kiranchandramohan wrote:

-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

The first version of this patch model the capture construct using one standalone operation (https://reviews.llvm.org/D115851?id=394750), but it is kind of hard to read.

[kiranchandramohan]

In D115851#3224817, @peixin wrote:

In D115851#3224776, @kiranchandramohan wrote:

-> Would it be better to model to capture x, expr, v, y etc as operands of the capture operation?
-> Would it be better to model this is a complex operation without a region?

The first version of this patch model the capture construct using one standalone operation (https://reviews.llvm.org/D115851?id=394750), but it is kind of hard to read.

OK. The only concern I have is that the operands of the operation are not captured in the operation definition. I don't know whether that will cause any issues. I think we can go ahead with this for now. Can you address the terminator comment?

[shraiysh]

Thank you for the reviews @kiranchandramohan and @peixin. I will address the terminator comment soon. I have been busy with some other work recently. Apologies for the delay.

[shraiysh]

Addressed SingleBlockImplicitTerminator comment.

[shraiysh]

Gentle ping!

[kiranchandramohan]

1. If the terminator is implicit will it create difficulties in lowering?
2. Should we update the examples in the description and the tests to be without the terminator?

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1569

If the terminator is implicit, does this hold?

[shraiysh]

Thanks for the comments @kiranchandramohan. IIUC SingleBlockImplicitTerminator just adds an error message if the region does not end with omp.terminator. I thought it would accept the following code but it gives a missing omp.terminator error instead

omp.atomic.capture {
  omp.atomic.read ...
  omp.atomic.update ...
}

I think lowering should not be a problem because the in-memory structure still has an omp.terminator at the end of the region.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

1569

Yes, it still holds (there is a test for this error message). I expected that with an implicit terminator, it will only have two, but this is not true. It looks like having SingleBlockImplicitTerminator just makes sure that we end a region with omp.terminator.

[kiranchandramohan]

Ahh, sorry about that. I kind of thought the terminator will not be visible.

LGTM.

[shraiysh]

Thanks! Rebase with main.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

let arguments = (ins OpenMP_PointerLikeType:$x,
                        AtomicBinOpKindAttr:$binop,
                        DefaultValuedAttr<I64Attr, "0">:$hint,
                        OptionalAttr<MemoryOrderKindAttr>:$memory_order);
   let parser = [{ return parseAtomicUpdateOp(parser, result); }];
   let printer = [{ return printAtomicUpdateOp(p, *this); }];
   let verifier = [{ return verifyAtomicUpdateOp(*this); }];
 }
 
+def AtomicCaptureOp : OpenMP_Op<"atomic.capture",
+    [SingleBlockImplicitTerminator<"TerminatorOp">]> {
+  let summary = "performs an atomic capture";
+  let description = [{
+    This operation performs an atomic capture.
+
+    `hint` is the value of hint (as used in the hint clause). It is a compile
+    time constant. As the name suggests, this is just a hint for optimization.
+
+    `memory_order` indicates the memory ordering behavior of the construct. It
+    can be one of `seq_cst`, `acq_rel`, `release`, `acquire` or `relaxed`.
+
+    The region has the following allowed forms:
+
+    ```
+      omp.atomic.capture {
+        omp.atomic.update ...
+        omp.atomic.read ...
+        omp.terminator
+      }
+
+      omp.atomic.capture {
+        omp.atomic.read ...
+        omp.atomic.update ...
+        omp.terminator
+      }
+
+      omp.atomic.capture {
+        omp.atomic.read ...
+        omp.atomic.write ...
+        omp.terminator
+      }
+    ```
+
+  }];
+
+  let arguments = (ins DefaultValuedAttr<I64Attr, "0">:$hint,
+                       OptionalAttr<MemoryOrderKind>:$memory_order);
+  let regions = (region SizedRegion<1>:$region);
+  let parser = [{ return parseAtomicCaptureOp(parser, result); }];
+  let printer = [{ return printAtomicCaptureOp(p, *this); }];
+  let verifier = [{ return verifyAtomicCaptureOp(*this); }];
+}
+
 //===----------------------------------------------------------------------===//
 // 2.19.5.7 declare reduction Directive
 //===----------------------------------------------------------------------===//
 
 def ReductionDeclareOp : OpenMP_Op<"reduction.declare", [Symbol]> {
   let summary = "declares a reduction kind";
 
   let description = [{

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 }
 
 /// Verifier for AtomicUpdateOp
 static LogicalResult verifyAtomicUpdateOp(AtomicUpdateOp op) {
   if (auto mo = op.memory_order()) {
     if (*mo == ClauseMemoryOrderKind::acq_rel ||
         *mo == ClauseMemoryOrderKind::acquire) {
       return op.emitError(
           "memory-order must not be acq_rel or acquire for atomic updates");
     }
   }
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// AtomicCaptureOp
+//===----------------------------------------------------------------------===//
+
+/// Parser for AtomicCaptureOp
+static LogicalResult parseAtomicCaptureOp(OpAsmParser &parser,
+                                          OperationState &result) {
+  SmallVector<ClauseType> clauses = {memoryOrderClause, hintClause};
+  SmallVector<int> segments;
+  if (parseClauses(parser, result, clauses, segments) ||
+      parser.parseRegion(*result.addRegion()))
+    return failure();
+  return success();
+}
+
+/// Printer for AtomicCaptureOp
+static void printAtomicCaptureOp(OpAsmPrinter &p, AtomicCaptureOp op) {
+  if (op.memory_order())
+    p << "memory_order(" << op.memory_order() << ") ";
+  if (op.hintAttr())
+    printSynchronizationHint(p, op, op.hintAttr());
+  p.printRegion(op.region());
+}
+
+/// Verifier for AtomicCaptureOp
+static LogicalResult verifyAtomicCaptureOp(AtomicCaptureOp op) {
+  Block::OpListType &ops = op.region().front().getOperations();
+  if (ops.size() != 3)
+    return emitError(op.getLoc())
+           << "expected three operations in omp.atomic.capture region (one "
+              "terminator, and two atomic ops)";
+  auto &firstOp = ops.front();
+  auto &secondOp = *ops.getNextNode(firstOp);
+  auto firstReadStmt = dyn_cast<AtomicReadOp>(firstOp);
+  auto firstUpdateStmt = dyn_cast<AtomicUpdateOp>(firstOp);
+  auto secondReadStmt = dyn_cast<AtomicReadOp>(secondOp);
+  auto secondUpdateStmt = dyn_cast<AtomicUpdateOp>(secondOp);
+  auto secondWriteStmt = dyn_cast<AtomicWriteOp>(secondOp);
+
+  if (!((firstUpdateStmt && secondReadStmt) ||
+        (firstReadStmt && secondUpdateStmt) ||
+        (firstReadStmt && secondWriteStmt)))
+    return emitError(ops.front().getLoc())
+           << "invalid sequence of operations in the capture region";
+  if (firstUpdateStmt && secondReadStmt &&
+      firstUpdateStmt.x() != secondReadStmt.x())
+    return emitError(firstUpdateStmt.getLoc())
+           << "updated variable in omp.atomic.update must be captured in "
+              "second operation";
+  if (firstReadStmt && secondUpdateStmt &&
+      firstReadStmt.x() != secondUpdateStmt.x())
+    return emitError(firstReadStmt.getLoc())
+           << "captured variable in omp.atomic.read must be updated in second "
+              "operation";
+  if (firstReadStmt && secondWriteStmt &&
+      firstReadStmt.x() != secondWriteStmt.address())
+    return emitError(firstReadStmt.getLoc())
+           << "captured variable in omp.atomic.read must be updated in "
+              "second operation";
+  return success();
+}
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/test/Dialect/OpenMP/invalid.mlir

 func @omp_atomic_update5(%x: memref<i32>, %expr: i32) {
   // expected-error @below {{use of value '%x' expects different type than prior uses: 'i32' vs 'memref<i32>'}}
   omp.atomic.update %x = %x add %expr : i32, memref<i32>
   return
 }
 
 // -----
 
+func @omp_atomic_capture(%x: memref<i32>, %v: memref<i32>, %expr: i32) {
+  // expected-error @below {{expected three operations in omp.atomic.capture region}}
+  omp.atomic.capture {
+    omp.atomic.read %v = %x : memref<i32>
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.read %v = %x : memref<i32>
+    omp.atomic.read %v = %x : memref<i32>
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.update %x = %x sub %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+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}}
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.atomic.read %v = %x : memref<i32>
+    omp.terminator
+  }
+  return
+}
+
+// -----
+
+func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
+  omp.atomic.capture {
+    // expected-error @below {{updated variable in omp.atomic.update must be captured in second operation}}
+    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.atomic.read %v = %y : memref<i32>
+    omp.terminator
+  }
+}
+
+// -----
+
+func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
+  omp.atomic.capture {
+    // expected-error @below {{captured variable in omp.atomic.read must be updated in second operation}}
+    omp.atomic.read %v = %y : memref<i32>
+    omp.atomic.update %x = %x add %expr : memref<i32>, i32
+    omp.terminator
+  }
+}
+
+// -----
+
+func @omp_atomic_capture(%x: memref<i32>, %y: memref<i32>, %v: memref<i32>, %expr: i32) {
+  omp.atomic.capture {
+    // expected-error @below {{captured variable in omp.atomic.read must be updated in second operation}}
+    omp.atomic.read %v = %x : memref<i32>
+    omp.atomic.write %y = %expr : memref<i32>, i32
+    omp.terminator
+  }
+}
+
+// -----
+
 func @omp_sections(%data_var1 : memref<i32>, %data_var2 : memref<i32>, %data_var3 : memref<i32>) -> () {
   // expected-error @below {{operand used in both private and firstprivate clauses}}
   omp.sections private(%data_var1 : memref<i32>) firstprivate(%data_var1 : memref<i32>) {
     omp.terminator
   }
   return
 }
 

mlir/test/Dialect/OpenMP/ops.mlir

func @omp_atomic_update(%x : memref<i32>, %expr : i32, %xBool : memref<i1>, %exprBool : i1) {
   omp.atomic.update %xBool =  %exprBool eqv %xBool : memref<i1>, i1
   // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] neqv %[[XBOOL]] : memref<i1>, i1
   omp.atomic.update %xBool =  %exprBool neqv %xBool : memref<i1>, i1
   // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] add %[[X]] memory_order(seq_cst) hint(speculative) : memref<i32>, i32
   omp.atomic.update %x = %expr add %x hint(speculative) memory_order(seq_cst) : memref<i32>, i32
   return
 }
 
+// CHECK-LABEL: omp_atomic_capture
+// CHECK-SAME: (%[[v:.*]]: memref<i32>, %[[x:.*]]: memref<i32>, %[[expr:.*]]: i32)
+func @omp_atomic_capture(%v: memref<i32>, %x: memref<i32>, %expr: i32) {
+  // CHECK: omp.atomic.capture{
+  // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
+  // CHECK-NEXT: omp.terminator
+  // CHECK-NEXT: }
+  omp.atomic.capture{
+    omp.atomic.update %x = %expr add %x : memref<i32>, i32
+    omp.atomic.read %v = %x : memref<i32>
+    omp.terminator
+  }
+  // CHECK: omp.atomic.capture{
+  // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
+  // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref<i32>, i32
+  // CHECK-NEXT: omp.terminator
+  // CHECK-NEXT: }
+  omp.atomic.capture{
+    omp.atomic.read %v = %x : memref<i32>
+    omp.atomic.update %x = %expr add %x : memref<i32>, i32
+    omp.terminator
+  }
+  // CHECK: omp.atomic.capture{
+  // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref<i32>
+  // CHECK-NEXT: omp.atomic.write %[[x]] = %[[expr]] : memref<i32>, i32
+  // CHECK-NEXT: omp.terminator
+  // CHECK-NEXT: }
+  omp.atomic.capture{
+    omp.atomic.read %v = %x : memref<i32>
+    omp.atomic.write %x = %expr : memref<i32>, i32
+    omp.terminator
+  }
+  return
+}
+
 // CHECK-LABEL: omp_sectionsop
 func @omp_sectionsop(%data_var1 : memref<i32>, %data_var2 : memref<i32>,
                      %data_var3 : memref<i32>, %redn_var : !llvm.ptr<f32>) {
 
   // CHECK: omp.sections private(%{{.*}} : memref<i32>) {
   "omp.sections" (%data_var1) ({
     // CHECK: omp.terminator
     omp.terminator
   }) {operand_segment_sizes = dense<[1,0,0,0,0,0]> : vector<6xi32>} : (memref<i32>) -> ()
 
   // CHECK: omp.sections firstprivate(%{{.*}} : memref<i32>) {
   "omp.sections" (%data_var1) ({
     // CHECK: omp.terminator
     omp.terminator
   }) {operand_segment_sizes = dense<[0,1,0,0,0,0]> : vector<6xi32>} : (memref<i32>) -> ()
 
   // CHECK: omp.sections lastprivate(%{{.*}} : memref<i32>) {
   "omp.sections" (%data_var1) ({
     // CHECK: omp.terminator