[MLIR][MemRef] Add verifier check for index count vs memref rank in generic_atomic_rmw by joker-eph · Pull Request #189229 · llvm/llvm-project

joker-eph · 2026-03-29T11:48:58Z

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

llvmbot · 2026-03-29T11:49:29Z

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-memref

Author: Mehdi Amini (joker-eph)

Changes

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 --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@+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@+1 {{unsupported layout}}

ftynse · 2026-04-01T08:53:53Z

@@ -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()))


Didn't double-check, let's ensure the op is already constrained to only accept ranked memrefs.

It is, and it reminded me I can write as:

MemRefType memrefType = getMemref().getType();

Instead of

auto memrefType = cast<MemRefType>(getMemref().getType());

Because getMemref() returns a TypedValue.

github-actions · 2026-04-01T10:33:31Z

🐧 Linux x64 Test Results

7750 tests passed
604 tests skipped
1 test failed

Failed Tests

(click on a test name to see its output)

MLIR

MLIR.Dialect/MemRef/invalid.mlir

Exit Code: 1

Command Output (stdout):
--
# RUN: at line 1
/home/gha/actions-runner/_work/llvm-project/llvm-project/build/bin/mlir-opt -allow-unregistered-dialect -split-input-file /home/gha/actions-runner/_work/llvm-project/llvm-project/mlir/test/Dialect/MemRef/invalid.mlir -verify-diagnostics
# executed command: /home/gha/actions-runner/_work/llvm-project/llvm-project/build/bin/mlir-opt -allow-unregistered-dialect -split-input-file /home/gha/actions-runner/_work/llvm-project/llvm-project/mlir/test/Dialect/MemRef/invalid.mlir -verify-diagnostics
# .---command stdout------------
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# `-----------------------------
# .---command stderr------------
# | /home/gha/actions-runner/_work/llvm-project/llvm-project/mlir/test/Dialect/MemRef/invalid.mlir:1130:8: error: unexpected error: 'memref.generic_atomic_rmw' op invalid number of indices for accessed memref, expected 0 but got 1
# |   %r = memref.generic_atomic_rmw %arg0[%idx] : memref<i32> {
# |        ^
# | /home/gha/actions-runner/_work/llvm-project/llvm-project/mlir/test/Dialect/MemRef/invalid.mlir:1129:6: error: expected error "index count (1) does not match memref rank (0)" was not produced
# |   // expected-error@+1 {{index count (1) does not match memref rank (0)}}
# |      ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# `-----------------------------
# error: command failed with exit status: 1

--

If these failures are unrelated to your changes (for example tests are broken or flaky at HEAD), please open an issue at https://github.com/llvm/llvm-project/issues and add the infrastructure label.

github-actions · 2026-04-01T10:33:31Z

🪟 Windows x64 Test Results

3598 tests passed
415 tests skipped
1 test failed

Failed Tests

(click on a test name to see its output)

MLIR

MLIR.Dialect/MemRef/invalid.mlir

Exit Code: 1

Command Output (stdout):
--
# RUN: at line 1
c:\_work\llvm-project\llvm-project\build\bin\mlir-opt.exe -allow-unregistered-dialect -split-input-file C:\_work\llvm-project\llvm-project\mlir\test\Dialect\MemRef\invalid.mlir -verify-diagnostics
# executed command: 'c:\_work\llvm-project\llvm-project\build\bin\mlir-opt.exe' -allow-unregistered-dialect -split-input-file 'C:\_work\llvm-project\llvm-project\mlir\test\Dialect\MemRef\invalid.mlir' -verify-diagnostics
# .---command stdout------------
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# | // -----
# `-----------------------------
# .---command stderr------------
# | C:\_work\llvm-project\llvm-project\mlir\test\Dialect\MemRef\invalid.mlir:1130:8: error: unexpected error: 'memref.generic_atomic_rmw' op invalid number of indices for accessed memref, expected 0 but got 1
# |   %r = memref.generic_atomic_rmw %arg0[%idx] : memref<i32> {
# |        ^
# | C:\_work\llvm-project\llvm-project\mlir\test\Dialect\MemRef\invalid.mlir:1129:6: error: expected error "index count (1) does not match memref rank (0)" was not produced
# |   // expected-error@+1 {{index count (1) does not match memref rank (0)}}
# |      ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# `-----------------------------
# error: command failed with exit status: 1

--

If these failures are unrelated to your changes (for example tests are broken or flaky at HEAD), please open an issue at https://github.com/llvm/llvm-project/issues and add the infrastructure label.

…eneric_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 llvm#178211 Assisted-by: Claude Code

joker-eph · 2026-04-01T10:51:16Z

Can't reproduce the failure locally, rebased and pushed just to see...

joker-eph · 2026-04-01T13:24:18Z

I didn't fetch before rebase, I reproduce the failure now.

This is because of b813b0b which changed this to use IndexedMemoryAccessOpInterface which already catch this case in the verifier now. So this change is obsolete and I can close the bug as fixed.

joker-eph requested review from hanhanW and kuhar March 29, 2026 11:48

llvmbot added mlir mlir:memref labels Mar 29, 2026

ftynse approved these changes Apr 1, 2026

View reviewed changes

joker-eph force-pushed the fix/issue-178211 branch from ecbf644 to 55c7ddd Compare April 1, 2026 10:24

joker-eph enabled auto-merge (squash) April 1, 2026 10:25

joker-eph force-pushed the fix/issue-178211 branch from 55c7ddd to ddfc72c Compare April 1, 2026 10:51

joker-eph closed this Apr 1, 2026

auto-merge was automatically disabled April 1, 2026 13:24
Pull request was closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MLIR][MemRef] Add verifier check for index count vs memref rank in generic_atomic_rmw#189229

[MLIR][MemRef] Add verifier check for index count vs memref rank in generic_atomic_rmw#189229
joker-eph wants to merge 1 commit into
llvm:mainfrom
joker-eph:fix/issue-178211

joker-eph commented Mar 29, 2026

Uh oh!

llvmbot commented Mar 29, 2026 •

edited

Loading

Uh oh!

ftynse Apr 1, 2026

Uh oh!

joker-eph Apr 1, 2026

Uh oh!

github-actions Bot commented Apr 1, 2026

Uh oh!

github-actions Bot commented Apr 1, 2026

Uh oh!

joker-eph commented Apr 1, 2026

Uh oh!

joker-eph commented Apr 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

joker-eph commented Mar 29, 2026

Uh oh!

llvmbot commented Mar 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftynse Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

joker-eph Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Apr 1, 2026

🐧 Linux x64 Test Results

Failed Tests

MLIR

Uh oh!

github-actions Bot commented Apr 1, 2026

🪟 Windows x64 Test Results

Failed Tests

MLIR

Uh oh!

joker-eph commented Apr 1, 2026

Uh oh!

joker-eph commented Apr 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

llvmbot commented Mar 29, 2026 •

edited

Loading