[mlir][OpenMP] - MLIR to LLVMIR translation support for delayed privatization of allocatables in omp.target ops

flang/lib/Optimizer/OpenMP/MapsForPrivatizedSymbols.cpp

@@ -49,13 +49,6 @@ class MapsForPrivatizedSymbolsPass
     : public flangomp::impl::MapsForPrivatizedSymbolsPassBase<
           MapsForPrivatizedSymbolsPass> {
 
-  bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) {
-    Region &allocRegion = privatizer.getAllocRegion();
-    Value blockArg0 = allocRegion.getArgument(0);
-    if (blockArg0.use_empty())
-      return false;
-    return true;
-  }
   omp::MapInfoOp createMapInfo(Location loc, Value var,
                                fir::FirOpBuilder &builder) {
     uint64_t mapTypeTo = static_cast<
@@ -132,9 +125,9 @@ class MapsForPrivatizedSymbolsPass
         omp::PrivateClauseOp privatizer =
             SymbolTable::lookupNearestSymbolFrom<omp::PrivateClauseOp>(
                 targetOp, privatizerName);
-        if (!privatizerNeedsMap(privatizer)) {
+        if (!privatizer.needsMap())
           continue;
-        }
+
         builder.setInsertionPoint(targetOp);
         Location loc = targetOp.getLoc();
         omp::MapInfoOp mapInfoOp = createMapInfo(loc, privVar, builder);

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

@@ -135,6 +135,17 @@ def PrivateClauseOp : OpenMP_Op<"private", [IsolatedFromAbove, RecipeInterface]>
       auto &region = getDeallocRegion();
       return region.empty() ? nullptr : region.getArgument(0);
     }
+    /// privatizer is a PrivateClauseOp that privatizes an MLIR value.
+    /// privatizerNeedsMap returns true if the value being privatized in an
+    /// omp.target p should additionally be mapped to the target region
+    /// using a MapInfoOp. This is most common when an allocatable is privatized.
+    /// In such cases, the descriptor is use in privatization and needs to be
+    /// mapped on to the device.
+    bool needsMap() {
+      Value blockArg0 = getAllocRegion().getArgument(0);
+      return !blockArg0.use_empty();
+    }
+
   }];
 
   let hasRegionVerifier = 1;

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -3316,19 +3316,25 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
 
   return builder.saveIP();
 }
-static bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) {
-  Region &allocRegion = privatizer.getAllocRegion();
-  Value blockArg0 = allocRegion.getArgument(0);
-  return !blockArg0.use_empty();
-}
-
-// Return the llvm::Value * corresponding to the privateVar that
-// is being privatized. It isn't always as simple as looking up
-// moduleTranslation with privateVar. For instance, in case of
-// an allocatable, the descriptor for the allocatable is privatized.
-// This descriptor is mapped using an MapInfoOp. So, this function
-// will return a pointer to the llvm::Value corresponding to the
-// block argument for the mapped descriptor.
+/// privatizer is a PrivateClauseOp that privatizes an MLIR value.
+/// privatizerNeedsMap returns true if the value being privatized in an
+/// omp.target p should additionally be mapped to the target region
+/// using a MapInfoOp. This is most common when an allocatable is privatized.
+/// In such cases, the descriptor is use in privatization and needs to be
+/// mapped on to the device.
+// static bool privatizerNeedsMap(omp::PrivateClauseOp &privatizer) {
+//   Region &allocRegion = privatizer.getAllocRegion();
+//   Value blockArg0 = allocRegion.getArgument(0);
+//   return !blockArg0.use_empty();
+// }
+
+/// Return the llvm::Value * corresponding to the privateVar that
+/// is being privatized. It isn't always as simple as looking up
+/// moduleTranslation with privateVar. For instance, in case of
+/// an allocatable, the descriptor for the allocatable is privatized.
+/// This descriptor is mapped using an MapInfoOp. So, this function
+/// will return a pointer to the llvm::Value corresponding to the
+/// block argument for the mapped descriptor.
 static llvm::Value *
 findHostAssociatedValue(Value privateVar, omp::TargetOp targetOp,
                         llvm::DenseMap<Value, int> &mappedPrivateVars,
@@ -3414,7 +3420,7 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
       SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym);
       omp::PrivateClauseOp privatizer =
           findPrivatizer(targetOp, privatizerName);
-      if (!privatizerNeedsMap(privatizer))
+      if (!privatizer.needsMap())
         continue;
 
       // The MapInfoOp defining the map var isn't really needed later.

mlir/test/Target/LLVMIR/openmp-target-multiple-private.mlir

@@ -53,16 +53,14 @@ llvm.func @target_allocatable_(%arg0: !llvm.ptr {fir.bindc_name = "lb"}, %arg1:
     %65 = llvm.alloca %64 x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
     %67 = llvm.alloca %64 x i32 : (i32) -> !llvm.ptr
     %66 = llvm.mlir.constant(19 : i32) : i32
-    %68 = llvm.mlir.constant(18 : i32) : i32
     %69 = llvm.mlir.constant(10 : i32) : i32
-    %70 = llvm.mlir.constant(5 : i32) : i32
-    llvm.store %70, %arg3 : i32, !llvm.ptr
+    llvm.store %64, %arg3 : i32, !llvm.ptr
     llvm.store %69, %67 : i32, !llvm.ptr
     %75 = llvm.mlir.undef : !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)>
     %90 = llvm.insertvalue %67, %75[0] : !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)> 
     llvm.store %90, %65 : !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8)>, !llvm.ptr
     %91 = llvm.mlir.zero : !llvm.ptr
-    %92 = llvm.call @_FortranAAssign(%arg6, %65, %91, %68) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i32) -> !llvm.struct<()>
+    %92 = llvm.call @_FortranAAssign(%arg6, %65, %91, %66) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i32) -> !llvm.struct<()>
     %93 = llvm.call @_FortranAAssign(%arg7, %65, %91, %66) : (!llvm.ptr, !llvm.ptr, !llvm.ptr, i32) -> !llvm.struct<()>
     omp.terminator
   }