[Mlir-commits] [mlir] [MLIR][MemRef] Emit error on atomic generic result op defined outside the region (PR #172190)

[Asher Mancinelli]

https://github.com/ashermancinelli created https://github.com/llvm/llvm-project/pull/172190

While figuring out how to perform an atomic exchange on a memref, I tried the generic atomic rmw with the yielded value captured from the enclosing scope (instead of a plain atomic_rmw with `arith::AtomicRMWKind::assign`). Instead of segfaulting, this PR changes the pass to produce an error when the result is not found in the region's IR map.

It might be more useful to give a suggestion to the user, but giving an error message instead of a crash is at least an imrovement, I think.

See: #172184

>From 022674a9556b0c5ebf82bc3758740e7e12e83591 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Sat, 13 Dec 2025 16:36:00 -0800
Subject: [PATCH] [MLIR][MemRef] Add error on invalid atomic generic

While figuring out how to perform an atomic exchange on a memref,
I tried the generic atomic rmw with the yielded value captured
from the enclosing scope. Instead of segfaulting, this produces
an error.

It might be more useful to give a suggestion to the user, but
giving an error message instead of a crash is at least an imrovement,
I think.

See: #172184
---
 mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp |  7 ++++++-
 mlir/test/Conversion/MemRefToLLVM/invalid.mlir    | 12 ++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index a6f816aa07377..91a0c4b55fa84 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -741,7 +741,12 @@ struct GenericAtomicRMWOpLowering
       Operation *clone = rewriter.clone(nestedOp, mapping);
       mapping.map(nestedOp.getResults(), clone->getResults());
     }
-    Value result = mapping.lookup(entryBlock.getTerminator()->getOperand(0));
+
+    Value result =
+        mapping.lookupOrNull(entryBlock.getTerminator()->getOperand(0));
+    if (!result) {
+      return atomicOp.emitError("result not defined in region");
+    }
 
     // Prepare the epilog of the loop block.
     // Append the cmpxchg op to the end of the loop block.
diff --git a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir b/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
index 0d04bba96bcdb..5462d3278d9e6 100644
--- a/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/invalid.mlir
@@ -40,3 +40,15 @@ func.func @issue_70160() {
   memref.store %0, %alloc1[] : memref<i32>
   func.return
 }
+
+
+// -----
+
+func.func @test_atomic_exch(%arg0: memref<?xi32>, %idx: index, %value: i32) {
+  // expected-error @+1 {{result not defined in region}}
+  %1 = memref.generic_atomic_rmw %arg0[%idx] : memref<?xi32> {
+  ^bb0(%arg3: i32):
+    memref.atomic_yield %value : i32
+  }
+  func.return
+}