[Flang][OpenMP][MLIR] Add option to switch off deferred mapping of descriptors

Author: agozillon

clang/include/clang/Driver/Options.td

@@ -7305,6 +7305,15 @@ def fhermetic_module_files : Flag<["-"], "fhermetic-module-files">, Group<f_Grou
 def fdo_concurrent_to_openmp_EQ : Joined<["-"], "fdo-concurrent-to-openmp=">,
   HelpText<"Try to map `do concurrent` loops to OpenMP [none|host|device]">,
       Values<"none, host, device">;
+
+def fdefer_desc_map : Flag<["-"], "fdefer-desc-map">, Group<f_Group>,
+  HelpText<"Enable deferred descriptor mapping, which puts off top-level descriptor "
+    "mapping until target regions, this is the default behaviour">;
+
+def fno_defer_desc_map : Flag<["-"], "fno-defer-desc-map">, Group<f_Group>,
+  HelpText<"Disable deferred deferred descriptor mapping, which puts off top-level "
+          " descriptor mapping till target regions">;
+
 } // let Visibility = [FC1Option, FlangOption]
 
 def J : JoinedOrSeparate<["-"], "J">,

clang/lib/Driver/ToolChains/Flang.cpp

@@ -184,7 +184,8 @@ void Flang::addCodegenOptions(const ArgList &Args,
        options::OPT_frepack_arrays_contiguity_EQ,
        options::OPT_fstack_repack_arrays, options::OPT_fno_stack_repack_arrays,
        options::OPT_ftime_report, options::OPT_ftime_report_EQ,
-       options::OPT_funroll_loops, options::OPT_fno_unroll_loops});
+       options::OPT_funroll_loops, options::OPT_fno_unroll_loops,
+       options::OPT_fdefer_desc_map, options::OPT_fno_defer_desc_map});
 }
 
 void Flang::addPicOptions(const ArgList &Args, ArgStringList &CmdArgs) const {

flang/include/flang/Frontend/CodeGenOptions.def

@@ -42,6 +42,8 @@ CODEGENOPT(AliasAnalysis, 1, 0) ///< Enable alias analysis pass
 
 CODEGENOPT(Underscoring, 1, 1)
 CODEGENOPT(OffloadGlobalFiltering, 1, 0)
+CODEGENOPT(DeferDescriptorMapping, 1, 0) ///< Fortran OpenMP specific optimisation for delaying descriptor mapping until target/target data regions
+
 ENUM_CODEGENOPT(RelocationModel, llvm::Reloc::Model, 3, llvm::Reloc::PIC_) ///< Name of the relocation model to use.
 ENUM_CODEGENOPT(DebugInfo,  llvm::codegenoptions::DebugInfoKind, 4,  llvm::codegenoptions::NoDebugInfo) ///< Level of debug info to generate
 ENUM_CODEGENOPT(VecLib, llvm::driver::VectorLibrary, 4, llvm::driver::VectorLibrary::NoLibrary) ///< Vector functions library to use

flang/include/flang/Optimizer/OpenMP/Passes.h

@@ -33,6 +33,8 @@ namespace flangomp {
 bool shouldUseWorkshareLowering(mlir::Operation *op);
 
 std::unique_ptr<mlir::Pass> createDoConcurrentConversionPass(bool mapToDevice);
+
+std::unique_ptr<mlir::Pass> createMapInfoFinalizationPass(bool deferDescMap);
 } // namespace flangomp
 
 #endif // FORTRAN_OPTIMIZER_OPENMP_PASSES_H

flang/include/flang/Optimizer/OpenMP/Passes.td

@@ -21,6 +21,12 @@ def MapInfoFinalizationPass
     explicit individual mapping by the OpenMP runtime.
   }];
   let dependentDialects = ["mlir::omp::OpenMPDialect"];
+
+  let options = [Option<"deferDescMapping", "opt-defer-desc-mapping",
+                      "bool", /*default=*/"true",
+                      "Activates or deactivates deferred descriptor mapping, "
+                      "which delays mapping of top-level descriptors to target "
+                      "regions and target data regions">];
 }
 
 def MapsForPrivatizedSymbolsPass

flang/include/flang/Optimizer/Passes/Pipelines.h

@@ -134,6 +134,11 @@ struct OpenMPFIRPassPipelineOpts {
   bool isTargetDevice;
   bool enableOffloadGlobalFiltering;
 
+  /// Deactivates or activates MapInfoFinalization passes removal of
+  /// top-level descriptor mapping for non-Target Data/Target region
+  /// directives.
+  bool deferDescMap;
+
   /// Controls how to map `do concurrent` loops; to device, host, or none at
   /// all.
   Fortran::frontend::CodeGenOptions::DoConcurrentMappingKind

flang/lib/Frontend/CompilerInvocation.cpp

@@ -488,6 +488,10 @@ static void parseCodeGenArgs(Fortran::frontend::CodeGenOptions &opts,
   }
 
   parseDoConcurrentMapping(opts, args, diags);
+
+  opts.DeferDescriptorMapping =
+      args.hasFlag(clang::driver::options::OPT_fdefer_desc_map,
+                   clang::driver::options::OPT_fno_defer_desc_map, true);
 }
 
 /// Parses all target input arguments and populates the target

flang/lib/Frontend/FrontendActions.cpp

@@ -307,6 +307,8 @@ bool CodeGenAction::beginSourceFileAction() {
       ci.getInvocation().getCodeGenOpts().getDoConcurrentMapping();
   opts.enableOffloadGlobalFiltering =
       ci.getInvocation().getCodeGenOpts().OffloadGlobalFiltering;
+  opts.deferDescMap =
+      ci.getInvocation().getCodeGenOpts().DeferDescriptorMapping;
 
   if (opts.doConcurrentMappingKind != DoConcurrentMappingKind::DCMK_None &&
       !isOpenMPEnabled) {

flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp

@@ -58,6 +58,13 @@ namespace {
 class MapInfoFinalizationPass
     : public flangomp::impl::MapInfoFinalizationPassBase<
           MapInfoFinalizationPass> {
+public:
+  MapInfoFinalizationPass() = default;
+
+  MapInfoFinalizationPass(
+      const flangomp::MapInfoFinalizationPassOptions &options)
+      : MapInfoFinalizationPassBase(options) {}
+
   /// Helper class tracking a members parent and its
   /// placement in the parents member list
   struct ParentAndPlacement {
@@ -954,17 +961,20 @@ class MapInfoFinalizationPass
       // within a target region. At which point we map the relevant descriptor
       // data and the runtime should correctly associate the data with the
       // descriptor and bind together and allow clean mapping and execution.
-      func->walk([&](mlir::omp::MapInfoOp op) {
-        if (fir::isTypeWithDescriptor(op.getVarType()) ||
-            mlir::isa_and_present<fir::BoxAddrOp>(
-                op.getVarPtr().getDefiningOp())) {
-          mlir::Operation *targetUser = getFirstTargetUser(op);
-          assert(targetUser && "expected user of map operation was not found");
-          builder.setInsertionPoint(op);
-          removeTopLevelDescriptor(op, builder, targetUser);
-          addImplictDescriptorMapToTargetDataOp(op, builder, targetUser);
-        }
-      });
+      if (deferDescMapping) {
+        func->walk([&](mlir::omp::MapInfoOp op) {
+          if (fir::isTypeWithDescriptor(op.getVarType()) ||
+              mlir::isa_and_present<fir::BoxAddrOp>(
+                  op.getVarPtr().getDefiningOp())) {
+            mlir::Operation *targetUser = getFirstTargetUser(op);
+            assert(targetUser &&
+                   "expected user of map operation was not found");
+            builder.setInsertionPoint(op);
+            removeTopLevelDescriptor(op, builder, targetUser);
+            addImplictDescriptorMapToTargetDataOp(op, builder, targetUser);
+          }
+        });
+      }
 
       // Wait until after we have generated all of our maps to add them onto
       // the target's block arguments, simplifying the process as there would be
@@ -977,5 +987,11 @@ class MapInfoFinalizationPass
     });
   }
 };
-
 } // namespace
+
+std::unique_ptr<mlir::Pass>
+flangomp::createMapInfoFinalizationPass(bool deferDescMap) {
+  MapInfoFinalizationPassOptions options;
+  options.deferDescMapping = deferDescMap;
+  return std::make_unique<MapInfoFinalizationPass>(options);
+}

flang/lib/Optimizer/Passes/Pipelines.cpp

@@ -312,7 +312,7 @@ void createOpenMPFIRPassPipeline(mlir::PassManager &pm,
   // underlying data which is necessary to access the data on the offload
   // target device.
   pm.addPass(flangomp::createMapsForPrivatizedSymbolsPass());
-  pm.addPass(flangomp::createMapInfoFinalizationPass());
+  pm.addPass(flangomp::createMapInfoFinalizationPass(opts.deferDescMap));
   pm.addPass(flangomp::createMarkDeclareTargetPass());
   pm.addPass(flangomp::createGenericLoopConversionPass());
   if (opts.isTargetDevice) {