[Mlir-commits] [mlir] [MLIR][OpenACC] Fix host fallback for acc.atomic.update (PR #207597)

[llvmlistbot at llvm.org]

https://github.com/ES3Q created https://github.com/llvm/llvm-project/pull/207597

The host fallback for `acc.atomic.update` only processed the first operation in the region and used its result for the store, ignoring the remaining operations and the `acc.yield` terminator. This generated invalid IR when the region contained multiple operations.

1.Fix this by cloning all operations in the region and using the operand of `acc.yield` as the final result to store.
2.Add tests for atomic read, write, update, and capture operations to cover the host fallback path.

>From b2e0378372c5dd221c0cfe95c9ed62dba8b292e4 Mon Sep 17 00:00:00 2001
From: ES3Q <ES3Q at QQ.COM>
Date: Mon, 6 Jul 2026 00:08:25 +0800
Subject: [PATCH] [MLIR][OpenACC] Fix host fallback for acc.atomic.update

The host fallback for  incorrectly cloned only the
first operation in the region and used its result for the store. This
generated invalid IR when the region contained multiple operations
leading up to the yield.

Fix this by cloning all operations in the region and using the operand
of  as the final result to store.
---
 .../Transforms/ACCSpecializeForHost.cpp       |  8 ++-
 .../acc-specialize-for-host-fallback.mlir     | 63 +++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
index d3c8b683b76cb..5f41a1175f21f 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp
@@ -142,9 +142,13 @@ class ACCOrphanAtomicUpdateOpConversion
       }
       IRMapping mapping;
       mapping.map(atomicUpdateOp.getRegion().front().getArgument(0), loadOp);
-      Operation *expr = rewriter.clone(*atomicUpdateOp.getFirstOp(), mapping);
+      Block &block = atomicUpdateOp.getRegion().front();
+      for (Operation &op : block.without_terminator())
+        rewriter.clone(op, mapping);
+      auto yieldOp = cast<acc::YieldOp>(block.getTerminator());
+      Value result = mapping.lookup(yieldOp.getOperand(0));
       if (!ptrLikeType.genStore(rewriter, atomicUpdateOp.getLoc(),
-                                expr->getResult(0), xTyped)) {
+                                result, xTyped)) {
         accSupport.emitNYI(atomicUpdateOp.getLoc(),
                            "failed to generate store for atomic update");
         return failure();
diff --git a/mlir/test/Dialect/OpenACC/acc-specialize-for-host-fallback.mlir b/mlir/test/Dialect/OpenACC/acc-specialize-for-host-fallback.mlir
index 59269b71bf61c..1a02f9fd710db 100644
--- a/mlir/test/Dialect/OpenACC/acc-specialize-for-host-fallback.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-specialize-for-host-fallback.mlir
@@ -155,3 +155,66 @@ func.func @declare_enter_exit(%arg0 : memref<i32>) attributes {acc.routine_info
   acc.declare_exit token(%token) dataOperands(%0 : memref<i32>)
   return
 }
+
+//===----------------------------------------------------------------------===//
+// Atomic operations (host fallback)
+//===----------------------------------------------------------------------===//
+
+acc.routine @acc_routine_atomic_read func(@atomic_read) seq
+// CHECK-LABEL: func.func @atomic_read
+// CHECK-NOT:   acc.atomic
+// CHECK:       memref.load
+func.func @atomic_read(%src : memref<f64>, %dst : memref<f64>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_atomic_read]>} {
+  acc.atomic.read %dst = %src : memref<f64>, memref<f64>, f64
+  return
+}
+
+acc.routine @acc_routine_atomic_write func(@atomic_write) seq
+// CHECK-LABEL: func.func @atomic_write
+// CHECK-NOT:   acc.atomic
+// CHECK:       memref.store
+func.func @atomic_write(%addr : memref<f64>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_atomic_write]>} {
+  %val = arith.constant 5.0e-01 : f64
+  acc.atomic.write %addr = %val : memref<f64>, f64
+  return
+}
+
+acc.routine @acc_routine_atomic_update func(@atomic_update) seq
+// CHECK-LABEL: func.func @atomic_update
+// CHECK-NOT:   acc.atomic
+// CHECK:       arith.select
+func.func @atomic_update(%x : memref<f64>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_atomic_update]>} {
+  %a = arith.constant 5.0e-01 : f64
+  acc.parallel {
+    acc.atomic.update %x : memref<f64> {
+    ^bb0(%arg0: f64):
+      %c = arith.cmpf ogt, %a, %arg0 fastmath<contract> : f64
+      %r = arith.select %c, %a, %arg0 : f64
+      acc.yield %r : f64
+    }
+    acc.terminator
+  }
+  return
+}
+
+acc.routine @acc_routine_atomic_capture func(@atomic_capture) seq
+// CHECK-LABEL: func.func @atomic_capture
+// CHECK-NOT:   acc.atomic
+// CHECK:       arith.select
+func.func @atomic_capture(%x : memref<f64>, %dst : memref<f64>) attributes {acc.routine_info = #acc.routine_info<[@acc_routine_atomic_capture]>} {
+  %a = arith.constant 5.0e-01 : f64
+  acc.parallel {
+    acc.atomic.capture {
+      acc.atomic.update %x : memref<f64> {
+      ^bb0(%arg0: f64):
+        %c = arith.cmpf ogt, %a, %arg0 fastmath<contract> : f64
+        %r = arith.select %c, %a, %arg0 : f64
+        acc.yield %r : f64
+      }
+      acc.atomic.read %dst = %x : memref<f64>, memref<f64>, f64
+      acc.terminator
+    }
+    acc.terminator
+  }
+  return
+}