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

[Slava Zakharin]

================
@@ -1710,30 +1710,76 @@ bool ClauseProcessor::processLink(
 
 static mlir::omp::MapInfoOp
 createMapInfoOp(fir::FirOpBuilder &builder, mlir::Location loc,
-                mlir::Value baseAddr, std::stringstream &name,
-                mlir::SmallVector<mlir::Value> bounds, uint64_t mapType,
-                mlir::omp::VariableCaptureKind mapCaptureType,
-                mlir::Type retTy) {
-  mlir::Value varPtr, varPtrPtr;
-  mlir::TypeAttr varType;
-
+                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();
   }
 
-  varPtr = baseAddr;
-  varType = mlir::TypeAttr::get(
+  mlir::TypeAttr varType = mlir::TypeAttr::get(
       llvm::cast<mlir::omp::PointerLikeType>(retTy).getElementType());
 
   mlir::omp::MapInfoOp op = builder.create<mlir::omp::MapInfoOp>(
-      loc, retTy, varPtr, varType, varPtrPtr, bounds,
+      loc, retTy, baseAddr, varType, varPtrPtr, members, bounds,
       builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
       builder.getAttr<mlir::omp::VariableCaptureKindAttr>(mapCaptureType),
-      builder.getStringAttr(name.str()));
+      builder.getStringAttr(name));
+
   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 (fir::isAssumedShape(fir::unwrapRefType(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 abose
+  // 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());
----------------
vzakhari wrote:

I really think it should be `ByCopy`.  Consider this case:
```
program main
  real, pointer :: a(:)
  allocate(a(4))
  call test(a)
  print *, a
contains
  subroutine test(a)
    use omp_lib
    real, pointer :: a(:)
    integer :: i
    !$omp parallel num_threads(2) private(i) shared(a)
    i = omp_get_thread_num()
    print *, i
    !$omp target map(a(2*i+1:2*i+2))
    a(2*i+1:2*i+2) = i + 1
    !$omp end target
    !$omp end parallel
  end subroutine test
end program main
```

The same in-memory descriptor is used for the `map`.  The the first host thread executing `target` will create its data and the descriptor in the device data environment.  Then the second thread executing `target` will create its data, but it will not create a new descriptor in the device data environment, because it is already present.  There is no way for the two different data pointers to exist in the same descriptor on the device at the same time.  If I understand it correctly, using `ByCopy` (i.e. `OMP_MAP_LITERAL`) should allow creating two different private copies of the descriptor on the device.

I might be missing something, but this looks like a valid OpenMP to me :)

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