[flang][OpenMP] Extend do concurrent mapping to device. #50

For simple loops, we can now choose to map `do concurrent` to either the host (i.e. `omp parallel do`) or the device (i.e. `omp target teams distribute parallel do`).

In order to use this from `flang-new`, you can pass: `-fdo-concurrent-parallel=[none|host|device]`.

Author: ergawy

flang/lib/Lower/OpenMP/Clauses.h renamed to flang/include/flang/Lower/OpenMP/Clauses.h

@@ -59,7 +59,7 @@ struct OmpMapMemberIndicesData {
 };
 
 mlir::omp::MapInfoOp
-createMapInfoOp(fir::FirOpBuilder &builder, mlir::Location loc,
+createMapInfoOp(mlir::OpBuilder &builder, mlir::Location loc,
                 mlir::Value baseAddr, mlir::Value varPtrPtr, std::string name,
                 mlir::ArrayRef<mlir::Value> bounds,
                 mlir::ArrayRef<mlir::Value> members,
@@ -102,6 +102,15 @@ void genObjectList(const ObjectList &objects,
                    Fortran::lower::AbstractConverter &converter,
                    llvm::SmallVectorImpl<mlir::Value> &operands);
 
+// TODO: consider moving this to the `omp.loop_nest` op. Would be something like
+// this:
+//
+// ```
+// mlir::Value LoopNestOp::calculateTripCount(mlir::OpBuilder &builder,
+// mlir::OpBuilder::InsertPoint ip)
+// ```
+mlir::Value calculateTripCount(fir::FirOpBuilder &builder, mlir::Location loc,
+                               const mlir::omp::CollapseClauseOps &ops);
 } // namespace omp
 } // namespace lower
 } // namespace Fortran

flang/lib/Lower/OpenMP/Utils.h renamed to flang/include/flang/Lower/OpenMP/Utils.h

@@ -38,6 +38,7 @@ namespace fir {
 #define GEN_PASS_DECL_ARRAYVALUECOPY
 #define GEN_PASS_DECL_CHARACTERCONVERSION
 #define GEN_PASS_DECL_CFGCONVERSION
+#define GEN_PASS_DECL_DOCONCURRENTCONVERSIONPASS
 #define GEN_PASS_DECL_EXTERNALNAMECONVERSION
 #define GEN_PASS_DECL_MEMREFDATAFLOWOPT
 #define GEN_PASS_DECL_SIMPLIFYINTRINSICS
@@ -88,7 +89,7 @@ createFunctionAttrPass(FunctionAttrTypes &functionAttr, bool noInfsFPMath,
                        bool noNaNsFPMath, bool approxFuncFPMath,
                        bool noSignedZerosFPMath, bool unsafeFPMath);
 
-std::unique_ptr<mlir::Pass> createDoConcurrentConversionPass();
+std::unique_ptr<mlir::Pass> createDoConcurrentConversionPass(bool mapToDevice);
 
 void populateCfgConversionRewrites(mlir::RewritePatternSet &patterns,
                                    bool forceLoopToExecuteOnce = false);

flang/include/flang/Optimizer/Transforms/Passes.h

@@ -416,8 +416,12 @@ def DoConcurrentConversionPass : Pass<"fopenmp-do-concurrent-conversion", "mlir:
          target.
   }];
 
-  let constructor = "::fir::createDoConcurrentConversionPass()";
   let dependentDialects = ["mlir::omp::OpenMPDialect"];
+
+  let options = [
+    Option<"mapTo", "map-to", "std::string", "",
+           "Try to map `do concurrent` loops to OpenMP (on host or device)">,
+  ];
 }
 
 #endif // FLANG_OPTIMIZER_TRANSFORMS_PASSES

flang/include/flang/Optimizer/Transforms/Passes.td

@@ -324,6 +324,9 @@ inline void createHLFIRToFIRPassPipeline(
   pm.addPass(hlfir::createConvertHLFIRtoFIRPass());
 }
 
+using DoConcurrentMappingKind =
+    Fortran::frontend::CodeGenOptions::DoConcurrentMappingKind;
+
 /// Create a pass pipeline for handling certain OpenMP transformations needed
 /// prior to FIR lowering.
 ///
@@ -333,8 +336,12 @@ inline void createHLFIRToFIRPassPipeline(
 /// \param pm - MLIR pass manager that will hold the pipeline definition.
 /// \param isTargetDevice - Whether code is being generated for a target device
 /// rather than the host device.
-inline void createOpenMPFIRPassPipeline(
-    mlir::PassManager &pm, bool isTargetDevice) {
+inline void createOpenMPFIRPassPipeline(mlir::PassManager &pm,
+    bool isTargetDevice, DoConcurrentMappingKind doConcurrentMappingKind) {
+  if (doConcurrentMappingKind != DoConcurrentMappingKind::DCMK_None)
+    pm.addPass(fir::createDoConcurrentConversionPass(
+        doConcurrentMappingKind == DoConcurrentMappingKind::DCMK_Device));
+
   pm.addPass(fir::createOMPMapInfoFinalizationPass());
   pm.addPass(fir::createOMPMarkDeclareTargetPass());
   if (isTargetDevice)

flang/include/flang/Tools/CLOptions.inc

@@ -320,41 +320,34 @@ bool CodeGenAction::beginSourceFileAction() {
   // Add OpenMP-related passes
   // WARNING: These passes must be run immediately after the lowering to ensure
   // that the FIR is correct with respect to OpenMP operations/attributes.
-  bool isOpenMPEnabled = ci.getInvocation().getFrontendOpts().features.IsEnabled(
+  bool isOpenMPEnabled =
+      ci.getInvocation().getFrontendOpts().features.IsEnabled(
           Fortran::common::LanguageFeature::OpenMP);
+
+  using DoConcurrentMappingKind =
+      Fortran::frontend::CodeGenOptions::DoConcurrentMappingKind;
+  DoConcurrentMappingKind doConcurrentMappingKind =
+      ci.getInvocation().getCodeGenOpts().getDoConcurrentMapping();
+
+  if (doConcurrentMappingKind != DoConcurrentMappingKind::DCMK_None &&
+      !isOpenMPEnabled) {
+    unsigned diagID = ci.getDiagnostics().getCustomDiagID(
+        clang::DiagnosticsEngine::Warning,
+        "lowering `do concurrent` loops to OpenMP is only supported if "
+        "OpenMP is enabled");
+    ci.getDiagnostics().Report(diagID);
+  }
+
   if (isOpenMPEnabled) {
     bool isDevice = false;
     if (auto offloadMod = llvm::dyn_cast<mlir::omp::OffloadModuleInterface>(
             mlirModule->getOperation()))
       isDevice = offloadMod.getIsTargetDevice();
+
     // WARNING: This pipeline must be run immediately after the lowering to
     // ensure that the FIR is correct with respect to OpenMP operations/
     // attributes.
-    fir::createOpenMPFIRPassPipeline(pm, isDevice);
-  }
-
-  using DoConcurrentMappingKind =
-      Fortran::frontend::CodeGenOptions::DoConcurrentMappingKind;
-  DoConcurrentMappingKind selectedKind = ci.getInvocation().getCodeGenOpts().getDoConcurrentMapping();
-  if (selectedKind != DoConcurrentMappingKind::DCMK_None) {
-    if (!isOpenMPEnabled) {
-      unsigned diagID = ci.getDiagnostics().getCustomDiagID(
-          clang::DiagnosticsEngine::Warning,
-          "lowering `do concurrent` loops to OpenMP is only supported if "
-          "OpenMP is enabled");
-      ci.getDiagnostics().Report(diagID);
-    } else {
-      bool mapToDevice = selectedKind == DoConcurrentMappingKind::DCMK_Device;
-
-      if (mapToDevice) {
-        unsigned diagID = ci.getDiagnostics().getCustomDiagID(
-            clang::DiagnosticsEngine::Warning,
-            "TODO: lowering `do concurrent` loops to OpenMP device is not "
-            "supported yet");
-        ci.getDiagnostics().Report(diagID);
-      } else
-        pm.addPass(fir::createDoConcurrentConversionPass());
-    }
+    fir::createOpenMPFIRPassPipeline(pm, isDevice, doConcurrentMappingKind);
   }
 
   pm.enableVerifier(/*verifyPasses=*/true);

flang/lib/Frontend/FrontendActions.cpp

@@ -11,8 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "ClauseProcessor.h"
-#include "Clauses.h"
 
+#include "flang/Lower/OpenMP/Clauses.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Parser/tools.h"
 #include "flang/Semantics/tools.h"

flang/lib/Lower/OpenMP/ClauseProcessor.cpp

@@ -12,12 +12,12 @@
 #ifndef FORTRAN_LOWER_CLAUASEPROCESSOR_H
 #define FORTRAN_LOWER_CLAUASEPROCESSOR_H
 
-#include "Clauses.h"
 #include "DirectivesCommon.h"
 #include "ReductionProcessor.h"
-#include "Utils.h"
 #include "flang/Lower/AbstractConverter.h"
 #include "flang/Lower/Bridge.h"
+#include "flang/Lower/OpenMP/Clauses.h"
+#include "flang/Lower/OpenMP/Utils.h"
 #include "flang/Optimizer/Builder/Todo.h"
 #include "flang/Parser/dump-parse-tree.h"
 #include "flang/Parser/parse-tree.h"

flang/lib/Lower/OpenMP/ClauseProcessor.h

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "Clauses.h"
+#include "flang/Lower/OpenMP/Clauses.h"
 
 #include "flang/Common/idioms.h"
 #include "flang/Evaluate/expression.h"

flang/lib/Lower/OpenMP/Clauses.cpp

@@ -12,7 +12,7 @@
 
 #include "DataSharingProcessor.h"
 
-#include "Utils.h"
+#include "flang/Lower/OpenMP/Utils.h"
 #include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/SymbolMap.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"