From 9585fa21b24971f0a0a0b9665c3545b080c5fbb4 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy@amd.com>
Date: Thu, 11 Jul 2024 00:09:51 -0500
Subject: [PATCH] [flang][OpenMP] Further handling of `target ... do` fixup
 logic

Cloning `hlfir.declare` ops during fix up can result in type mismatches.
For example, a `fir.ref<fir.array<...>>` variable is mapped a s
`fir.box<fir.array<...>>` which results in the fix up logic creating
invalid IR. This PR fixes the issue by avoiding cloning `hlfir.declare`
ops in the first place since they are specially handled during the
fix-up process.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             | 17 +++++++++--
 .../OpenMP/target-do-loop-control-exprs.f90   | 29 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index ca96e2618f506b..43982657471289 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -179,10 +179,20 @@ class HostClausesInsertionGuard {
 
     mlir::IRMapping mapper;
     builder.setInsertionPoint(escapingOperand->getOwner());
-    mlir::Operation *lastSliceOp;
+    mlir::Operation *lastSliceOp = nullptr;
+
+    for (auto *op : backwardSlice) {
+      // DeclareOps need special handling by searching for the corresponding ops
+      // in the host. Therefore, do not clone them since this special handling
+      // is done later in the fix-up process.
+      //
+      // TODO this might need a more elaborate handling in the future but for
+      // now this seems sufficient for our purposes.
+      if (llvm::isa<hlfir::DeclareOp>(op))
+        break;
 
-    for (auto *op : backwardSlice)
       lastSliceOp = builder.clone(*op, mapper);
+    }
 
     builder.restoreInsertionPoint(ip);
     return lastSliceOp;
@@ -201,6 +211,9 @@ class HostClausesInsertionGuard {
              "a block argument)");
       mlir::Operation *lastSliceOp = cloneOperandSliceOutsideTargetOp(operand);
 
+      if (lastSliceOp == nullptr)
+        continue;
+
       // Find the index of the operand in the list of results produced by its
       // defining op.
       unsigned operandResultIdx = 0;
diff --git a/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90 b/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
index 027251801ff5df..8f4813bb27cd12 100644
--- a/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
+++ b/flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
@@ -33,3 +33,32 @@ subroutine foo(upper_bound)
 ! CHECK:   omp.target
 
 ! CHECK: }
+
+subroutine foo_with_dummy_arg(nodes)
+  implicit none
+  integer, intent(inout) :: nodes( : )
+  integer :: i
+
+  !$omp target teams distribute parallel do
+    do i = 1, ubound(nodes, 1)
+      nodes(i) = i
+    end do
+  !$omp end target teams distribute parallel do
+end subroutine
+
+! CHECK: func.func @_QPfoo_with_dummy_arg(%[[FUNC_ARG:.*]]: !fir.box<!fir.array<?xi32>> {fir.bindc_name = "nodes"}) {
+
+! CHECK:   %[[ARR_DECL:.*]]:2 = hlfir.declare %[[FUNC_ARG]] dummy_scope
+
+! CHECK:   omp.map.info
+! CHECK:   omp.map.info
+! CHECK:   omp.map.info
+
+! Verify that we get the box dims of the host array declaration not the target
+! one.
+
+! CHECK:   fir.box_dims %[[ARR_DECL]]
+
+! CHECK:   omp.target
+
+! CHECK: }