[MLIR] Add ReturnLike trait to memref.atomic_yield

[joker-eph]

Without this, the yield isn't considered as the region terminator and the dataflow framework does not consider it live.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

Without this, the yield isn't considered as the region terminator and the dataflow framework does not consider it live.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+1-2)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+15)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 77e3074661abf..fe3b7a80c3dea 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1078,8 +1078,7 @@ def GenericAtomicRMWOp : MemRef_Op<"generic_atomic_rmw", [
 
 def AtomicYieldOp : MemRef_Op<"atomic_yield", [
       HasParent<"GenericAtomicRMWOp">,
-      Pure,
-      Terminator
+      Pure, Terminator, ReturnLike
     ]> {
   let summary = "yield operation for GenericAtomicRMWOp";
   let description = [{
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 8c2a1cf7546f3..631b507c993e4 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -510,3 +510,18 @@ module {
 // CHECK: %[[yield:.*]] = arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK: linalg.yield %[[yield]] : f32
 // CHECK-NOT: arith.subf
+
+// -----
+
+// CHECK-LABEL: func.func @test_atomic_yield
+func.func @test_atomic_yield(%I: memref<10xf32>, %idx : index) {
+  // CHECK: memref.generic_atomic_rmw
+  %x = memref.generic_atomic_rmw %I[%idx] : memref<10xf32> {
+  ^bb0(%current_value : f32):
+  // CHECK: arith.constant
+    %c1 = arith.constant 1.0 : f32
+  // CHECK: memref.atomic_yield
+    memref.atomic_yield %c1 : f32
+  }
+  func.return
+}

[llvmbot]

@llvm/pr-subscribers-mlir-memref

Author: Mehdi Amini (joker-eph)

Changes

Without this, the yield isn't considered as the region terminator and the dataflow framework does not consider it live.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td (+1-2)
• (modified) mlir/test/Transforms/remove-dead-values.mlir (+15)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 77e3074661abf..fe3b7a80c3dea 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -1078,8 +1078,7 @@ def GenericAtomicRMWOp : MemRef_Op<"generic_atomic_rmw", [
 
 def AtomicYieldOp : MemRef_Op<"atomic_yield", [
       HasParent<"GenericAtomicRMWOp">,
-      Pure,
-      Terminator
+      Pure, Terminator, ReturnLike
     ]> {
   let summary = "yield operation for GenericAtomicRMWOp";
   let description = [{
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 8c2a1cf7546f3..631b507c993e4 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -510,3 +510,18 @@ module {
 // CHECK: %[[yield:.*]] = arith.addf %{{.*}}, %{{.*}} : f32
 // CHECK: linalg.yield %[[yield]] : f32
 // CHECK-NOT: arith.subf
+
+// -----
+
+// CHECK-LABEL: func.func @test_atomic_yield
+func.func @test_atomic_yield(%I: memref<10xf32>, %idx : index) {
+  // CHECK: memref.generic_atomic_rmw
+  %x = memref.generic_atomic_rmw %I[%idx] : memref<10xf32> {
+  ^bb0(%current_value : f32):
+  // CHECK: arith.constant
+    %c1 = arith.constant 1.0 : f32
+  // CHECK: memref.atomic_yield
+    memref.atomic_yield %c1 : f32
+  }
+  func.return
+}

[Copilot]

Pull Request Overview

This PR ensures that memref.atomic_yield is recognized as a region terminator by adding the ReturnLike trait, so the dataflow framework treats yields as live terminators.

• Add a new test case in remove-dead-values.mlir to verify atomic_yield behavior in dead-value elimination.
• Update MemRefOps.td to attach the ReturnLike trait to AtomicYieldOp.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
mlir/test/Transforms/remove-dead-values.mlir	Added a CHECK for memref.atomic_yield in the dead-value test.
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td	Added ReturnLike trait to AtomicYieldOp definition.

Comments suppressed due to low confidence (1)

mlir/test/Transforms/remove-dead-values.mlir:523

• Consider adding a negative check (e.g., // CHECK-NOT: arith.subf) after memref.atomic_yield to ensure no operations remain live past the atomic yield.

  // CHECK: memref.atomic_yield