[flang-commits] [flang] [Flang][OpenMP] Support unstructured code in target data (PR #71051)

[Kiran Chandramohan via flang-commits]

https://github.com/kiranchandramohan created https://github.com/llvm/llvm-project/pull/71051

None

>From 2cacf68ef56e69232bfcb8bed0a01dcbd59df39e Mon Sep 17 00:00:00 2001
From: Kiran Chandramohan <kiran.chandramohan at arm.com>
Date: Thu, 2 Nov 2023 11:55:20 +0000
Subject: [PATCH] [Flang][OpenMP] Support unstructured code in target data

---
 flang/lib/Lower/OpenMP.cpp         | 36 ++++++++++++++++++++++++------
 flang/test/Lower/OpenMP/target.f90 | 23 +++++++++++++++++++
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/flang/lib/Lower/OpenMP.cpp b/flang/lib/Lower/OpenMP.cpp
index 1b9a03f74ac479f..0ca28f10bba0300 100644
--- a/flang/lib/Lower/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP.cpp
@@ -2173,7 +2173,8 @@ static void createBodyOfOp(
 }
 
 static void createBodyOfTargetDataOp(
-    Fortran::lower::AbstractConverter &converter, mlir::omp::DataOp &dataOp,
+    Fortran::lower::AbstractConverter &converter,
+    Fortran::lower::pft::Evaluation &eval, mlir::omp::DataOp &dataOp,
     const llvm::SmallVector<mlir::Type> &useDeviceTypes,
     const llvm::SmallVector<mlir::Location> &useDeviceLocs,
     const llvm::SmallVector<const Fortran::semantics::Symbol *>
@@ -2183,8 +2184,6 @@ static void createBodyOfTargetDataOp(
   mlir::Region &region = dataOp.getRegion();
 
   firOpBuilder.createBlock(&region, {}, useDeviceTypes, useDeviceLocs);
-  firOpBuilder.create<mlir::omp::TerminatorOp>(currentLocation);
-  firOpBuilder.setInsertionPointToStart(&region.front());
 
   unsigned argIndex = 0;
   for (const Fortran::semantics::Symbol *sym : useDeviceSymbols) {
@@ -2213,6 +2212,27 @@ static void createBodyOfTargetDataOp(
     }
     argIndex++;
   }
+
+  // Insert dummy instruction to remember the insertion position. The
+  // marker will be deleted since there are not uses.
+  // In the HLFIR flow there are hlfir.declares inserted above while
+  // setting block arguments.
+  mlir::Value undefMarker = firOpBuilder.create<fir::UndefOp>(
+      dataOp.getOperation()->getLoc(), firOpBuilder.getIndexType());
+
+  // Create blocks for unstructured regions. This has to be done since
+  // the original block allocation are created with the function as
+  // the parent region of blocks.
+  if (eval.lowerAsUnstructured()) {
+    Fortran::lower::createEmptyRegionBlocks<mlir::omp::TerminatorOp,
+                                            mlir::omp::YieldOp>(
+        firOpBuilder, eval.getNestedEvaluations());
+  }
+
+  firOpBuilder.create<mlir::omp::TerminatorOp>(currentLocation);
+
+  // Create the insertion point after the marker.
+  firOpBuilder.setInsertionPointAfter(undefMarker.getDefiningOp());
 }
 
 template <typename OpTy, typename... Args>
@@ -2360,6 +2380,7 @@ genTaskGroupOp(Fortran::lower::AbstractConverter &converter,
 
 static mlir::omp::DataOp
 genDataOp(Fortran::lower::AbstractConverter &converter,
+          Fortran::lower::pft::Evaluation &eval,
           Fortran::semantics::SemanticsContext &semanticsContext,
           mlir::Location currentLocation,
           const Fortran::parser::OmpClauseList &clauseList) {
@@ -2386,8 +2407,8 @@ genDataOp(Fortran::lower::AbstractConverter &converter,
   auto dataOp = converter.getFirOpBuilder().create<mlir::omp::DataOp>(
       currentLocation, ifClauseOperand, deviceOperand, devicePtrOperands,
       deviceAddrOperands, mapOperands);
-  createBodyOfTargetDataOp(converter, dataOp, useDeviceTypes, useDeviceLocs,
-                           useDeviceSymbols, currentLocation);
+  createBodyOfTargetDataOp(converter, eval, dataOp, useDeviceTypes,
+                           useDeviceLocs, useDeviceSymbols, currentLocation);
   return dataOp;
 }
 
@@ -2603,7 +2624,7 @@ genOmpSimpleStandalone(Fortran::lower::AbstractConverter &converter,
     firOpBuilder.create<mlir::omp::TaskyieldOp>(currentLocation);
     break;
   case llvm::omp::Directive::OMPD_target_data:
-    genDataOp(converter, semanticsContext, currentLocation, opClauseList);
+    genDataOp(converter, eval, semanticsContext, currentLocation, opClauseList);
     break;
   case llvm::omp::Directive::OMPD_target_enter_data:
     genEnterExitDataOp<mlir::omp::EnterDataOp>(converter, semanticsContext,
@@ -2892,7 +2913,8 @@ genOMP(Fortran::lower::AbstractConverter &converter,
                 beginClauseList, directive.v);
     break;
   case llvm::omp::Directive::OMPD_target_data:
-    genDataOp(converter, semanticsContext, currentLocation, beginClauseList);
+    genDataOp(converter, eval, semanticsContext, currentLocation,
+              beginClauseList);
     break;
   case llvm::omp::Directive::OMPD_task:
     genTaskOp(converter, eval, currentLocation, beginClauseList);
diff --git a/flang/test/Lower/OpenMP/target.f90 b/flang/test/Lower/OpenMP/target.f90
index a94859a3a1aa7d8..754e13d1a36b504 100644
--- a/flang/test/Lower/OpenMP/target.f90
+++ b/flang/test/Lower/OpenMP/target.f90
@@ -265,6 +265,29 @@ subroutine omp_target_device_addr
    !CHECK: }
 end subroutine omp_target_device_addr
 
+
+!===============================================================================
+! Target Data with unstructured code
+!===============================================================================
+!CHECK-LABEL: func.func @_QPsb
+subroutine sb
+   integer :: i = 1
+   integer :: j = 11
+!CHECK: omp.target_data map_entries(%{{.*}}, %{{.*}} : !fir.ref<i32>, !fir.ref<i32>)
+   !$omp target data map(tofrom: i, j)
+     j = j - 1
+!CHECK: %[[J_VAL:.*]] = arith.subi
+!CHECK: hlfir.assign %[[J_VAL]]
+!CHECK: cf.br ^[[BB:.*]]
+!CHECK: ^[[BB]]:
+     goto 20
+20   i = i + 1
+!CHECK: %[[I_VAL:.*]] = arith.addi
+!CHECK: hlfir.assign %[[I_VAL]]
+!CHECK: omp.terminator
+   !$omp end target data
+end subroutine
+
 !===============================================================================
 ! Target with parallel loop
 !===============================================================================