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

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/189229

>From ddfc72ce69018f25ea47cb0f8aec5b5b72274e78 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Mar 2026 08:58:03 -0700
Subject: [PATCH] [MLIR][MemRef] Add verifier check for index count vs memref
 rank in generic_atomic_rmw

`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).

Add a verifier check to `GenericAtomicRMWOp::verify()` that emits a
proper error diagnostic instead of crashing in the lowering pattern.

Fixes #178211

Assisted-by: Claude Code
---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp |  6 ++++++
 mlir/test/Dialect/MemRef/invalid.mlir    | 11 +++++++++++
 2 files changed, 17 insertions(+)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index f31811ad7b98e..7b1d3b0bdac7c 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");
 
+  MemRefType 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}}