[Mlir-commits] [mlir] [MLIR][MemRef] Add verifier check for index count vs memref rank in generic_atomic_rmw (PR #189229)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-memref

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`memref.generic_atomic_rmw` did not verify that the number of index operands matches the rank of the memref. This caused a crash during lowering to LLVM when an index count mismatch was present (e.g., a rank-0 memref accessed with one index).

Fixes #<!-- -->178211

Assisted-by: Claude Code

---
Full diff: https://github.com/llvm/llvm-project/pull/189229.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp (+6) 
- (modified) mlir/test/Dialect/MemRef/invalid.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 404b2aacf1450..8728e66919baf 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -1593,6 +1593,12 @@ LogicalResult GenericAtomicRMWOp::verify() {
   if (getResult().getType() != body.getArgument(0).getType())
     return emitOpError("expected block argument of the same type result type");
 
+  auto memrefType = cast<MemRefType>(getMemref().getType());
+  if (getIndices().size() != static_cast<size_t>(memrefType.getRank()))
+    return emitOpError("index count (")
+           << getIndices().size() << ") does not match memref rank ("
+           << memrefType.getRank() << ")";
+
   bool hasSideEffects =
       body.walk([&](Operation *nestedOp) {
             if (isMemoryEffectFree(nestedOp))
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index af068d8ca8e95..eeb1a80278f1a 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/mlir/test/Dialect/MemRef/invalid.mlir
@@ -1125,6 +1125,17 @@ func.func @atomic_yield_type_mismatch(%I: memref<10xf32>, %i : index) {
 
 // -----
 
+func.func @generic_atomic_rmw_rank_mismatch(%arg0: memref<i32>, %idx: index) {
+  // expected-error at +1 {{index count (1) does not match memref rank (0)}}
+  %r = memref.generic_atomic_rmw %arg0[%idx] : memref<i32> {
+  ^bb0(%v: i32):
+    memref.atomic_yield %v : i32
+  }
+  return
+}
+
+// -----
+
 #map0 = affine_map<(d0) -> (d0 floordiv 8, d0 mod 8)>
 func.func @memref_realloc_layout(%src : memref<256xf32, #map0>) -> memref<?xf32>{
   // expected-error at +1 {{unsupported layout}}

``````````

</details>


https://github.com/llvm/llvm-project/pull/189229