[MLIR] Verify there are no side-effecting ops in GenericAtomicRMWOp body.

[pifon2a]

Updated descrition in td.

[pifon2a]

Fixed typo.

[ftynse]

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

515

This looks like it also disallows ops with regions inside the atomic region. While I don't see an immediate use-case for it, I also don't think we should just ban them. Maybe walk the op instead and check for side effects?

[pifon2a]

Address the comment.

[pifon2a]

Remove whitespace.

[herhut]

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

520

It would be extra nice if we could get a note where the side-effecting operation is. But that is icing on the cake.

[mehdi_amini]

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

520

Having the emitOpError on the operation inside the walk would probably get this achieved nicely I think?

[pifon2a]

Addressed the comments.

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

def GenericAtomicRMWOp : Std_Op<"generic_atomic_rmw", [
   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.
+    performed. No side-effecting ops are allowed in the body of
+    `GenericAtomicRMWOp`.
 
     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

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

 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");
+
+  bool hasSideEffects =
+      op.body()
+          .walk([&](Operation *nestedOp) {
+            return MemoryEffectOpInterface::hasNoEffect(nestedOp)
+                       ? WalkResult::advance()
+                       : WalkResult::interrupt();
+          })
+          .wasInterrupted();
+  if (hasSideEffects)
+    return op.emitError("body of 'generic_atomic_rmw' should contain "
+                        "only operations with no side effects");
   return success();
 }
 
 static ParseResult parseGenericAtomicRMWOp(OpAsmParser &parser,
                                            OperationState &result) {
   OpAsmParser::OperandType memref;
   Type memrefType;
   SmallVector<OpAsmParser::OperandType, 4> ivs;

mlir/test/IR/invalid-ops.mlir

^bb0(%old_value: f32):
       %c1 = constant 1.0 : f32
       atomic_yield %c1 : f32
     }) : (memref<10xf32>, index) -> i32
   return
 }
 
 // -----
 
+func @generic_atomic_rmw_has_side_effects(%I: memref<10xf32>, %i : index) {
+  // expected-error@+1 {{should contain only operations with no side effects}}
+  %x = generic_atomic_rmw %I[%i] : memref<10xf32> {
+    ^bb0(%old_value : f32):
+      %c1 = constant 1.0 : f32
+      %buf = alloc() : memref<2048xf32>
+      atomic_yield %c1 : f32
+  }
+}
+
+// -----
+
 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