[MLIR][MemRef] Emit error on atomic generic result op defined outside the region

[ashermancinelli]

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

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Asher Mancinelli (ashermancinelli)

Changes

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

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp (+6-1)
• (modified) mlir/test/Conversion/MemRefToLLVM/invalid.mlir (+12)

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
+}