[Openmp-commits] [mlir] [openmp] [flang] [Flang][OpenMP] Initial mapping of Fortran pointers and allocatables for target devices (PR #71766)

[via Openmp-commits]

================
@@ -0,0 +1,145 @@
+#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"
+#include <iterator>
+
+namespace fir {
+#define GEN_PASS_DEF_OMPDESCRIPTORMAPINFOGENPASS
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+namespace {
+class OMPDescriptorMapInfoGenPass
+    : public fir::impl::OMPDescriptorMapInfoGenPassBase<
+          OMPDescriptorMapInfoGenPass> {
+
+  void genDescriptorMemberMaps(mlir::omp::MapInfoOp op,
+                               fir::FirOpBuilder &builder,
+                               mlir::Operation *target) {
+    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);
+
+    // Member of the descriptor pointing at the allocated data
+    mlir::Value baseAddr = builder.create<mlir::omp::MapInfoOp>(
+        loc, baseAddrAddr.getType(), baseAddrAddr,
+        llvm::cast<mlir::omp::PointerLikeType>(
+            fir::unwrapRefType(baseAddrAddr.getType()))
+            .getElementType(),
+        mlir::Value{}, mlir::SmallVector<mlir::Value>{}, op.getBounds(),
+        builder.getIntegerAttr(builder.getIntegerType(64, false),
+                               op.getMapType().value()),
+        builder.getAttr<mlir::omp::VariableCaptureKindAttr>(
+            mlir::omp::VariableCaptureKind::ByRef),
+        builder.getStringAttr("") /*name*/);
+
+    // 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().append(baseAddr);
+    op.getBoundsMutable().assign(llvm::SmallVector<mlir::Value>{});
----------------
agozillon wrote:

There's a reasonable amount of usage of the Mutable interface in transformation passes in the wider MLIR project, so I am not so sure if there's a requirement for it in transformation passes (but I wasn't aware there was for conversion patterns either).

I can make the change if you'd like though, but it is going to be a bit harder to read and maintain as the MapClauseOwningOpInterface segment will become more complex, you'd have to create specific checks for each of the operations that meet this interface and then create a new version of that operation to replace the old version with (e.g. TargetOp/DataEntryOp etc.), someone will also have to make sure that if they add a new operation that meets this interface, they appropriately update this pass to create a new version of their operation, as opposed to it generically being able to access the map operands list and append through the interface.   

https://github.com/llvm/llvm-project/pull/71766