[Flang][OpenMP] Generate correct present checks for implicit maps of optional allocatables

flang/include/flang/Optimizer/Builder/DirectivesCommon.h

@@ -243,6 +243,17 @@ genBaseBoundsOps(fir::FirOpBuilder &builder, mlir::Location loc,
   return bounds;
 }
 
+/// Checks if an argument is optional based on the fortran attributes
+/// that are tied to it.
+inline bool isOptionalArgument(mlir::Operation *op) {
+  if (auto declareOp = mlir::dyn_cast_or_null<hlfir::DeclareOp>(op))
+    if (declareOp.getFortranAttrs() &&
+        bitEnumContainsAny(*declareOp.getFortranAttrs(),
+                           fir::FortranVariableFlagsEnum::optional))
+      return true;
+  return false;
+}
+
 template <typename BoundsOp, typename BoundsType>
 llvm::SmallVector<mlir::Value>
 genImplicitBoundsOps(fir::FirOpBuilder &builder, AddrAndBoundsInfo &info,

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -2320,7 +2320,8 @@ genTargetOp(lower::AbstractConverter &converter, lower::SymMap &symTable,
 
       fir::factory::AddrAndBoundsInfo info =
           Fortran::lower::getDataOperandBaseAddr(
-              converter, firOpBuilder, sym, converter.getCurrentLocation());
+              converter, firOpBuilder, sym.GetUltimate(),
+              converter.getCurrentLocation());
       llvm::SmallVector<mlir::Value> bounds =
           fir::factory::genImplicitBoundsOps<mlir::omp::MapBoundsOp,
                                              mlir::omp::MapBoundsType>(

flang/lib/Optimizer/OpenMP/MapInfoFinalization.cpp

@@ -131,7 +131,8 @@ class MapInfoFinalizationPass
               boxMap.getVarPtr().getDefiningOp()))
         descriptor = addrOp.getVal();
 
-    if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()))
+    if (!mlir::isa<fir::BaseBoxType>(descriptor.getType()) &&
+        !fir::factory::isOptionalArgument(descriptor.getDefiningOp()))
       return descriptor;
 
     mlir::Value &slot = localBoxAllocas[descriptor.getDefiningOp()];
@@ -151,16 +152,22 @@ class MapInfoFinalizationPass
     mlir::Location loc = boxMap->getLoc();
     assert(allocaBlock && "No alloca block found for this top level op");
     builder.setInsertionPointToStart(allocaBlock);
-    auto alloca = builder.create<fir::AllocaOp>(loc, descriptor.getType());
+
+    mlir::Type allocaType = descriptor.getType();
+    if (fir::isBoxAddress(allocaType))
+      allocaType = fir::unwrapRefType(allocaType);
+    auto alloca = builder.create<fir::AllocaOp>(loc, allocaType);
     builder.restoreInsertionPoint(insPt);
     // We should only emit a store if the passed in data is present, it is
     // possible a user passes in no argument to an optional parameter, in which
     // case we cannot store or we'll segfault on the emitted memcpy.
     auto isPresent =
         builder.create<fir::IsPresentOp>(loc, builder.getI1Type(), descriptor);
     builder.genIfOp(loc, {}, isPresent, false)
-        .genThen(
-            [&]() { builder.create<fir::StoreOp>(loc, descriptor, alloca); })
+        .genThen([&]() {
+          descriptor = builder.loadIfRef(loc, descriptor);
+          builder.create<fir::StoreOp>(loc, descriptor, alloca);
+        })
         .end();
     return slot = alloca;
   }

flang/test/Lower/OpenMP/optional-argument-map-2.f90

@@ -0,0 +1,46 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp %s -o - | FileCheck %s
+
+module mod
+  implicit none
+contains
+  subroutine routine(a)
+    implicit none
+    real(4), allocatable, optional, intent(inout) :: a(:)
+    integer(4) :: i
+
+    !$omp target teams distribute parallel do shared(a)
+        do i=1,10
+            a(i) = i + a(i)
+        end do
+
+  end subroutine routine
+end module mod
+
+! CHECK-LABEL:   func.func @_QMmodProutine(
+! CHECK-SAME:      %[[ARG0:.*]]: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {fir.bindc_name = "a", fir.optional}) {
+! CHECK:           %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
+! CHECK:           %[[VAL_1:.*]] = fir.dummy_scope : !fir.dscope
+! CHECK:           %[[VAL_2:.*]]:2 = hlfir.declare %[[ARG0]] dummy_scope %[[VAL_1]] {fortran_attrs = #fir.var_attrs<allocatable, intent_inout, optional>, uniq_name = "_QMmodFroutineEa"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.dscope) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+! CHECK:           %[[VAL_8:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
+! CHECK:           %[[VAL_9:.*]]:5 = fir.if %[[VAL_8]] -> (index, index, index, index, index) {
+! CHECK:             %[[VAL_10:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK:             %[[VAL_11:.*]] = arith.constant 1 : index
+! CHECK:             %[[VAL_12:.*]] = arith.constant 0 : index
+! CHECK:             %[[VAL_13:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK:             %[[VAL_14:.*]] = arith.constant 0 : index
+! CHECK:             %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_13]], %[[VAL_14]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
+! CHECK:             %[[VAL_16:.*]]:3 = fir.box_dims %[[VAL_10]], %[[VAL_12]] : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
+! CHECK:             %[[VAL_17:.*]] = arith.constant 0 : index
+! CHECK:             %[[VAL_18:.*]] = arith.subi %[[VAL_16]]#1, %[[VAL_11]] : index
+! CHECK:             fir.result %[[VAL_17]], %[[VAL_18]], %[[VAL_16]]#1, %[[VAL_16]]#2, %[[VAL_15]]#0 : index, index, index, index, index
+! CHECK:           } else {
+! CHECK:             %[[VAL_19:.*]] = arith.constant 0 : index
+! CHECK:             %[[VAL_20:.*]] = arith.constant -1 : index
+! CHECK:             fir.result %[[VAL_19]], %[[VAL_20]], %[[VAL_19]], %[[VAL_19]], %[[VAL_19]] : index, index, index, index, index
+! CHECK:           }
+! CHECK:           %[[VAL_21:.*]] = omp.map.bounds lower_bound(%[[VAL_9]]#0 : index) upper_bound(%[[VAL_9]]#1 : index) extent(%[[VAL_9]]#2 : index) stride(%[[VAL_9]]#3 : index) start_idx(%[[VAL_9]]#4 : index) {stride_in_bytes = true}
+! CHECK:           %[[VAL_23:.*]] = fir.is_present %[[VAL_2]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> i1
+! CHECK:           fir.if %[[VAL_23]] {
+! CHECK:             %[[VAL_24:.*]] = fir.load %[[VAL_2]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK:             fir.store %[[VAL_24]] to %[[VAL_0]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK:           }

offload/test/offloading/fortran/optional-mapped-arguments-2.f90

@@ -0,0 +1,57 @@
+! OpenMP offloading regression test that checks we do not cause a segfault when
+! implicitly mapping a not present optional allocatable function argument and
+! utilise it in the target region. No results requiring checking other than
+! that the program compiles and runs to completion with no error.
+! REQUIRES: flang, amdgpu
+
+! RUN: %libomptarget-compile-fortran-run-and-check-generic
+module mod
+  implicit none
+contains
+  subroutine routine(a, b)
+    implicit none
+    real(4), allocatable, optional, intent(in) :: a(:)
+    real(4), intent(out) :: b(:)
+    integer(4) :: i, ia
+    if(present(a)) then
+       ia = 1
+       write(*,*) "a is present"
+    else
+       ia=0
+       write(*,*) "a is not present"
+    end if
+
+    !$omp target teams distribute parallel do shared(a,b,ia)
+    do i=1,10
+       if (ia>0) then
+          b(i) = b(i) + a(i)
+       end if
+    end do
+
+  end subroutine routine
+
+end module mod
+
+program main
+  use mod
+  implicit none
+  real(4), allocatable :: a(:)
+  real(4), allocatable :: b(:)
+  integer(4) :: i
+  allocate(b(10))
+  do i=1,10
+     b(i)=0
+  end do
+  !$omp target data map(from: b)
+
+  call routine(b=b)
+
+  !$omp end target data
+
+  deallocate(b)
+
+  print *, "success, no segmentation fault"
+end program main
+
+!CHECK: a is not present
+!CHECK: success, no segmentation fault