[MLIR][OpenMP] Allow omp.atomic_update in SPMD loops

ROCm · Jun 4, 2024 · 3e9f940 · 3e9f940
1 parent 9071627
commit 3e9f940
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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;