[Flang][mlir] - In fortran lowering, handle block arguments defined outside omp target

bhandarkar-pranav · bhandarkar-pranav · commit a482741cfc92 · 2025-05-02T17:23:59.000-05:00
This patch handles block arguments that are live-in into the target region of
an omp.target op. This is done by simply reusing the mapping mechanism in place
for values that are not block arguments. Further, in MapInfoFinalizationPass, this
patch adds bounds to maps that map `!fir.ref&lt;!fir.char&lt;k, ?&gt;&gt;` types.
Also, we don't clone bounds when binding entry block arguments any more.
diff --git a/flang/include/flang/Optimizer/Builder/DirectivesCommon.h b/flang/include/flang/Optimizer/Builder/DirectivesCommon.h
@@ -17,6 +17,8 @@
 #ifndef FORTRAN_OPTIMIZER_BUILDER_DIRECTIVESCOMMON_H_
 #define FORTRAN_OPTIMIZER_BUILDER_DIRECTIVESCOMMON_H_
 
+#include "BoxValue.h"
+#include "FIRBuilder.h"
 #include "flang/Optimizer/Builder/BoxValue.h"
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/Todo.h"
@@ -131,6 +133,31 @@ gatherBoundsOrBoundValues(fir::FirOpBuilder &builder, mlir::Location loc,
   }
   return values;
 }
+template <typename BoundsOp, typename BoundsType>
+mlir::Value
+genBoundsOpFromBoxChar(fir::FirOpBuilder &builder, mlir::Location loc,
+                       fir::ExtendedValue dataExv, AddrAndBoundsInfo &info) {
+  // TODO: Handle info.isPresent.
+  if (auto boxCharType =
+          mlir::dyn_cast<fir::BoxCharType>(info.addr.getType())) {
+    mlir::Type idxTy = builder.getIndexType();
+    mlir::Type lenType = builder.getCharacterLengthType();
+    mlir::Type refType = builder.getRefType(boxCharType.getEleTy());
+    auto unboxed =
+        builder.create<fir::UnboxCharOp>(loc, refType, lenType, info.addr);
+    mlir::Value zero = builder.createIntegerConstant(loc, idxTy, 0);
+    mlir::Value one = builder.createIntegerConstant(loc, idxTy, 1);
+    mlir::Value extent = unboxed.getResult(1);
+    mlir::Value stride = one;
+    mlir::Value ub = builder.create<mlir::arith::SubIOp>(loc, extent, one);
+    mlir::Type boundTy = builder.getType<mlir::omp::MapBoundsType>();
+    return builder.create<mlir::omp::MapBoundsOp>(
+        loc, boundTy, /*lower_bound=*/zero,
+        /*upper_bound=*/ub, /*extent=*/extent, /*stride=*/stride,
+        /*stride_in_bytes = */ true, /*start_idx=*/zero);
+  }
+  return mlir::Value{};
+}
 
 /// Generate the bounds operation from the descriptor information.
 template <typename BoundsOp, typename BoundsType>
@@ -258,6 +285,10 @@ genImplicitBoundsOps(fir::FirOpBuilder &builder, AddrAndBoundsInfo &info,
     bounds = genBaseBoundsOps<BoundsOp, BoundsType>(builder, loc, dataExv,
                                                     dataExvIsAssumedSize);
   }
+  if (characterWithDynamicLen(fir::unwrapRefType(baseOp.getType()))) {
+    bounds = {genBoundsOpFromBoxChar<BoundsOp, BoundsType>(builder, loc,
+                                                           dataExv, info)};
+  }
 
   return bounds;
 }
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -217,111 +217,36 @@ static void bindEntryBlockArgs(lower::AbstractConverter &converter,
   assert(args.isValid() && "invalid args");
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
 
-  auto bindSingleMapLike = [&converter,
-                            &firOpBuilder](const semantics::Symbol &sym,
-                                           const mlir::Value val,
-                                           const mlir::BlockArgument &arg) {
-    // Clones the `bounds` placing them inside the entry block and returns
-    // them.
-    auto cloneBound = [&](mlir::Value bound) {
-      if (mlir::isMemoryEffectFree(bound.getDefiningOp())) {
-        mlir::Operation *definingOp = bound.getDefiningOp();
-        mlir::Operation *clonedOp = firOpBuilder.clone(*definingOp);
-        // Todo: Do we need to check for more operation types?
-        // For now, specializing only for fir::UnboxCharOp
-        if (auto unboxCharOp = mlir::dyn_cast<fir::UnboxCharOp>(definingOp))
-          return clonedOp->getResult(1);
-        return clonedOp->getResult(0);
-      }
-      TODO(converter.getCurrentLocation(),
-           "target map-like clause operand unsupported bound type");
-    };
-
-    auto cloneBounds = [cloneBound](llvm::ArrayRef<mlir::Value> bounds) {
-      llvm::SmallVector<mlir::Value> clonedBounds;
-      llvm::transform(bounds, std::back_inserter(clonedBounds),
-                      [&](mlir::Value bound) { return cloneBound(bound); });
-      return clonedBounds;
-    };
-
+  auto bindSingleMapLike = [&converter](const semantics::Symbol &sym,
+                                        const mlir::BlockArgument &arg) {
     fir::ExtendedValue extVal = converter.getSymbolExtendedValue(sym);
     auto refType = mlir::dyn_cast<fir::ReferenceType>(arg.getType());
     if (refType && fir::isa_builtin_cptr_type(refType.getElementType())) {
       converter.bindSymbol(sym, arg);
     } else {
       extVal.match(
           [&](const fir::BoxValue &v) {
-            converter.bindSymbol(sym,
-                                 fir::BoxValue(arg, cloneBounds(v.getLBounds()),
-                                               v.getExplicitParameters(),
-                                               v.getExplicitExtents()));
+            converter.bindSymbol(sym, fir::BoxValue(arg, v.getLBounds(),
+                                                    v.getExplicitParameters(),
+                                                    v.getExplicitExtents()));
           },
           [&](const fir::MutableBoxValue &v) {
             converter.bindSymbol(
-                sym, fir::MutableBoxValue(arg, cloneBounds(v.getLBounds()),
+                sym, fir::MutableBoxValue(arg, v.getLBounds(),
                                           v.getMutableProperties()));
           },
           [&](const fir::ArrayBoxValue &v) {
-            converter.bindSymbol(
-                sym, fir::ArrayBoxValue(arg, cloneBounds(v.getExtents()),
-                                        cloneBounds(v.getLBounds()),
-                                        v.getSourceBox()));
+            converter.bindSymbol(sym, fir::ArrayBoxValue(arg, v.getExtents(),
+                                                         v.getLBounds(),
+                                                         v.getSourceBox()));
           },
           [&](const fir::CharArrayBoxValue &v) {
-            converter.bindSymbol(
-                sym, fir::CharArrayBoxValue(arg, cloneBound(v.getLen()),
-                                            cloneBounds(v.getExtents()),
-                                            cloneBounds(v.getLBounds())));
+            converter.bindSymbol(sym, fir::CharArrayBoxValue(arg, v.getLen(),
+                                                             v.getExtents(),
+                                                             v.getLBounds()));
           },
           [&](const fir::CharBoxValue &v) {
-            // In some cases, v.len could reference the input to the
-            // hlfir.declare which is the corresponding v.addr. While this isn't
-            // a big problem by itself, it is desirable to extract this out of
-            // v.addr itself since it's first result will be of type
-            // fir.boxchar<>. For example, consider the following
-            //
-            // func.func private @_QFPrealtest(%arg0: !fir.boxchar<1>)
-            //  %2 = fir.dummy_scope : !fir.dscope
-            //  %3:2 = fir.unboxchar %arg0 : (!fir.boxchar<1>) ->
-            //         (!fir.ref<!fir.char<1,?>>, index)
-            //  %4:2 = hlfir.declare (%3#0, %3#1, %2):(!fir.ref<!fir.char<1,?>>,
-            //                        index,!fir.dscope) ->
-            //                       (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
-
-            // In the case above,
-            // v.addr is
-            //  %4:2 = hlfir.declare (%3#0, %3#1, %2):(!fir.ref<!fir.char<1,?>>,
-            //                        index,!fir.dscope) ->
-            //                       (!fir.boxchar<1>, !fir.ref<!fir.char<1,?>>)
-            // v.len is
-            //  %3:2 = fir.unboxchar %arg0 : (!fir.boxchar<1>) ->
-            //         (!fir.ref<!fir.char<1,?>>, index)
-
-            // Mapping this to the target will create a use of %arg0 on the
-            // target. Since omp.target is IsolatedFromAbove, %arg0 will have to
-            // be mapped. Presently, OpenMP lowering of target barfs when it has
-            // to map a value that doesnt have a defining op. This can be fixed.
-            // Or we ensure that v.len is fir.unboxchar %4#0 which will
-            // cause %4#1 to be used on the target and consequently be
-            // mapped to the target. As such then, there wont be any use of the
-            // block argument %arg0 on the target.
-
-            mlir::Value len = v.getLen();
-            if (auto declareOp = val.getDefiningOp<hlfir::DeclareOp>()) {
-              mlir::Value base = declareOp.getBase();
-              if (auto boxCharType =
-                      mlir::dyn_cast<fir::BoxCharType>(base.getType())) {
-                mlir::Type lenType = firOpBuilder.getCharacterLengthType();
-                mlir::Type refType =
-                    firOpBuilder.getRefType(boxCharType.getEleTy());
-                mlir::Location loc = converter.getCurrentLocation();
-                auto unboxed = firOpBuilder.create<fir::UnboxCharOp>(
-                    loc, refType, lenType, base);
-                len = unboxed.getResult(1);
-              }
-            }
-            auto charBoxValue = fir::CharBoxValue(arg, cloneBound(len));
-            converter.bindSymbol(sym, charBoxValue);
+            converter.bindSymbol(sym, fir::CharBoxValue(arg, v.getLen()));
           },
           [&](const fir::UnboxedValue &v) { converter.bindSymbol(sym, arg); },
           [&](const auto &) {
@@ -333,7 +258,6 @@ static void bindEntryBlockArgs(lower::AbstractConverter &converter,
 
   auto bindMapLike =
       [&bindSingleMapLike](llvm::ArrayRef<const semantics::Symbol *> syms,
-                           llvm::ArrayRef<mlir::Value> vars,
                            llvm::ArrayRef<mlir::BlockArgument> args) {
         // Structure component symbols don't have bindings, and can only be
         // explicitly mapped individually. If a member is captured implicitly
@@ -342,8 +266,8 @@ static void bindEntryBlockArgs(lower::AbstractConverter &converter,
         llvm::copy_if(syms, std::back_inserter(processedSyms),
                       [](auto *sym) { return !sym->owner().IsDerivedType(); });
 
-        for (auto [sym, var, arg] : llvm::zip_equal(processedSyms, vars, args))
-          bindSingleMapLike(*sym, var, arg);
+        for (auto [sym, arg] : llvm::zip_equal(processedSyms, args))
+          bindSingleMapLike(*sym, arg);
       };
 
   auto bindPrivateLike = [&converter, &firOpBuilder](
@@ -374,20 +298,17 @@ static void bindEntryBlockArgs(lower::AbstractConverter &converter,
   // Process in clause name alphabetical order to match block arguments order.
   // Do not bind host_eval variables because they cannot be used inside of the
   // corresponding region, except for very specific cases handled separately.
-  bindMapLike(args.hasDeviceAddr.syms, args.hasDeviceAddr.vars,
-              op.getHasDeviceAddrBlockArgs());
+  bindMapLike(args.hasDeviceAddr.syms, op.getHasDeviceAddrBlockArgs());
   bindPrivateLike(args.inReduction.syms, args.inReduction.vars,
                   op.getInReductionBlockArgs());
-  bindMapLike(args.map.syms, args.map.vars, op.getMapBlockArgs());
+  bindMapLike(args.map.syms, op.getMapBlockArgs());
   bindPrivateLike(args.priv.syms, args.priv.vars, op.getPrivateBlockArgs());
   bindPrivateLike(args.reduction.syms, args.reduction.vars,
                   op.getReductionBlockArgs());
   bindPrivateLike(args.taskReduction.syms, args.taskReduction.vars,
                   op.getTaskReductionBlockArgs());
-  bindMapLike(args.useDeviceAddr.syms, args.useDeviceAddr.vars,
-              op.getUseDeviceAddrBlockArgs());
-  bindMapLike(args.useDevicePtr.syms, args.useDevicePtr.vars,
-              op.getUseDevicePtrBlockArgs());
+  bindMapLike(args.useDeviceAddr.syms, op.getUseDeviceAddrBlockArgs());
+  bindMapLike(args.useDevicePtr.syms, op.getUseDevicePtrBlockArgs());
 }
 
 /// Get the list of base values that the specified map-like variables point to.
@@ -1427,14 +1348,13 @@ static void genBodyOfTargetOp(
   while (!valuesDefinedAbove.empty()) {
     for (mlir::Value val : valuesDefinedAbove) {
       mlir::Operation *valOp = val.getDefiningOp();
-      assert(valOp != nullptr);
 
       // 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 (mlir::isMemoryEffectFree(valOp) &&
+      if (valOp && mlir::isMemoryEffectFree(valOp) &&
           !mlir::isa<fir::BoxDimsOp>(valOp)) {
         mlir::Operation *clonedOp = valOp->clone();
         entryBlock->push_front(clonedOp);
@@ -1447,7 +1367,13 @@ static void genBodyOfTargetOp(
         valOp->replaceUsesWithIf(clonedOp, replace);
       } else {
         auto savedIP = firOpBuilder.getInsertionPoint();
-        firOpBuilder.setInsertionPointAfter(valOp);
+
+        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);
diff --git a/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp b/flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp
@@ -47,6 +47,8 @@
 #include <iterator>
 #include <numeric>
 
+#define DEBUG_TYPE "omp-map-info-finalization"
+#define PDBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "]: ")
 namespace flangomp {
 #define GEN_PASS_DEF_MAPINFOFINALIZATIONPASS
 #include "flang/Optimizer/OpenMP/Passes.h.inc"
@@ -554,52 +556,31 @@ class MapInfoFinalizationPass
       // with an underlying type of fir.char<k, ?>, i.e a character
       // with dynamic length. If so, check if they need bounds added.
       func->walk([&](mlir::omp::MapInfoOp op) {
-        mlir::Value varPtr = op.getVarPtr();
-        mlir::Type underlyingVarType = fir::unwrapRefType(varPtr.getType());
-        if (!mlir::isa<fir::CharacterType>(underlyingVarType))
+        if (!op.getBounds().empty())
           return mlir::WalkResult::advance();
 
-        fir::CharacterType cType =
-            mlir::cast<fir::CharacterType>(underlyingVarType);
-        if (!cType.hasDynamicLen())
-          return mlir::WalkResult::advance();
+        mlir::Value varPtr = op.getVarPtr();
+        mlir::Type underlyingVarType = fir::unwrapRefType(varPtr.getType());
 
-        if (!op.getBounds().empty())
-          return mlir::WalkResult::advance();
-        // This means varPtr is a BlockArgument. I do not know how to get to a
-        // fir.boxchar<> type of mlir::Value for varPtr. So, skipping this for
-        // now.
-        mlir::Operation *definingOp = varPtr.getDefiningOp();
-        if (!definingOp)
+        if (!fir::characterWithDynamicLen(underlyingVarType))
           return mlir::WalkResult::advance();
 
-        if (auto declOp = mlir::dyn_cast<hlfir::DeclareOp>(definingOp)) {
-          mlir::Value base = declOp.getBase();
-          assert(mlir::isa<fir::BoxCharType>(base.getType()));
-          //          mlir::value unboxChar
-          builder.setInsertionPoint(op);
-          fir::BoxCharType boxCharType =
-              mlir::cast<fir::BoxCharType>(base.getType());
-          mlir::Type idxTy = builder.getIndexType();
-          mlir::Type lenType = builder.getCharacterLengthType();
-          mlir::Type refType = builder.getRefType(boxCharType.getEleTy());
-          mlir::Location location = op.getLoc();
-          auto unboxed = builder.create<fir::UnboxCharOp>(location, refType,
-                                                          lenType, base);
-          //          len = unboxed.getResult(1);
-          mlir::Value zero = builder.createIntegerConstant(location, idxTy, 0);
-          mlir::Value one = builder.createIntegerConstant(location, idxTy, 1);
-          mlir::Value extent = unboxed.getResult(1);
-          mlir::Value stride = one;
-          mlir::Value ub =
-              builder.create<mlir::arith::SubIOp>(location, extent, one);
-          mlir::Type boundTy = builder.getType<mlir::omp::MapBoundsType>();
-          mlir::Value boundsOp = builder.create<mlir::omp::MapBoundsOp>(
-              location, boundTy, /*lower_bound=*/zero,
-              /*upper_bound=*/ub, /*extent=*/extent, /*stride=*/stride,
-              /*stride_in_bytes = */ true, /*start_idx=*/zero);
-          op.getBoundsMutable().append({boundsOp});
-        }
+        fir::factory::AddrAndBoundsInfo info =
+            fir::factory::getDataOperandBaseAddr(
+                builder, varPtr, /*isOptional=*/false, varPtr.getLoc());
+        fir::ExtendedValue extendedValue =
+            hlfir::translateToExtendedValue(varPtr.getLoc(), builder,
+                                            hlfir::Entity{info.addr},
+                                            /*continguousHint=*/true)
+                .first;
+        builder.setInsertionPoint(op);
+        llvm::SmallVector<mlir::Value> boundsOps =
+            fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
+                                               mlir::omp::MapBoundsType>(
+                builder, info, extendedValue,
+                /*dataExvIsAssumedSize=*/false, varPtr.getLoc());
+
+        op.getBoundsMutable().append(boundsOps);
         return mlir::WalkResult::advance();
       });
 
diff --git a/flang/test/Lower/OpenMP/map-character.f90 b/flang/test/Lower/OpenMP/map-character.f90
