From 3e9f94076c95b822e39fe92a2688b9a7084ec801 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof@amd.com>
Date: Tue, 4 Jun 2024 16:37:13 +0100
Subject: [PATCH] [MLIR][OpenMP] Allow omp.atomic_update in SPMD loops

---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index f102730013050..29eb6119c6f03 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1543,6 +1543,15 @@ Operation *TargetOp::getInnermostCapturedOmpOp() {
 
 bool TargetOp::isTargetSPMDLoop() {
   Operation *capturedOp = getInnermostCapturedOmpOp();
+
+  // Allow an omp.atomic_update to be captured inside of the loop and still
+  // consider the parent omp.target operation to be potentially defining an SPMD
+  // loop.
+  // TODO: Potentially accept other captured OpenMP dialect operations as well,
+  // if they are allowed inside of an SPMD loop.
+  if (isa_and_present<AtomicUpdateOp>(capturedOp))
+    capturedOp = capturedOp->getParentOp();
+
   if (!isa_and_present<LoopNestOp>(capturedOp))
     return false;