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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-openacc

@llvm/pr-subscribers-mlir-openacc

Author: ES3Q

<details>
<summary>Changes</summary>

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.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCSpecializeForHost.cpp (+6-2) 
- (modified) mlir/test/Dialect/OpenACC/acc-specialize-for-host-fallback.mlir (+63) 


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

``````````

</details>


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