From 40f16d57d0f6139b29c8b02881b9eaa9391701e9 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy@amd.com>
Date: Fri, 2 Aug 2024 03:06:44 -0500
Subject: [PATCH] [flang][OpenMP] Retrieve correct alloca block for target
 ws-loops

Fix an issue uncovered by Genasis where we were retrieving an alloc
block that lies outside the `target` region to insert an allocation.
This also adds an HLFIR to FIR test to make sure we properly validate
the logic related to OpenMP in that layer.
---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp    |  5 ++
 .../Lower/OpenMP/hlfir-to-fir-conv-omp.mlir   | 70 +++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 flang/test/Lower/OpenMP/hlfir-to-fir-conv-omp.mlir

diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index b0dcef18466430..c45d3b9d7a2f3d 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -264,6 +264,11 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
     if (!parallelOp || !llvm::isa_and_present<mlir::omp::DistributeOp>(
                            parallelOp->getParentOp()))
       return ompOutlineableIface.getAllocaBlock();
+
+    if (auto parentOutlineable =
+            parallelOp
+                ->getParentOfType<mlir::omp::OutlineableOpenMPOpInterface>())
+      return parentOutlineable.getAllocaBlock();
   }
 
   if (auto recipeIface =
diff --git a/flang/test/Lower/OpenMP/hlfir-to-fir-conv-omp.mlir b/flang/test/Lower/OpenMP/hlfir-to-fir-conv-omp.mlir
new file mode 100644
index 00000000000000..62f93efde9c643
--- /dev/null
+++ b/flang/test/Lower/OpenMP/hlfir-to-fir-conv-omp.mlir
@@ -0,0 +1,70 @@
+// Tests HLFIR-to-FIR conversion aspects relevant to OpenMP. For example, that
+// the correct alloca block is chosen for OMP regions.
+
+// RUN: fir-opt --convert-hlfir-to-fir %s -o - | \
+// RUN: FileCheck %s
+
+fir.global internal @_QQro.1xi4.0(dense<42> : tensor<1xi32>) constant : !fir.array<1xi32>
+
+func.func @_QPfoo() {
+  %c1 = arith.constant 1 : index
+  %host_alloc = fir.alloca !fir.array<1xi32> {bindc_name = "arr", uniq_name = "_QFfooEarr"}
+
+  %1 = fir.shape %c1 : (index) -> !fir.shape<1>
+  %host_decl:2 = hlfir.declare %host_alloc(%1) {uniq_name = "_QFfooEarr"} : (!fir.ref<!fir.array<1xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<1xi32>>, !fir.ref<!fir.array<1xi32>>)
+  %map_info = omp.map.info var_ptr(%host_decl#1 : !fir.ref<!fir.array<1xi32>>, !fir.array<1xi32>) map_clauses(implicit, tofrom) capture(ByRef)  -> !fir.ref<!fir.array<1xi32>> {name = "arr"}
+
+  // CHECK: omp.target
+  omp.target map_entries(%map_info -> %arg1 : !fir.ref<!fir.array<1xi32>>)  {
+  ^bb0(%arg1: !fir.ref<!fir.array<1xi32>>):
+
+    // CHECK: %[[TO_BOX_ALLOC:.*]] = fir.alloca !fir.box<!fir.array<1xi32>> {pinned}
+    %c1_2 = arith.constant 1 : index
+    %21 = fir.shape %c1_2 : (index) -> !fir.shape<1>
+
+    // CHECK: %[[TARGET_DECL:.*]] = fir.declare
+    %target_decl:2 = hlfir.declare %arg1(%21) {uniq_name = "_QFfooEarr"} : (!fir.ref<!fir.array<1xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<1xi32>>, !fir.ref<!fir.array<1xi32>>)
+
+    // CHECK: omp.teams
+    omp.teams {
+      %c1_3 = arith.constant 1 : i32
+      %c10 = arith.constant 10 : i32
+
+      // CHECK: omp.distribute
+      omp.distribute {
+        // CHECK: omp.parallel
+        omp.parallel {
+          // CHECK: omp.wsloop
+          omp.wsloop {
+            // CHECK: omp.loop_nest
+            omp.loop_nest (%arg2) : i32 = (%c1_3) to (%c10) inclusive step (%c1_3) {
+              %25 = fir.address_of(@_QQro.1xi4.0) : !fir.ref<!fir.array<1xi32>>
+              %26 = fir.shape %c1_2 : (index) -> !fir.shape<1>
+              %27:2 = hlfir.declare %25(%26) {fortran_attrs = #fir.var_attrs<parameter>, uniq_name = "_QQro.1xi4.0"} : (!fir.ref<!fir.array<1xi32>>, !fir.shape<1>) -> (!fir.ref<!fir.array<1xi32>>, !fir.ref<!fir.array<1xi32>>)
+
+
+              // CHECK: %[[EMBOX:.*]] = fir.embox %[[TARGET_DECL]]
+              // CHECK: fir.store %[[EMBOX]] to %[[TO_BOX_ALLOC]]
+              // CHECK: %[[BOX_ALLOC_CONV:.*]] = fir.convert %[[TO_BOX_ALLOC]] : (!fir.ref<!fir.box<!fir.array<1xi32>>>) -> !fir.ref<!fir.box<none>>
+              // CHECK: fir.call @_FortranAAssign(%[[BOX_ALLOC_CONV]], {{.*}})
+              hlfir.assign %27#0 to %target_decl#0 : !fir.ref<!fir.array<1xi32>>, !fir.ref<!fir.array<1xi32>>
+              // CHECK: omp.yield
+              omp.yield
+            }
+            // CHECK: omp.terminator
+            omp.terminator
+          }
+          // CHECK: omp.terminator
+          omp.terminator
+        }
+         // CHECK: omp.terminator
+        omp.terminator
+      }
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+  return
+}