[flang-commits] [flang] [llvm] [mlir] [Flang][OpenMP] Fix Fortran automap handling (PR #162501)

[via flang-commits]

================
@@ -113,25 +110,89 @@ class AutomapToTargetDataPass
       if (needsBoundsOps(memOp.getMemref()))
         genBoundsOps(builder, memOp.getMemref(), bounds);
 
+      mlir::Value boxValue;
+      if (auto storeOp = mlir::dyn_cast<fir::StoreOp>(memOp.getOperation()))
+        boxValue = storeOp.getValue();
+      else
+        boxValue = mlir::cast<fir::LoadOp>(memOp.getOperation()).getResult();
+
+      mlir::Value baseAddr =
+          fir::BoxAddrOp::create(builder, memOp.getLoc(), boxValue);
+      mlir::Value dataAddr = builder.createConvert(
+          memOp.getLoc(),
+          builder.getRefType(fir::unwrapRefType(baseAddr.getType())), baseAddr);
+      mlir::Type baseTy = fir::unwrapRefType(dataAddr.getType());
+      if (mlir::Type eleTy = fir::dyn_cast_ptrOrBoxEleTy(baseTy))
+        baseTy = eleTy;
+      if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(baseTy))
+        if (seqTy.hasDynamicExtents())
+          baseTy = seqTy.getEleTy();
+
       omp::TargetEnterExitUpdateDataOperands clauses;
+      bool isAlloc = isa<fir::StoreOp>(memOp);
+
+      auto createDescriptorMap =
+          [&](mlir::omp::ClauseMapFlags mapType) -> mlir::omp::MapInfoOp {
+        mlir::Type descriptorTy =
+            fir::unwrapRefType(memOp.getMemref().getType());
+        // The attach entry expects the descriptor object to already have a
+        // device mapping, but this raw object map must not be expanded as a
+        // Fortran descriptor member map.
+        return mlir::omp::MapInfoOp::create(
+            builder, memOp.getLoc(), memOp.getMemref().getType(),
+            memOp.getMemref(), TypeAttr::get(descriptorTy),
+            builder.getAttr<omp::ClauseMapFlagsAttr>(mapType),
+            builder.getAttr<omp::VariableCaptureKindAttr>(
+                omp::VariableCaptureKind::ByRef),
+            /*var_ptr_ptr=*/mlir::Value{},
+            /*members=*/SmallVector<Value>{},
+            /*members_index=*/ArrayAttr{},
+            /*bounds=*/SmallVector<Value>{},
+            /*mapperId=*/mlir::FlatSymbolRefAttr(), globalOp.getSymNameAttr(),
+            builder.getBoolAttr(true));
+      };
+
+      if (isAlloc)
+        clauses.mapVars.push_back(createDescriptorMap(
+            omp::ClauseMapFlags::to | omp::ClauseMapFlags::always));
+
+      mlir::omp::ClauseMapFlags mapType =
+          isAlloc ? omp::ClauseMapFlags::storage : omp::ClauseMapFlags::del;
       mlir::omp::MapInfoOp mapInfo = mlir::omp::MapInfoOp::create(
-          builder, memOp.getLoc(), memOp.getMemref().getType(),
-          memOp.getMemref(),
-          TypeAttr::get(fir::unwrapRefType(memOp.getMemref().getType())),
-          builder.getAttr<omp::ClauseMapFlagsAttr>(
-              isa<fir::StoreOp>(memOp) ? omp::ClauseMapFlags::to
-                                       : omp::ClauseMapFlags::del),
+          builder, memOp.getLoc(), dataAddr.getType(), dataAddr,
+          TypeAttr::get(baseTy),
+          builder.getAttr<omp::ClauseMapFlagsAttr>(mapType),
           builder.getAttr<omp::VariableCaptureKindAttr>(
-              omp::VariableCaptureKind::ByCopy),
+              omp::VariableCaptureKind::ByRef),
           /*var_ptr_ptr=*/mlir::Value{},
           /*members=*/SmallVector<Value>{},
           /*members_index=*/ArrayAttr{}, bounds,
           /*mapperId=*/mlir::FlatSymbolRefAttr(), globalOp.getSymNameAttr(),
           builder.getBoolAttr(false));
       clauses.mapVars.push_back(mapInfo);
-      isa<fir::StoreOp>(memOp)
-          ? omp::TargetEnterDataOp::create(builder, memOp.getLoc(), clauses)
-          : omp::TargetExitDataOp::create(builder, memOp.getLoc(), clauses);
+
+      if (isAlloc) {
+        mlir::omp::MapInfoOp attachInfo = mlir::omp::MapInfoOp::create(
+            builder, memOp.getLoc(), dataAddr.getType(), dataAddr,
----------------
agozillon wrote:

I am probably missing something, but is there a reason we can't just emit a regular map for the BoxType here with the right top-level map types and let MapInfoFinalization.cpp expand it? If you also wanted a more fine grain way of detailing what bits specifically do or do not get mapped, ref_ptr/ref_ptee allows you to specify if you just want the descriptor on device for the map, or the base address, and attach_none allows you to omit the attach

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