[NFC][flang][OpenMP] Extract target region utils to map or clone outside values (#155754)

Following up on #154483, this PR introduces further refactoring to
extract some shared utils between OpenMP lowering and `do concurrent`
conversion pass. In particular, this PR extracts 2 utils that handle
mapping or cloning values used inside target regions but defined
outside.

Later `do concurrent` PR(s) will also use these utils.

PR stack:
- https://github.com/llvm/llvm-project/pull/155754 ◀️
- https://github.com/llvm/llvm-project/pull/155987
- https://github.com/llvm/llvm-project/pull/155992
- https://github.com/llvm/llvm-project/pull/155993
- https://github.com/llvm/llvm-project/pull/156589
- https://github.com/llvm/llvm-project/pull/156610
- https://github.com/llvm/llvm-project/pull/156837

1 files changed, 113 insertions, 0 deletions

diff --git a/flang/lib/Utils/OpenMP.cpp b/flang/lib/Utils/OpenMP.cpp
index e1681e9c3487..2261912fec22 100644
--- a/flang/lib/Utils/OpenMP.cpp
+++ b/flang/lib/Utils/OpenMP.cpp
@@ -8,10 +8,14 @@
 
 #include "flang/Utils/OpenMP.h"
 
+#include "flang/Lower/ConvertExprToHLFIR.h"
+#include "flang/Optimizer/Builder/DirectivesCommon.h"
+#include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 namespace Fortran::utils::openmp {
 mlir::omp::MapInfoOp createMapInfoOp(mlir::OpBuilder &builder,
@@ -44,4 +48,113 @@ mlir::omp::MapInfoOp createMapInfoOp(mlir::OpBuilder &builder,
           builder.getStringAttr(name), builder.getBoolAttr(partialMap));
   return op;
 }
+
+mlir::Value mapTemporaryValue(fir::FirOpBuilder &firOpBuilder,
+    mlir::omp::TargetOp targetOp, mlir::Value val, llvm::StringRef name) {
+  mlir::OpBuilder::InsertionGuard guard(firOpBuilder);
+  mlir::Operation *valOp = val.getDefiningOp();
+
+  if (valOp)
+    firOpBuilder.setInsertionPointAfter(valOp);
+  else
+    // This means val is a block argument
+    firOpBuilder.setInsertionPoint(targetOp);
+
+  auto copyVal = firOpBuilder.createTemporary(val.getLoc(), val.getType());
+  firOpBuilder.createStoreWithConvert(copyVal.getLoc(), val, copyVal);
+
+  fir::factory::AddrAndBoundsInfo info = fir::factory::getDataOperandBaseAddr(
+      firOpBuilder, val, /*isOptional=*/false, val.getLoc());
+  llvm::SmallVector<mlir::Value> bounds =
+      fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
+          mlir::omp::MapBoundsType>(firOpBuilder, info,
+          hlfir::translateToExtendedValue(
+              val.getLoc(), firOpBuilder, hlfir::Entity{val})
+              .first,
+          /*dataExvIsAssumedSize=*/false, val.getLoc());
+
+  firOpBuilder.setInsertionPoint(targetOp);
+
+  llvm::omp::OpenMPOffloadMappingFlags mapFlag =
+      llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_IMPLICIT;
+  mlir::omp::VariableCaptureKind captureKind =
+      mlir::omp::VariableCaptureKind::ByRef;
+
+  mlir::Type eleType = copyVal.getType();
+  if (auto refType = mlir::dyn_cast<fir::ReferenceType>(copyVal.getType())) {
+    eleType = refType.getElementType();
+  }
+
+  if (fir::isa_trivial(eleType) || fir::isa_char(eleType)) {
+    captureKind = mlir::omp::VariableCaptureKind::ByCopy;
+  } else if (!fir::isa_builtin_cptr_type(eleType)) {
+    mapFlag |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_TO;
+  }
+
+  mlir::Value mapOp = createMapInfoOp(firOpBuilder, copyVal.getLoc(), copyVal,
+      /*varPtrPtr=*/mlir::Value{}, name.str(), bounds,
+      /*members=*/llvm::SmallVector<mlir::Value>{},
+      /*membersIndex=*/mlir::ArrayAttr{},
+      static_cast<std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+          mapFlag),
+      captureKind, copyVal.getType());
+
+  auto argIface = llvm::cast<mlir::omp::BlockArgOpenMPOpInterface>(*targetOp);
+  mlir::Region &region = targetOp.getRegion();
+
+  // Get the index of the first non-map argument before modifying mapVars,
+  // then append an element to mapVars and an associated entry block
+  // argument at that index.
+  unsigned insertIndex =
+      argIface.getMapBlockArgsStart() + argIface.numMapBlockArgs();
+  targetOp.getMapVarsMutable().append(mapOp);
+  mlir::Value clonedValArg =
+      region.insertArgument(insertIndex, copyVal.getType(), copyVal.getLoc());
+
+  mlir::Block *entryBlock = &region.getBlocks().front();
+  firOpBuilder.setInsertionPointToStart(entryBlock);
+  auto loadOp =
+      firOpBuilder.create<fir::LoadOp>(clonedValArg.getLoc(), clonedValArg);
+  return loadOp.getResult();
+}
+
+void cloneOrMapRegionOutsiders(
+    fir::FirOpBuilder &firOpBuilder, mlir::omp::TargetOp targetOp) {
+  mlir::Region &region = targetOp.getRegion();
+  mlir::Block *entryBlock = &region.getBlocks().front();
+
+  llvm::SetVector<mlir::Value> valuesDefinedAbove;
+  mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
+  while (!valuesDefinedAbove.empty()) {
+    for (mlir::Value val : valuesDefinedAbove) {
+      mlir::Operation *valOp = val.getDefiningOp();
+
+      // NOTE: We skip BoxDimsOp's as the lesser of two evils is to map the
+      // indices separately, as the alternative is to eventually map the Box,
+      // which comes with a fairly large overhead comparatively. We could be
+      // more robust about this and check using a BackwardsSlice to see if we
+      // run the risk of mapping a box.
+      if (valOp && mlir::isMemoryEffectFree(valOp) &&
+          !mlir::isa<fir::BoxDimsOp>(valOp)) {
+        mlir::Operation *clonedOp = valOp->clone();
+        entryBlock->push_front(clonedOp);
+
+        auto replace = [entryBlock](mlir::OpOperand &use) {
+          return use.getOwner()->getBlock() == entryBlock;
+        };
+
+        valOp->getResults().replaceUsesWithIf(clonedOp->getResults(), replace);
+        valOp->replaceUsesWithIf(clonedOp, replace);
+      } else {
+        mlir::Value mappedTemp = mapTemporaryValue(firOpBuilder, targetOp, val,
+            /*name=*/{});
+        val.replaceUsesWithIf(mappedTemp, [entryBlock](mlir::OpOperand &use) {
+          return use.getOwner()->getBlock() == entryBlock;
+        });
+      }
+    }
+    valuesDefinedAbove.clear();
+    mlir::getUsedValuesDefinedAbove(region, valuesDefinedAbove);
+  }
+}
 } // namespace Fortran::utils::openmp