From d4989a22dfc8dc1ff4cdf515ba38fc5ab04096af Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof@amd.com>
Date: Mon, 23 Sep 2024 14:23:53 +0100
Subject: [PATCH] [MLIR][OpenMP] More robust support for target SPMD

This patch fixes `TargetOp::getInnermostCapturedOmpOp()` to avoid detecting as
captured by the top target construct other constructs nested inside of a loop.
This prevents the `omp.target` verifier from incorrectly flagging valid SPMD
loops, like in the following example:

```f90
subroutine foo(n)
  implicit none
  integer, intent(in) :: n
  integer :: i, j

  !$omp target teams distribute parallel do
  do i=1,n
    !$omp simd
    do j=1,n
      call bar()
    enddo
  enddo
end subroutine foo
```
---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td |  9 ++++++---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  | 19 ++++++++-----------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index aeed1ddf7ea9cf..e31832869a22ab 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1127,13 +1127,16 @@ def TargetOp : OpenMP_Op<"target", traits = [
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{
-    /// Returns the innermost OpenMP dialect operation nested inside of this
-    /// operation's region. For an operation to be detected as captured, it must
-    /// be inside a (possibly multi-level) nest of OpenMP dialect operation's
+    /// Returns the innermost OpenMP dialect operation captured by this target
+    /// construct. For an operation to be detected as captured, it must be
+    /// inside a (possibly multi-level) nest of OpenMP dialect operation's
     /// regions where none of these levels contain other operations considered
     /// not-allowed for these purposes (i.e. only terminator operations are
     /// allowed from the OpenMP dialect, and other dialect's operations are
     /// allowed as long as they don't have a memory write effect).
+    ///
+    /// If there are omp.loop_nest operations in the sequence of nested
+    /// operations, the top level one will be the one captured.
     Operation *getInnermostCapturedOmpOp();
 
     /// Tells whether this target region represents a single worksharing loop
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 5dfe27327e8357..bc08d41da655cb 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1540,6 +1540,13 @@ Operation *TargetOp::getInnermostCapturedOmpOp() {
     if (op == *this)
       return;
 
+    // Reset captured op if crossing through an omp.loop_nest, so that the top
+    // level one will be the one captured.
+    if (llvm::isa<LoopNestOp>(op)) {
+      capturedOp = nullptr;
+      capturedParentRegion = nullptr;
+    }
+
     bool isOmpDialect = op->getDialect() == ompDialect;
     bool hasRegions = op->getNumRegions() > 0;
 
@@ -1583,21 +1590,11 @@ 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;
 
-  Operation *workshareOp = capturedOp->getParentOp();
-
   // Accept optional SIMD leaf construct.
+  Operation *workshareOp = capturedOp->getParentOp();
   if (isa_and_present<SimdOp>(workshareOp))
     workshareOp = workshareOp->getParentOp();