[Flang][MLIR] Move majority of OpenMP lowering of descriptor types into Opt Pass

Author: agozillon

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

@@ -76,6 +76,7 @@ std::unique_ptr<mlir::Pass>
 createAlgebraicSimplificationPass(const mlir::GreedyRewriteConfig &config);
 std::unique_ptr<mlir::Pass> createPolymorphicOpConversionPass();
 
+std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass();
 std::unique_ptr<mlir::Pass> createOMPFunctionFilteringPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 createOMPMarkDeclareTargetPass();

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

@@ -318,6 +318,17 @@ def LoopVersioning : Pass<"loop-versioning", "mlir::func::FuncOp"> {
   let dependentDialects = [ "fir::FIROpsDialect" ];
 }
 
+def OMPDescriptorMapInfoGenPass
+    : Pass<"omp-descriptor-map-info-gen", "mlir::ModuleOp"> {
+  let summary = "expands OpenMP MapInfo operations containing descriptors";
+  let description = [{
+    Expands MapInfo operations containing descriptor types into multiple MapInfo's for each pointer element in 
+    the descriptor that requires explicit individual mapping by the OpenMP runtime.
+  }];
+  let constructor = "::fir::createOMPDescriptorMapInfoGenPass()";
+  let dependentDialects = ["mlir::omp::OpenMPDialect"];
+}
+
 def OMPMarkDeclareTargetPass
     : Pass<"omp-mark-declare-target", "mlir::ModuleOp"> {
   let summary = "Marks all functions called by an OpenMP declare target function as declare target";

flang/include/flang/Tools/CLOptions.inc

@@ -274,6 +274,7 @@ inline void createHLFIRToFIRPassPipeline(
 /// rather than the host device.
 inline void createOpenMPFIRPassPipeline(
     mlir::PassManager &pm, bool isTargetDevice) {
+  pm.addPass(fir::createOMPDescriptorMapInfoGenPass());
   pm.addPass(fir::createOMPMarkDeclareTargetPass());
   if (isTargetDevice)
     pm.addPass(fir::createOMPFunctionFilteringPass());

flang/lib/Lower/OpenMP.cpp

@@ -1740,54 +1740,6 @@ createMapInfoOp(fir::FirOpBuilder &builder, mlir::Location loc,
   return op;
 }
 
-static mlir::omp::MapInfoOp processDescriptorTypeMappings(
-    Fortran::semantics::SemanticsContext &semanticsContext,
-    Fortran::lower::StatementContext &stmtCtx,
-    Fortran::lower::AbstractConverter &converter, mlir::Location loc,
-    mlir::Value descriptorAddr, mlir::Value descDataBaseAddr,
-    mlir::ValueRange bounds, std::string asFortran,
-    llvm::omp::OpenMPOffloadMappingFlags mapCaptureType) {
-  llvm::SmallVector<mlir::Value> descriptorBaseAddrMembers;
-  fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
-
-  mlir::Value descriptor = descriptorAddr;
-
-  // The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as
-  // allowing it to access non-reference box operations can cause some
-  // problematic SSA IR. However, in the case of assumed shape's the type
-  // is not a !fir.ref, in these cases to retrieve the appropriate
-  // !fir.ref<!fir.box<...>> to access the data we need to map we must
-  // perform an alloca and then store to it and retrieve the data from the new
-  // alloca.
-  if (mlir::isa<fir::BaseBoxType>(descriptorAddr.getType())) {
-    mlir::OpBuilder::InsertPoint insPt = firOpBuilder.saveInsertionPoint();
-    firOpBuilder.setInsertionPointToStart(firOpBuilder.getAllocaBlock());
-    descriptor =
-        firOpBuilder.create<fir::AllocaOp>(loc, descriptorAddr.getType());
-    firOpBuilder.restoreInsertionPoint(insPt);
-    firOpBuilder.create<fir::StoreOp>(loc, descriptorAddr, descriptor);
-  }
-
-  mlir::Value baseAddrAddr = firOpBuilder.create<fir::BoxOffsetOp>(
-      loc, descriptor, fir::BoxFieldAttr::base_addr);
-
-  descriptorBaseAddrMembers.push_back(createMapInfoOp(
-      firOpBuilder, loc, descDataBaseAddr, baseAddrAddr, asFortran, bounds, {},
-      static_cast<std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
-          mapCaptureType),
-      mlir::omp::VariableCaptureKind::ByRef, descDataBaseAddr.getType()));
-
-  // TODO: map the addendum segment of the descriptor, similarly to the above
-  // base address/data pointer member.
-
-  return createMapInfoOp(
-      firOpBuilder, loc, descriptor, mlir::Value{}, asFortran, {},
-      descriptorBaseAddrMembers,
-      static_cast<std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
-          mapCaptureType),
-      mlir::omp::VariableCaptureKind::ByRef, descriptor.getType());
-}
-
 bool ClauseProcessor::processMap(
     mlir::Location currentLocation, const llvm::omp::Directive &directive,
     Fortran::semantics::SemanticsContext &semanticsContext,
@@ -1857,24 +1809,20 @@ bool ClauseProcessor::processMap(
 
           auto origSymbol =
               converter.getSymbolAddress(*getOmpObjectSymbol(ompObject));
-          mlir::Value mapOp, symAddr;
-          if (origSymbol && fir::isTypeWithDescriptor(origSymbol.getType())) {
+          mlir::Value symAddr = info.addr;
+          if (origSymbol && fir::isTypeWithDescriptor(origSymbol.getType()))
             symAddr = origSymbol;
-            mapOp = processDescriptorTypeMappings(
-                semanticsContext, stmtCtx, converter, clauseLocation,
-                origSymbol, info.addr, bounds, asFortran.str(), mapTypeBits);
-          } else {
-            // Explicit map captures are captured ByRef by default,
-            // optimisation passes may alter this to ByCopy or other capture
-            // types to optimise
-            symAddr = info.addr;
-            mapOp = createMapInfoOp(
-                firOpBuilder, clauseLocation, info.addr, mlir::Value{},
-                asFortran.str(), bounds, {},
-                static_cast<std::underlying_type_t<
-                    llvm::omp::OpenMPOffloadMappingFlags>>(mapTypeBits),
-                mlir::omp::VariableCaptureKind::ByRef, info.addr.getType());
-          }
+
+          // Explicit map captures are captured ByRef by default,
+          // optimisation passes may alter this to ByCopy or other capture
+          // types to optimise
+          mlir::Value mapOp = createMapInfoOp(
+              firOpBuilder, clauseLocation, symAddr, mlir::Value{},
+              asFortran.str(), bounds, {},
+              static_cast<
+                  std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+                  mapTypeBits),
+              mlir::omp::VariableCaptureKind::ByRef, symAddr.getType());
 
           mapOperands.push_back(mapOp);
           if (mapSymTypes)
@@ -1989,22 +1937,20 @@ bool ClauseProcessor::processMotionClauses(
 
           auto origSymbol =
               converter.getSymbolAddress(*getOmpObjectSymbol(ompObject));
-          mlir::Value mapOp;
-          if (origSymbol && fir::isTypeWithDescriptor(origSymbol.getType())) {
-            mapOp = processDescriptorTypeMappings(
-                semanticsContext, stmtCtx, converter, clauseLocation,
-                origSymbol, info.addr, bounds, asFortran.str(), mapTypeBits);
-          } else {
-            // Explicit map captures are captured ByRef by default,
-            // optimisation passes may alter this to ByCopy or other capture
-            // types to optimise
-            mapOp = createMapInfoOp(
-                firOpBuilder, clauseLocation, info.addr, mlir::Value{},
-                asFortran.str(), bounds, {},
-                static_cast<std::underlying_type_t<
-                    llvm::omp::OpenMPOffloadMappingFlags>>(mapTypeBits),
-                mlir::omp::VariableCaptureKind::ByRef, info.addr.getType());
-          }
+          mlir::Value symAddr = info.addr;
+          if (origSymbol && fir::isTypeWithDescriptor(origSymbol.getType()))
+            symAddr = origSymbol;
+
+          // Explicit map captures are captured ByRef by default,
+          // optimisation passes may alter this to ByCopy or other capture
+          // types to optimise
+          mlir::Value mapOp = createMapInfoOp(
+              firOpBuilder, clauseLocation, symAddr, mlir::Value{},
+              asFortran.str(), bounds, {},
+              static_cast<
+                  std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+                  mapTypeBits),
+              mlir::omp::VariableCaptureKind::ByRef, symAddr.getType());
 
           mapOperands.push_back(mapOp);
         }
@@ -2819,20 +2765,13 @@ genTargetOp(Fortran::lower::AbstractConverter &converter,
           mapFlag |= llvm::omp::OpenMPOffloadMappingFlags::OMP_MAP_FROM;
         }
 
-        mlir::Value mapOp;
-        if (fir::isTypeWithDescriptor(baseOp.getType())) {
-          mapOp = processDescriptorTypeMappings(
-              semanticsContext, stmtCtx, converter, baseOp.getLoc(), baseOp,
-              info.addr, bounds, name.str(), mapFlag);
-        } else {
-          mapOp = createMapInfoOp(
-              converter.getFirOpBuilder(), baseOp.getLoc(), baseOp,
-              mlir::Value{}, name.str(), bounds, {},
-              static_cast<
-                  std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
-                  mapFlag),
-              captureKind, baseOp.getType());
-        }
+        mlir::Value mapOp = createMapInfoOp(
+            converter.getFirOpBuilder(), baseOp.getLoc(), baseOp, mlir::Value{},
+            name.str(), bounds, {},
+            static_cast<
+                std::underlying_type_t<llvm::omp::OpenMPOffloadMappingFlags>>(
+                mapFlag),
+            captureKind, baseOp.getType());
 
         mapOperands.push_back(mapOp);
         mapSymTypes.push_back(baseOp.getType());

flang/lib/Optimizer/Transforms/CMakeLists.txt

@@ -17,6 +17,7 @@ add_flang_library(FIRTransforms
   AddDebugFoundation.cpp
   PolymorphicOpConversion.cpp
   LoopVersioning.cpp
+  OMPDescriptorMapInfoGen.cpp
   OMPFunctionFiltering.cpp
   OMPMarkDeclareTarget.cpp
   VScaleAttr.cpp

flang/lib/Optimizer/Transforms/OMPDescriptorMapInfoGen.cpp

@@ -0,0 +1,123 @@
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace fir {
+#define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+class OMPDescriptorMapInfoGenPass
+    : public fir::impl::OMPDescriptorMapInfoGenPassBase<
+          OMPDescriptorMapInfoGenPass> {
+
+  mlir::omp::MapInfoOp
+  createMapInfoOp(fir::FirOpBuilder &builder, mlir::Location loc,
+                  mlir::Value baseAddr, mlir::Value varPtrPtr, std::string name,
+                  mlir::SmallVector<mlir::Value> bounds,
+                  mlir::SmallVector<mlir::Value> members, uint64_t mapType,
+                  mlir::omp::VariableCaptureKind mapCaptureType,
+                  mlir::Type retTy, bool isVal = false) {
+    if (auto boxTy = baseAddr.getType().dyn_cast<fir::BaseBoxType>()) {
+      baseAddr = builder.create<fir::BoxAddrOp>(loc, baseAddr);
+      retTy = baseAddr.getType();
+    }
+
+    mlir::TypeAttr varType = mlir::TypeAttr::get(
+        llvm::cast<mlir::omp::PointerLikeType>(retTy).getElementType());
+
+    mlir::omp::MapInfoOp op = builder.create<mlir::omp::MapInfoOp>(
+        loc, retTy, baseAddr, varType, varPtrPtr, members, bounds,
+        builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
+        builder.getAttr<mlir::omp::VariableCaptureKindAttr>(mapCaptureType),
+        builder.getStringAttr(name));
+
+    return op;
+  }
+
+  void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
+                               fir::FirOpBuilder &builder) {
+    llvm::SmallVector<mlir::Value> descriptorBaseAddrMembers;
+    mlir::Location loc = builder.getUnknownLoc();
+    mlir::Value descriptor = op.getVarPtr();
+
+    // If we enter this function, but the mapped type itself is not the
+    // descriptor, then it's likely the address of the descriptor so we
+    // must retrieve the descriptor SSA.
+    if (!fir::isTypeWithDescriptor(op.getVarType())) {
+      if (auto addrOp = mlir::dyn_cast_if_present<fir::BoxAddrOp>(
+              op.getVarPtr().getDefiningOp())) {
+        descriptor = addrOp.getVal();
+      }
+    }
+
+    // The fir::BoxOffsetOp only works with !fir.ref<!fir.box<...>> types, as
+    // allowing it to access non-reference box operations can cause some
+    // problematic SSA IR. However, in the case of assumed shape's the type
+    // is not a !fir.ref, in these cases to retrieve the appropriate
+    // !fir.ref<!fir.box<...>> to access the data we need to map we must
+    // perform an alloca and then store to it and retrieve the data from the new
+    // alloca.
+    if (mlir::isa<fir::BaseBoxType>(descriptor.getType())) {
+      mlir::OpBuilder::InsertPoint insPt = builder.saveInsertionPoint();
+      builder.setInsertionPointToStart(builder.getAllocaBlock());
+      auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
+      builder.restoreInsertionPoint(insPt);
+      builder.create<fir::StoreOp>(loc, descriptor, alloca);
+      descriptor = alloca;
+    }
+
+    mlir::Value baseAddrAddr = builder.create<fir::BoxOffsetOp>(
+        loc, descriptor, fir::BoxFieldAttr::base_addr);
+
+    descriptorBaseAddrMembers.push_back(createMapInfoOp(
+        builder, loc, baseAddrAddr, {}, "", op.getBounds(), {},
+        op.getMapType().value(), mlir::omp::VariableCaptureKind::ByRef,
+        fir::unwrapRefType(baseAddrAddr.getType())));
+
+    // TODO: map the addendum segment of the descriptor, similarly to the above
+    // base address/data pointer member.
+
+    op.getVarPtrMutable().assign(descriptor);
+    op.setVarType(fir::unwrapRefType(descriptor.getType()));
+    op.getMembersMutable().assign(descriptorBaseAddrMembers);
+    op.getBoundsMutable().assign(llvm::SmallVector<mlir::Value>{});
+  }
+
+  // This pass executes on mlir::ModuleOp's finding omp::MapInfoOp's containing
+  // descriptor based types (allocatables, pointers, assumed shape etc.) and
+  // expanding them into multiple omp::MapInfoOp's for each pointer member
+  // contained within the descriptor.
+  void runOnOperation() override {
+    fir::KindMapping kindMap = fir::getKindMapping(getOperation());
+    fir::FirOpBuilder builder{getOperation(), std::move(kindMap)};
+
+    getOperation()->walk([&](mlir::omp::MapInfoOp op) {
+      if (fir::isTypeWithDescriptor(op.getVarType()) ||
+          mlir::isa<fir::BoxAddrOp>(op.getVarPtr().getDefiningOp())) {
+        builder.setInsertionPoint(op);
+        genDescriptorMemberMaps(op, builder);
+      }
+    });
+  }
+};
+
+} // namespace
+
+namespace fir {
+std::unique_ptr<mlir::Pass> createOMPDescriptorMapInfoGenPass() {
+  return std::make_unique<OMPDescriptorMapInfoGenPass>();
+}
+} // namespace fir

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -1584,6 +1584,8 @@ mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
   if (failed(translator.convertFunctions()))
     return nullptr;
 
+  // translator.llvmModule->dump();
+
   if (llvm::verifyModule(*translator.llvmModule, &llvm::errs()))
     return nullptr;