[flang][OpenMP] Use enum for DoConcurrentConversionPass::mapTo (#96)

Cleans-up do-concurrent conversion code a bit by moving from using `std::string` for the `mapTo` option to using an `enum. The same `enum` is used across the code for `flang-new`, `fir-opt`, and `bbc`.
ROCm · Jun 4, 2024 · a400728 · a400728
1 parent 9071627
commit a400728
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 13 deletions.
diff --git a/flang/include/flang/Frontend/CodeGenOptions.h b/flang/include/flang/Frontend/CodeGenOptions.h
@@ -15,6 +15,7 @@
 #ifndef FORTRAN_FRONTEND_CODEGENOPTIONS_H
 #define FORTRAN_FRONTEND_CODEGENOPTIONS_H
 
+#include "flang/Optimizer/Transforms/Utils.h"
 #include "llvm/Frontend/Debug/Options.h"
 #include "llvm/Frontend/Driver/CodeGenOptions.h"
 #include "llvm/Support/CodeGen.h"
@@ -131,11 +132,7 @@ class CodeGenOptions : public CodeGenOptionsBase {
 
   /// Optionally map `do concurrent` loops to OpenMP. This is only valid of
   /// OpenMP is enabled.
-  enum class DoConcurrentMappingKind {
-    DCMK_None,  // Do not lower `do concurrent` to OpenMP.
-    DCMK_Host,  // Lower to run in parallel on the CPU.
-    DCMK_Device // Lower to run in parallel on the GPU.
-  };
+  using DoConcurrentMappingKind = fir::omp::DoConcurrentMappingKind;
 
   // Define accessors/mutators for code generation options of enumeration type.
 #define CODEGENOPT(Name, Bits, Default)

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.h b/flang/include/flang/Optimizer/Transforms/Passes.h
@@ -10,10 +10,12 @@
 #define FORTRAN_OPTIMIZER_TRANSFORMS_PASSES_H
 
 #include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Transforms/Utils.h"
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
+
 #include <memory>
 
 namespace mlir {

diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -15,6 +15,7 @@
 #define FLANG_OPTIMIZER_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
+include "mlir/IR/EnumAttr.td"
 
 def AbstractResultOpt
   : Pass<"abstract-result"> {
@@ -413,16 +414,23 @@ def DoConcurrentConversionPass : Pass<"fopenmp-do-concurrent-conversion", "mlir:
 
      Still to TODO:
        - More extensive testing.
-       - Mapping to `target teams distribute parallel do`.
-       - Allowing the user to control mapping behavior: either to the host or
-         target.
   }];
 
   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)">,
+    Option<"mapTo", "map-to",
+           "fir::omp::DoConcurrentMappingKind",
+           /*default=*/"fir::omp::DoConcurrentMappingKind::DCMK_None",
+           "Try to map `do concurrent` loops to OpenMP (on host or device)",
+           [{::llvm::cl::values(
+               clEnumValN(fir::omp::DoConcurrentMappingKind::DCMK_None,
+                          "none", "Do not lower `do concurrent` to OpenMP"),
+               clEnumValN(fir::omp::DoConcurrentMappingKind::DCMK_Host,
+                          "host", "Lower to run in parallel on the CPU"),
+               clEnumValN(fir::omp::DoConcurrentMappingKind::DCMK_Device,
+                          "device", "Lower to run in parallel on the GPU")
+           )}]>,
   ];
 }
 

diff --git a/flang/include/flang/Optimizer/Transforms/Utils.h b/flang/include/flang/Optimizer/Transforms/Utils.h
@@ -13,8 +13,13 @@
 #ifndef FORTRAN_OPTIMIZER_TRANSFORMS_UTILS_H
 #define FORTRAN_OPTIMIZER_TRANSFORMS_UTILS_H
 
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+
 namespace fir {
 
+class FirOpBuilder;
+
 using MinlocBodyOpGeneratorTy = llvm::function_ref<mlir::Value(
     fir::FirOpBuilder &, mlir::Location, const mlir::Type &, mlir::Value,
     mlir::Value, mlir::Value, const llvm::SmallVectorImpl<mlir::Value> &)>;
@@ -33,6 +38,13 @@ void genMinMaxlocReductionLoop(fir::FirOpBuilder &builder, mlir::Value array,
                                mlir::Type maskElemType, mlir::Value resultArr,
                                bool maskMayBeLogicalScalar);
 
+namespace omp {
+enum class DoConcurrentMappingKind {
+  DCMK_None,  // Do not lower `do concurrent` to OpenMP.
+  DCMK_Host,  // Lower to run in parallel on the CPU.
+  DCMK_Device // Lower to run in parallel on the GPU.
+};
+}
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_TRANSFORMS_UTILS_H
diff --git a/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp b/flang/lib/Optimizer/Transforms/DoConcurrentConversion.cpp
@@ -581,15 +581,17 @@ class DoConcurrentConversionPass
 
     auto *context = &getContext();
 
-    if (mapTo != "host" && mapTo != "device") {
+    if (mapTo != fir::omp::DoConcurrentMappingKind::DCMK_Host &&
+        mapTo != fir::omp::DoConcurrentMappingKind::DCMK_Device) {
       mlir::emitWarning(mlir::UnknownLoc::get(context),
                         "DoConcurrentConversionPass: invalid `map-to` value. "
                         "Valid values are: `host` or `device`");
       return;
     }
 
     mlir::RewritePatternSet patterns(context);
-    patterns.insert<DoConcurrentConversion>(context, mapTo == "device");
+    patterns.insert<DoConcurrentConversion>(
+        context, mapTo == fir::omp::DoConcurrentMappingKind::DCMK_Device);
     mlir::ConversionTarget target(*context);
     target.addLegalDialect<fir::FIROpsDialect, hlfir::hlfirDialect,
                            mlir::arith::ArithDialect, mlir::func::FuncDialect,
@@ -611,7 +613,8 @@ class DoConcurrentConversionPass
 std::unique_ptr<mlir::Pass>
 fir::createDoConcurrentConversionPass(bool mapToDevice) {
   DoConcurrentConversionPassOptions options;
-  options.mapTo = mapToDevice ? "device" : "host";
+  options.mapTo = mapToDevice ? fir::omp::DoConcurrentMappingKind::DCMK_Device
+                              : fir::omp::DoConcurrentMappingKind::DCMK_Host;
 
   return std::make_unique<DoConcurrentConversionPass>(options);
 }