diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td --- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td +++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td @@ -338,14 +338,12 @@ def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator, - ParentOneOf<["WsLoopOp", "ReductionDeclareOp"]>]> { + ParentOneOf<["WsLoopOp", "ReductionDeclareOp", "AtomicUpdateOp"]>]> { let summary = "loop yield and termination operation"; let description = [{ "omp.yield" yields SSA values from the OpenMP dialect op region and terminates the region. The semantics of how the values are yielded is defined by the parent operation. - If "omp.yield" has any operands, the operands must match the parent - operation's results. }]; let arguments = (ins Variadic:$results); @@ -663,39 +661,35 @@ let symbolToStringFnName = "AtomicBinOpKindToString"; } -def AtomicUpdateOp : OpenMP_Op<"atomic.update"> { +def AtomicUpdateOp : OpenMP_Op<"atomic.update", + [SingleBlockImplicitTerminator<"YieldOp">]> { let summary = "performs an atomic update"; let description = [{ This operation performs an atomic update. - The operands `x` and `expr` are exactly the same as the operands `x` and - `expr` in the OpenMP Standard. The operand `x` is the address of the - variable that is being updated. `x` is atomically read/written. The - evaluation of `expr` need not be atomic w.r.t the read or write of the - location designated by `x`. In general, type(x) must dereference to - type(expr). - - The attribute `isXBinopExpr` is - - true when the expression is of the form `x binop expr` on RHS - - false when the expression is of the form `expr binop x` on RHS - - The attribute `binop` is the binary operation being performed atomically. + The operand `x` is exactly the same as the operand `x` in the OpenMP + Standard (OpenMP 5.0, section 2.17.7). It is the address of the variable + that is being updated. `x` is atomically read/written. `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 describes how to update the value of `x`. It takes the value at + `x` as an input and must yield the updated value. Only the update to `x` is + atomic. Generally the region must have only one instruction, but can + potentially have more than one instructions too. The update is sematically + similar to a compare-exchange loop based atomic update. }]; let arguments = (ins OpenMP_PointerLikeType:$x, - AnyType:$expr, - UnitAttr:$isXBinopExpr, - AtomicBinOpKindAttr:$binop, DefaultValuedAttr:$hint, OptionalAttr:$memory_order); + let regions = (region SizedRegion<1>:$region); let hasCustomAssemblyFormat = 1; let hasVerifier = 1; } @@ -769,6 +763,9 @@ Note that the MLIR type system does not allow for type-polymorphic reductions. Separate reduction declarations should be created for different element and accumulator types. + + For initializer and reduction regions, the operand to `omp.yield` must + match the parent operation's results. }]; let arguments = (ins SymbolNameAttr:$sym_name, diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -1528,63 +1528,30 @@ /// Parser for AtomicUpdateOp /// -/// operation ::= `omp.atomic.update` atomic-clause-list region +/// operation ::= `omp.atomic.update` atomic-clause-list ssa-id-and-type region ParseResult AtomicUpdateOp::parse(OpAsmParser &parser, OperationState &result) { SmallVector clauses = {memoryOrderClause, hintClause}; SmallVector segments; - OpAsmParser::OperandType x, y, z; + OpAsmParser::OperandType x, expr; Type xType, exprType; - StringRef binOp; - - // x = y `op` z : xtype, exprtype - if (parser.parseOperand(x) || parser.parseEqual() || parser.parseOperand(y) || - parser.parseKeyword(&binOp) || parser.parseOperand(z) || - parseClauses(parser, result, clauses, segments) || parser.parseColon() || - parser.parseType(xType) || parser.parseComma() || - parser.parseType(exprType) || - parser.resolveOperand(x, xType, result.operands)) { - return failure(); - } - auto binOpEnum = AtomicBinOpKindToEnum(binOp.upper()); - if (!binOpEnum) - return parser.emitError(parser.getNameLoc()) - << "invalid atomic bin op in atomic update\n"; - auto attr = - parser.getBuilder().getI64IntegerAttr((int64_t)binOpEnum.getValue()); - result.addAttribute("binop", attr); - - OpAsmParser::OperandType expr; - if (x.name == y.name && x.number == y.number) { - expr = z; - result.addAttribute("isXBinopExpr", parser.getBuilder().getUnitAttr()); - } else if (x.name == z.name && x.number == z.number) { - expr = y; - } else { - return parser.emitError(parser.getNameLoc()) - << "atomic update variable " << x.name - << " not found in the RHS of the assignment statement in an" - " atomic.update operation"; - } - return parser.resolveOperand(expr, exprType, result.operands); + if (parseClauses(parser, result, clauses, segments) || + parser.parseOperand(x) || parser.parseColon() || + parser.parseType(xType) || + parser.resolveOperand(x, xType, result.operands) || + parser.parseRegion(*result.addRegion())) + return failure(); + return success(); } void AtomicUpdateOp::print(OpAsmPrinter &p) { - p << " " << x() << " = "; - Value y, z; - if (isXBinopExpr()) { - y = x(); - z = expr(); - } else { - y = expr(); - z = x(); - } - p << y << " " << AtomicBinOpKindToString(binop()).lower() << " " << z << " "; + p << " "; if (auto mo = memory_order()) p << "memory_order(" << stringifyClauseMemoryOrderKind(*mo) << ") "; if (hintAttr()) printSynchronizationHint(p, *this, hintAttr()); - p << ": " << x().getType() << ", " << expr().getType(); + p << x() << " : " << x().getType(); + p.printRegion(region()); } /// Verifier for AtomicUpdateOp @@ -1596,6 +1563,22 @@ "memory-order must not be acq_rel or acquire for atomic updates"); } } + + if (region().getNumArguments() != 1) + return emitError("the region must accept exactly one argument"); + + if (x().getType().cast().getElementType() != + region().getArgument(0).getType()) { + return emitError("the type of the operand must be a pointer type whose " + "element type is the same as that of the region argument"); + } + + YieldOp yieldOp = *region().getOps().begin(); + + if (yieldOp.results().size() != 1) + return emitError("only updated value must be returned"); + if (yieldOp.results().front().getType() != region().getArgument(0).getType()) + return emitError("input and yielded value must have the same type"); return success(); } diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir --- a/mlir/test/Dialect/OpenMP/invalid.mlir +++ b/mlir/test/Dialect/OpenMP/invalid.mlir @@ -609,17 +609,26 @@ // ----- -func @omp_atomic_update1(%x: memref, %expr: i32, %foo: memref) { - // expected-error @below {{atomic update variable %x not found in the RHS of the assignment statement in an atomic.update operation}} - omp.atomic.update %x = %foo add %expr : memref, i32 +func @omp_atomic_update1(%x: memref, %expr: f32) { + // expected-error @below {{the type of the operand must be a pointer type whose element type is the same as that of the region argument}} + omp.atomic.update %x : memref { + ^bb0(%xval: f32): + %newval = llvm.add %xval, %expr : f32 + omp.yield (%newval : f32) + } return } // ----- func @omp_atomic_update2(%x: memref, %expr: i32) { - // expected-error @below {{invalid atomic bin op in atomic update}} - omp.atomic.update %x = %x invalid %expr : memref, i32 + // expected-error @+2 {{op expects regions to end with 'omp.yield', found 'omp.terminator'}} + // expected-note @below {{in custom textual format, the absence of terminator implies 'omp.yield'}} + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.terminator + } return } @@ -627,7 +636,11 @@ func @omp_atomic_update3(%x: memref, %expr: i32) { // expected-error @below {{memory-order must not be acq_rel or acquire for atomic updates}} - omp.atomic.update %x = %x add %expr memory_order(acq_rel) : memref, i32 + omp.atomic.update memory_order(acq_rel) %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } return } @@ -635,7 +648,11 @@ func @omp_atomic_update4(%x: memref, %expr: i32) { // expected-error @below {{memory-order must not be acq_rel or acquire for atomic updates}} - omp.atomic.update %x = %x add %expr memory_order(acquire) : memref, i32 + omp.atomic.update memory_order(acquire) %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } return } @@ -644,7 +661,35 @@ // expected-note @below {{prior use here}} func @omp_atomic_update5(%x: memref, %expr: i32) { // expected-error @below {{use of value '%x' expects different type than prior uses: 'i32' vs 'memref'}} - omp.atomic.update %x = %x add %expr : i32, memref + omp.atomic.update %x : i32 { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } + return +} + +// ----- + +func @omp_atomic_update6(%x: memref, %expr: i32) { + // expected-error @below {{only updated value must be returned}} + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval, %expr : i32, i32) + } + return +} + +// ----- + +func @omp_atomic_update7(%x: memref, %expr: i32, %y: f32) { + // expected-error @below {{input and yielded value must have the same type}} + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%y: f32) + } return } @@ -676,8 +721,16 @@ func @omp_atomic_capture(%x: memref, %v: memref, %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 - omp.atomic.update %x = %x sub %expr : memref, i32 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } omp.terminator } return @@ -701,7 +754,11 @@ omp.atomic.capture { // expected-error @below {{invalid sequence of operations in the capture region}} omp.atomic.write %x = %expr : memref, i32 - omp.atomic.update %x = %x add %expr : memref, i32 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } omp.terminator } return @@ -712,7 +769,11 @@ func @omp_atomic_capture(%x: memref, %v: memref, %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 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } omp.atomic.write %x = %expr : memref, i32 omp.terminator } @@ -736,7 +797,11 @@ func @omp_atomic_capture(%x: memref, %y: memref, %v: memref, %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 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } omp.atomic.read %v = %y : memref omp.terminator } @@ -748,7 +813,11 @@ omp.atomic.capture { // expected-error @below {{captured variable in omp.atomic.read must be updated in second operation}} omp.atomic.read %v = %y : memref - omp.atomic.update %x = %x add %expr : memref, i32 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield (%newval : i32) + } omp.terminator } } diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -541,61 +541,55 @@ // CHECK-LABEL: omp_atomic_update // CHECK-SAME: (%[[X:.*]]: memref, %[[EXPR:.*]]: i32, %[[XBOOL:.*]]: memref, %[[EXPRBOOL:.*]]: i1) func @omp_atomic_update(%x : memref, %expr : i32, %xBool : memref, %exprBool : i1) { - // CHECK: omp.atomic.update %[[X]] = %[[X]] add %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x add %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] sub %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x sub %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] mul %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x mul %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] div %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x div %expr : memref, i32 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] and %[[EXPRBOOL]] : memref, i1 - omp.atomic.update %xBool = %xBool and %exprBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] or %[[EXPRBOOL]] : memref, i1 - omp.atomic.update %xBool = %xBool or %exprBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] xor %[[EXPRBOOL]] : memref, i1 - omp.atomic.update %xBool = %xBool xor %exprBool : memref, i1 - // CHECK: omp.atomic.update %[[X]] = %[[X]] shiftr %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x shiftr %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] shiftl %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x shiftl %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] max %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x max %expr : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[X]] min %[[EXPR]] : memref, i32 - omp.atomic.update %x = %x min %expr : memref, i32 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] eqv %[[EXPRBOOL]] : memref, i1 - omp.atomic.update %xBool = %xBool eqv %exprBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[XBOOL]] neqv %[[EXPRBOOL]] : memref, i1 - omp.atomic.update %xBool = %xBool neqv %exprBool : memref, i1 - - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] add %[[X]] : memref, i32 - omp.atomic.update %x = %expr add %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] sub %[[X]] : memref, i32 - omp.atomic.update %x = %expr sub %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] mul %[[X]] : memref, i32 - omp.atomic.update %x = %expr mul %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] div %[[X]] : memref, i32 - omp.atomic.update %x = %expr div %x : memref, i32 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] and %[[XBOOL]] : memref, i1 - omp.atomic.update %xBool = %exprBool and %xBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] or %[[XBOOL]] : memref, i1 - omp.atomic.update %xBool = %exprBool or %xBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] xor %[[XBOOL]] : memref, i1 - omp.atomic.update %xBool = %exprBool xor %xBool : memref, i1 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] shiftr %[[X]] : memref, i32 - omp.atomic.update %x = %expr shiftr %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] shiftl %[[X]] : memref, i32 - omp.atomic.update %x = %expr shiftl %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] max %[[X]] : memref, i32 - omp.atomic.update %x = %expr max %x : memref, i32 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] min %[[X]] : memref, i32 - omp.atomic.update %x = %expr min %x : memref, i32 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] eqv %[[XBOOL]] : memref, i1 - omp.atomic.update %xBool = %exprBool eqv %xBool : memref, i1 - // CHECK: omp.atomic.update %[[XBOOL]] = %[[EXPRBOOL]] neqv %[[XBOOL]] : memref, i1 - omp.atomic.update %xBool = %exprBool neqv %xBool : memref, i1 - // CHECK: omp.atomic.update %[[X]] = %[[EXPR]] add %[[X]] memory_order(seq_cst) hint(speculative) : memref, i32 - omp.atomic.update %x = %expr add %x hint(speculative) memory_order(seq_cst) : memref, i32 + // CHECK: omp.atomic.update %[[X]] : memref + // CHECK-NEXT: (%[[XVAL:.*]]: i32): + // CHECK-NEXT: %[[NEWVAL:.*]] = llvm.add %[[XVAL]], %[[EXPR]] : i32 + // CHECK-NEXT: omp.yield(%[[NEWVAL]] : i32) + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield(%newval : i32) + } + // CHECK: omp.atomic.update %[[XBOOL]] : memref + // CHECK-NEXT: (%[[XVAL:.*]]: i1): + // CHECK-NEXT: %[[NEWVAL:.*]] = llvm.and %[[XVAL]], %[[EXPRBOOL]] : i1 + // CHECK-NEXT: omp.yield(%[[NEWVAL]] : i1) + omp.atomic.update %xBool : memref { + ^bb0(%xval: i1): + %newval = llvm.and %xval, %exprBool : i1 + omp.yield(%newval : i1) + } + // CHECK: omp.atomic.update %[[X]] : memref + // CHECK-NEXT: (%[[XVAL:.*]]: i32): + // CHECK-NEXT: %[[NEWVAL:.*]] = llvm.shl %[[XVAL]], %[[EXPR]] : i32 + // CHECK-NEXT: omp.yield(%[[NEWVAL]] : i32) + // CHECK-NEXT: } + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.shl %xval, %expr : i32 + omp.yield(%newval : i32) + } + // CHECK: omp.atomic.update %[[X]] : memref + // CHECK-NEXT: (%[[XVAL:.*]]: i32): + // CHECK-NEXT: %[[NEWVAL:.*]] = "llvm.intr.smax"(%[[XVAL]], %[[EXPR]]) : (i32, i32) -> i32 + // CHECK-NEXT: omp.yield(%[[NEWVAL]] : i32) + // CHECK-NEXT: } + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = "llvm.intr.smax"(%xval, %expr) : (i32, i32) -> i32 + omp.yield(%newval : i32) + } + + // CHECK: omp.atomic.update %[[XBOOL]] : memref + // CHECK-NEXT: (%[[XVAL:.*]]: i1): + // CHECK-NEXT: %[[NEWVAL:.*]] = llvm.icmp "eq" %[[XVAL]], %[[EXPRBOOL]] : i1 + // CHECK-NEXT: omp.yield(%[[NEWVAL]] : i1) + // } + omp.atomic.update %xBool : memref { + ^bb0(%xval: i1): + %newval = llvm.icmp "eq" %xval, %exprBool : i1 + omp.yield(%newval : i1) + } return } @@ -603,23 +597,39 @@ // CHECK-SAME: (%[[v:.*]]: memref, %[[x:.*]]: memref, %[[expr:.*]]: i32) func @omp_atomic_capture(%v: memref, %x: memref, %expr: i32) { // CHECK: omp.atomic.capture{ - // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref, i32 + // CHECK-NEXT: omp.atomic.update %[[x]] : memref + // CHECK-NEXT: (%[[xval:.*]]: i32): + // CHECK-NEXT: %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32 + // CHECK-NEXT: omp.yield(%[[newval]] : i32) + // CHECK-NEXT: } // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref // CHECK-NEXT: omp.terminator // CHECK-NEXT: } omp.atomic.capture{ - omp.atomic.update %x = %expr add %x : memref, i32 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield(%newval : i32) + } omp.atomic.read %v = %x : memref omp.terminator } // CHECK: omp.atomic.capture{ // CHECK-NEXT: omp.atomic.read %[[v]] = %[[x]] : memref - // CHECK-NEXT: omp.atomic.update %[[x]] = %[[expr]] add %[[x]] : memref, i32 + // CHECK-NEXT: omp.atomic.update %[[x]] : memref + // CHECK-NEXT: (%[[xval:.*]]: i32): + // CHECK-NEXT: %[[newval:.*]] = llvm.add %[[xval]], %[[expr]] : i32 + // CHECK-NEXT: omp.yield(%[[newval]] : i32) + // CHECK-NEXT: } // CHECK-NEXT: omp.terminator // CHECK-NEXT: } omp.atomic.capture{ omp.atomic.read %v = %x : memref - omp.atomic.update %x = %expr add %x : memref, i32 + omp.atomic.update %x : memref { + ^bb0(%xval: i32): + %newval = llvm.add %xval, %expr : i32 + omp.yield(%newval : i32) + } omp.terminator } // CHECK: omp.atomic.capture{