[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/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");
+
+  // Check if there are any ops with side effects in the body.
+  for (auto &block : op.body()) {
+    for (auto &nestedOp : block.getOperations()) {
+      if (MemoryEffectOpInterface::hasNoEffect(&nestedOp) &&
+          nestedOp.getNumRegions() == 0)
+        continue;
+      return nestedOp.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@+4 {{body of 'generic_atomic_rmw' 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