[MLIR] Add GenericAtomicRMWOp.

[pifon2a]

Summary

https://llvm.discourse.group/t/rfc-add-std-atomic-rmw-op/489

[pifon2a]

Added block arg for the old value.

[pifon2a]

Added `getCurrentValue() method and updated the doc.

[pifon2a]

Remove whitespace in GenericAtomicRMW printer.

[herhut]

Thanks. This looks great to me. Has there been consensus on https://llvm.discourse.group/t/rfc-add-std-atomic-rmw-op/489 yet? There were no more replies but it seems nobody is opposed either.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

507

x -> %x

mlir/test/IR/invalid-ops.mlir

1171

Could you add a test for the case where the memref and result types disagree?

[pifon2a]

Addressed the comments.

[pifon2a]

In D78352#1991742, @herhut wrote:

Thanks. This looks great to me. Has there been consensus on https://llvm.discourse.group/t/rfc-add-std-atomic-rmw-op/489 yet? There were no more replies but it seems nobody is opposed either.

I implemented exactly what we agreed to with flaub@. There were no additional replies after that.

[herhut]

Thanks for adding this!

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

502

We should mention that no side effecting operation is allowed in the region.

536

This is the kind of method I expect provided by a Traits otherwise we're losing consistency across all operations with regions.

[flaub]

LGTM

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

511

Could we add a check that all the ops in the region are side-effect free?

[herhut]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

536

Do you have a new trait in mind or should this be supported by the SingleBlockImplicitTerminator trait? The body builder makes only sense if there is a single region.

These are somewhat widespread already. I'd rather have it as a separate cleanup.

[pifon2a]

Thanks !

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

502

https://reviews.llvm.org/D78559

536

I am happy to do the cleanup as soon as there is some consensus.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

511

done! https://reviews.llvm.org/D78559

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

536

Adding it to SingleBlockImplicitTerminator looks fine to me

[pifon2a]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

536

Will do.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def AtomicRMWOp : Std_Op<"atomic_rmw", [
 
   let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return memref().getType().cast<MemRefType>();
     }
   }];
 }
 
+def GenericAtomicRMWOp : Std_Op<"generic_atomic_rmw", [
+      SingleBlockImplicitTerminator<"AtomicYieldOp">,
+      TypesMatchWith<"result type matches element type of memref",
+                     "memref", "result",
+                     "$_self.cast<MemRefType>().getElementType()">
+    ]> {
+  let summary = "atomic read-modify-write operation with a region";
+  let description = [{
+    The `atomic_rmw` operation provides a way to perform a read-modify-write
+    sequence that is free from data races. The memref operand represents the
+    buffer that the read and write will be performed against, as accessed by
+    the specified indices. The arity of the indices is the rank of the memref.
+    The result represents the latest value that was stored. The region contains
+    the code for the modification itself. The entry block has a single argument
+    that represents the value stored in `memref[indices]` before the write is
+    performed.
+
+    Example:
+
+    ```mlir
+    %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
+      ^bb0(%current_value : f32):
+        %c1 = constant 1.0 : f32
+        %inc = addf %c1, %current_value : f32
+        atomic_yield %inc : f32
+    }
+    ```
+  }];
+
+  let arguments = (ins
+      MemRefOf<[AnySignlessInteger, AnyFloat]>:$memref,
+      Variadic<Index>:$indices);
+
+  let results = (outs
+      AnyTypeOf<[AnySignlessInteger, AnyFloat]>:$result);
+
+  let regions = (region AnyRegion:$body);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<"Builder *builder, OperationState &result, "
+              "Value memref, ValueRange ivs">
+  ];
+
+  let extraClassDeclaration = [{
+    OpBuilder getBodyBuilder() {
+      assert(!body().empty() && "Unexpected empty 'body' region.");
+      Block &block = body().front();
+      return OpBuilder(&block, block.end());
+    }
+    // The value stored in memref[ivs].
+    Value getCurrentValue() {
+      return body().front().getArgument(0);
+    }
+  }];
+}
+
+def AtomicYieldOp : Std_Op<"atomic_yield", [
+      HasParent<"GenericAtomicRMWOp">,
+      NoSideEffect,
+      Terminator
+    ]> {
+  let summary = "yield operation for GenericAtomicRMWOp";
+  let description = [{
+    "atomic_yield" yields an SSA value from a GenericAtomicRMWOp region.
+  }];
+
+  let arguments = (ins AnyType:$result);
+  let assemblyFormat = "$result attr-dict `:` type($result)";
+}
+
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
 def BranchOp : Std_Op<"br",
     [DeclareOpInterfaceMethods<BranchOpInterface>, NoSideEffect, Terminator]> {
   let summary = "branch operation";
   let description = [{

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

case AtomicRMWKind::muli:
     break;
   default:
     break;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// GenericAtomicRMWOp
+//===----------------------------------------------------------------------===//
+
+void GenericAtomicRMWOp::build(Builder *builder, OperationState &result,
+                               Value memref, ValueRange ivs) {
+  result.addOperands(memref);
+  result.addOperands(ivs);
+
+  if (auto memrefType = memref.getType().dyn_cast<MemRefType>()) {
+    Type elementType = memrefType.getElementType();
+    result.addTypes(elementType);
+
+    Region *bodyRegion = result.addRegion();
+    bodyRegion->push_back(new Block());
+    bodyRegion->front().addArgument(elementType);
+  }
+}
+
+static LogicalResult verify(GenericAtomicRMWOp op) {
+  auto &block = op.body().front();
+  if (block.getNumArguments() != 1)
+    return op.emitOpError("expected single number of entry block arguments");
+
+  if (op.getResult().getType() != block.getArgument(0).getType())
+    return op.emitOpError(
+        "expected block argument of the same type result type");
+  return success();
+}
+
+static ParseResult parseGenericAtomicRMWOp(OpAsmParser &parser,
+                                           OperationState &result) {
+  OpAsmParser::OperandType memref;
+  Type memrefType;
+  SmallVector<OpAsmParser::OperandType, 4> ivs;
+
+  Type indexType = parser.getBuilder().getIndexType();
+  if (parser.parseOperand(memref) ||
+      parser.parseOperandList(ivs, OpAsmParser::Delimiter::Square) ||
+      parser.parseColonType(memrefType) ||
+      parser.resolveOperand(memref, memrefType, result.operands) ||
+      parser.resolveOperands(ivs, indexType, result.operands))
+    return failure();
+
+  Region *body = result.addRegion();
+  if (parser.parseRegion(*body, llvm::None, llvm::None))
+    return failure();
+  result.types.push_back(memrefType.cast<MemRefType>().getElementType());
+  return success();
+}
+
+static void print(OpAsmPrinter &p, GenericAtomicRMWOp op) {
+  p << op.getOperationName() << ' ' << op.memref() << "[" << op.indices()
+    << "] : " << op.memref().getType();
+  p.printRegion(op.body());
+  p.printOptionalAttrDict(op.getAttrs());
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicYieldOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(AtomicYieldOp op) {
+  Type parentType = op.getParentOp()->getResultTypes().front();
+  Type resultType = op.result().getType();
+  if (parentType != resultType)
+    return op.emitOpError() << "types mismatch between yield op: " << resultType
+                            << " and its parent: " << parentType;
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Simplify a branch to a block that has a single predecessor. This effectively
 /// merges the two blocks.
 struct SimplifyBrToBlockWithSinglePred : public OpRewritePattern<BranchOp> {
   using OpRewritePattern<BranchOp>::OpRewritePattern;

mlir/test/IR/core-ops.mlir

func @tensor_load_store(%0 : memref<4x4xi32>) {
   // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<4x4xi32>
   %1 = tensor_load %0 : memref<4x4xi32>
   // CHECK: tensor_store %[[TENSOR]], %[[MEMREF]] : memref<4x4xi32>
   tensor_store %1, %0 : memref<4x4xi32>
   return
 }
 
 // CHECK-LABEL: func @atomic_rmw
+// CHECK-SAME: ([[BUF:%.*]]: memref<10xf32>, [[VAL:%.*]]: f32, [[I:%.*]]: index)
 func @atomic_rmw(%I: memref<10xf32>, %val: f32, %i : index) {
-  // CHECK: %{{.*}} = atomic_rmw "addf" %{{.*}}, %{{.*}}[%{{.*}}]
   %x = atomic_rmw "addf" %val, %I[%i] : (f32, memref<10xf32>) -> f32
+  // CHECK: atomic_rmw "addf" [[VAL]], [[BUF]]{{\[}}[[I]]]
+  return
+}
+
+// CHECK-LABEL: func @generic_atomic_rmw
+// CHECK-SAME: ([[BUF:%.*]]: memref<1x2xf32>, [[I:%.*]]: index, [[J:%.*]]: index)
+func @generic_atomic_rmw(%I: memref<1x2xf32>, %i : index, %j : index) {
+  %x = generic_atomic_rmw %I[%i, %j] : memref<1x2xf32> {
+  // CHECK-NEXT: generic_atomic_rmw [[BUF]]{{\[}}[[I]], [[J]]] : memref
+    ^bb0(%old_value : f32):
+      %c1 = constant 1.0 : f32
+      %out = addf %c1, %old_value : f32
+      atomic_yield %out : f32
+  }
   return
 }
 
 // CHECK-LABEL: func @assume_alignment
 // CHECK-SAME: %[[MEMREF:.*]]: memref<4x4xf16>
 func @assume_alignment(%0: memref<4x4xf16>) {
   // CHECK: assume_alignment %[[MEMREF]], 16 : memref<4x4xf16>
   assume_alignment %0, 16 : memref<4x4xf16>
   return
 }

mlir/test/IR/invalid-ops.mlir

 func @atomic_rmw_expects_int(%I: memref<16x10xf32>, %i : index, %val : f32) {
   // expected-error@+1 {{expects an integer type}}
   %x = atomic_rmw "addi" %val, %I[%i, %i] : (f32, memref<16x10xf32>) -> f32
   return
 }
 
 // -----
 
+func @generic_atomic_rmw_wrong_arg_num(%I: memref<10xf32>, %i : index) {
+  // expected-error@+1 {{expected single number of entry block arguments}}
+  %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
+    ^bb0(%arg0 : f32, %arg1 : f32):
+      %c1 = constant 1.0 : f32
+      atomic_yield %c1 : f32
+  }
+  return
+}
+
+// -----
+
+func @generic_atomic_rmw_wrong_arg_type(%I: memref<10xf32>, %i : index) {
+  // expected-error@+1 {{expected block argument of the same type result type}}
+  %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
+    ^bb0(%old_value : i32):
+      %c1 = constant 1.0 : f32
+      atomic_yield %c1 : f32
+  }
+  return
+}
+
+// -----
+
+func @generic_atomic_rmw_result_type_mismatch(%I: memref<10xf32>, %i : index) {
+ // expected-error@+1 {{failed to verify that result type matches element type of memref}}
+ %0 = "std.generic_atomic_rmw"(%I, %i) ( {
+    ^bb0(%old_value: f32):
+      %c1 = constant 1.0 : f32
+      atomic_yield %c1 : f32
+    }) : (memref<10xf32>, index) -> i32
+  return
+}
+
+// -----
+
+func @atomic_yield_type_mismatch(%I: memref<10xf32>, %i : index) {
+  // expected-error@+4 {{op types mismatch between yield op: 'i32' and its parent: 'f32'}}
+  %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
+    ^bb0(%old_value : f32):
+      %c1 = constant 1 : i32
+      atomic_yield %c1 : i32
+  }
+  return
+}
+
+// -----
+
 // alignment is not power of 2.
 func @assume_alignment(%0: memref<4x4xf16>) {
   // expected-error@+1 {{alignment must be power of 2}}
   std.assume_alignment %0, 12 : memref<4x4xf16>
   return
 }
 
 // -----